Earlier, we discussed that, using automatic differentiation, Torch helps us compute the gradient of the loss \(L_{\hat {p}_{\text {data}}}(\theta )\) with respect to each parameter that requires gradients, in particular the model parameters. The gradient points in the direction of steepest increase, so we simply take a step in the opposite direction. For a parameter \(\mathbf {W}\), the update is
where \(\mathbf {W}_0\) is the initial location in the loss landscape, \(\theta _t\) includes \(\mathbf {W}_t\), \(\mathbf {W}_t\) is a parameter of a given layer, and \(\alpha _t\) is the learning rate. The update continues until a stopping criterion is met.
Computing \(\frac {\partial L}{\partial \mathbf {W}}\) over the entire dataset can be costly. The expected loss can be written as
where \(N\) is the number of samples in the dataset. Computing this quantity and its gradient over the entire dataset at every update can be prohibitively expensive, especially since deep learning datasets can be very large. In general, the more data we have at hand, the better the true data distribution can be approximated.
Instead of using the entire dataset for every update, a randomly sampled mini-batch \(\mathbb {B}_t \subset \mathbb {D}\) is used to approximate the full-dataset gradient:
The parameter update therefore becomes
Using only a randomly sampled mini-batch introduces noise into the gradient estimate, making it a less precise approximation of the full-dataset gradient. This stochastic approximation is why the method is referred to as stochastic gradient descent rather than full-batch gradient descent.
So far, we have discussed how tensors are represented, how the computation graph is constructed, and how gradients are propagated during the backward pass. The final step of a training iteration is updating the model parameters. Before understanding how an optimizer performs this update, it is useful to see how PyTorch keeps track of a model’s parameters in the first place.
Every nn.Module owns a hidden field named self._parameters, which stores all learnable parameters that belong directly to that module, such as the weight matrix and bias of a linear layer (PyTorch automatically registers parameters when they are assigned as module attributes).
def __setattr__(self, name: str, value: Union[Tensor, "Module"]) -> None: ... if isinstance(value, Parameter): if params is None: raise AttributeError( "cannot assign parameters before Module.__init__() call" ) ... self.register_parameter(name, value) ....
Every module also owns self._modules, which stores all child modules registered as attributes of the current module. This allows PyTorch to recursively traverse all parameters in a model. For example, an nn.Linear layer is itself an nn.Module.
def named_parameters( self, prefix: str = "", recurse: bool = True, remove_duplicate: bool = True ) -> Iterator[tuple[str, Parameter]]: r"""Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself. Args: prefix (str): prefix to prepend to all parameter names. recurse (bool): if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module. remove_duplicate (bool, optional): whether to remove the duplicated parameters in the result. Defaults to True. Yields: (str, Parameter): Tuple containing the name and parameter Example:: >>> # xdoctest: +SKIP("undefined vars") >>> for name, param in self.named_parameters(): >>> if name in ['bias']: >>> print(param.size()) """ gen = self._named_members( lambda module: module._parameters.items(), prefix=prefix, recurse=recurse, remove_duplicate=remove_duplicate, ) yield from gen
Once the backward pass has completed, every parameter contains its accumulated gradient inside its .grad field.
The optimizer can therefore iterate over every parameter and update it according to the chosen optimization algorithm. In its simplest form, stochastic gradient descent or SGD updates each parameter by moving it in the opposite direction of its gradient. Strictly speaking, the gradient with respect to a parameter is often a tensor rather than a vector (e.g., the gradient of the loss with respect to a weight matrix has the same shape as the matrix itself). More generally, one can think of it as the derivative (or Jacobian) of the loss with respect to the parameter. Regardless of its shape, the jacobian or the gradient describes how the loss changes as the parameter changes (i.e rate of change), and moving in its opposite direction yields the direction of steepest decrease.
def step(self, closure=None): """Perform a single optimization step. Args: closure (Callable, optional): A closure that reevaluates the model and returns the loss. """ loss = None if closure is not None: with torch.enable_grad(): loss = closure() for group in self.param_groups: params: list[Tensor] = [] grads: list[Tensor] = [] momentum_buffer_list: list[Tensor | None] = [] has_sparse_grad = self._init_group( group, params, grads, momentum_buffer_list ) sgd( params, grads, momentum_buffer_list, weight_decay=group["weight_decay"], momentum=group["momentum"], lr=group["lr"], dampening=group["dampening"], nesterov=group["nesterov"], maximize=group["maximize"], has_sparse_grad=has_sparse_grad, foreach=group["foreach"], fused=group["fused"], grad_scale=getattr(self, "grad_scale", None), found_inf=getattr(self, "found_inf", None), ) if group["momentum"] != 0: # update momentum_buffers in state for p, momentum_buffer in zip( params, momentum_buffer_list, strict=True ): state = self.state[p] state["momentum_buffer"] = momentum_buffer return loss