Optimizers in Model Training

Before you can run a model, you need to know whether the hardware you have can hold it. Knowing how much memory a model needs is what lets you choose the right hardware, and often the practical question is whether a setup you already own can afford to run a particular model. Training a model is more demanding than running it, and a large part of that extra demand comes from one idea: which parameters are allowed to change. This article builds that idea up step by step, starting from what a parameter is and ending with why the number of trainable parameters shapes memory use during finetuning.

The source describes three situations that differ in which parameters change. During pre-training, all model parameters are updated. During inference, no model parameters are updated. During finetuning, some or all model parameters may be updated. Of the three, finetuning is the case where the split between trainable and frozen parameters matters most, because the parameters you leave frozen are exactly the ones that do not need to be updated.

This is the first clue about where extra work comes from. Every parameter that is allowed to change has to be handled during training. Every parameter that is frozen can be left alone. Keep this distinction in mind, because the rest of the article depends on it.

As of this writing, neural networks are typically trained using a mechanism called backpropagation. With backpropagation, each training step consists of two phases. The first is the forward pass, which is the process of computing the output from the input. The second is the backward pass, which is the process of updating the model's weights using aggregated signals from the forward pass. Read that carefully: the backward pass does not simply reverse the forward pass. It uses signals gathered during the forward pass to decide how the weights should change.

The order of events within one step is the thing to hold on to. The forward pass runs from the input through the model to a computed output. The backward pass then starts from the comparison between that computed output and the expected output, and works backward to adjust the trainable parameters. The diagram below traces that order across the actors involved.

Whether both phases run depends on what you are doing with the model. During inference, only the forward pass is executed. During training, both passes are executed. This is the core reason inference and training have different memory needs: training carries the extra work of the backward pass.

At a high level, the backward pass works in three stages, and each stage produces something the next stage needs. First, the backward pass computes the loss, which is the measured gap between what the model produced and what was expected. Second, it computes a gradient for each trainable parameter. Third, it uses the optimizer to decide how to adjust each parameter. The loss is the starting signal, the gradient is the per-parameter signal derived from it, and the optimizer turns those signals into an actual adjustment.

Notice what the diagram does not contain: frozen parameters. Only trainable parameters get a gradient and pass through the optimizer. This is why the split between the two kinds matters for what training computes, not just for what it is allowed to change.

The reason the number of trainable parameters matters is what the backward pass needs to keep around. During the backward pass, each trainable parameter comes with additional values: its gradient and its optimizer states. Each trainable parameter therefore brings gradient and optimizer state storage with it, and the more trainable parameters there are, the more memory is needed to store these additional values. This is why one key factor that determines a model's memory footprint during finetuning is its number of trainable parameters.

The source frames memory use in general terms. A model's memory footprint depends on the model, the workload, and the optimization techniques used to reduce memory usage. Because it is impossible to account for every optimization technique and workload, the source outlines only approximate calculations that give a rough idea of how much memory is needed to operate a model during inference and training. The useful takeaway for this article is directional: adding trainable parameters adds the per-parameter gradient and optimizer state burden, and freezing parameters avoids it for those parameters.

Consider a small model with two groups of parameters: some are trainable and some are frozen. Before the update, every parameter holds its current value. After one training step, the optimizer has adjusted the trainable parameters based on their gradients, so their values have moved. The frozen parameters still hold the values they started with. The change is selective, and the selection is determined entirely by the trainable and frozen split.

This example is about inference, where only the forward pass runs. Finetuning adds the backward pass and the per-trainable-parameter storage described earlier, so the same kind of memory check becomes more involved when training is on the table.