How Neural Networks Learn: Rolling Down the Free Energy Landscape with Gradient Descent

Upon Completion of this Topic

You will understand how the 33.62 million dials (weights and biases) we saw in Part 2 are automatically adjusted to complete training. We will connect the relationships between the three components – cost function, gradient, and learning rate – into a single story, and explain why this process is essentially the same as the common free energy minimization process in nature.

This part serves as the foundation for Part 4 (Backpropagation) and the fundamental principle of RLHF, which will be discussed in Part 11.

The Dilemma of a Graduate Student Discovering a New Enzyme

Let's say you are in the lab, characterizing a newly discovered enzyme. This enzyme catalyzes a specific metabolic reaction, but the optimal reaction conditions are not yet known. Your task is to experimentally determine when this enzyme is most active.

Let's first assume that only two conditions affect the enzyme's activity: temperature and pH. All other conditions (substrate concentration, ionic strength, presence or absence of cofactors) are kept constant.

If you had unlimited time and resources, how would you proceed? You could test all combinations, varying the temperature from 4°C to 60°C in 1°C increments and the pH from 4.0 to 10.0 in 0.1 increments. That would be 57 × 61 = 3,477 conditions. If you repeated each condition three times, you would have approximately 10,000 experiments. It would take a few months, but eventually, you would find the optimum.

However, the real situation is much worse. What if there were not two, but five, ten, or even one hundred parameters? If you use a grid search for each parameter, the number of combinations will explode exponentially. Even with just 100 parameters, each with 10 steps, you would have 10^100 combinations. This requires more experiments than the number of atoms in the universe (approximately 10^80), which is obviously impossible.

The neural network we are training has 33.62 million parameters, as we saw in Part 2. A grid search is physically impossible. We need a much more clever method.

The method we will learn here is Gradient Descent. In fact, this method is not just a special technique for training neural networks, but a principle that nature has been using for a long time in chemical reactions, protein folding, and ecosystem balance. How does nature find the optimal state without searching through a multidimensional condition space? The answer is that neural networks borrow this principle directly.

Free Energy Landscape: Nature's Minimization Problem

You may have encountered the concept of the free energy landscape somewhere in your undergraduate chemistry or biochemistry courses. It is a curve showing the progress of a chemical reaction on the x-axis and the free energy on the y-axis. It is a rugged landscape with hills, valleys, and saddle points.

How does the reaction proceed in this landscape? It rolls down towards lower free energy. In some cases, it must overcome an activation energy barrier, but eventually, it settles into a more stable state with lower free energy. This principle explains organic reactions, redox reactions, protein folding, and self-assembly.

The key is this: Even if you don't know the overall picture of the landscape, you can move towards the minimum point if you know at each point "which direction will lower the free energy." A rolling ball does not see the entire landscape in advance. It only senses the slope under its feet.

Learning in neural networks is conceptually the same process.

We have:

• Instead of a free energy landscape, we have the landscape of the cost function.
• Instead of reaction coordinates, we have a parameter space (millions or billions of dimensions).
• Instead of the state of minimum free energy, we have the goal of minimum cost (most accurate prediction).

If you keep this perspective in mind, all subsequent concepts will connect naturally.

Cost Function: Quantifying How Bad the Model Is

First, let's talk about the cost function. The cost function or loss function is a function that numerically represents how different the current prediction of the neural network is from the correct answer. The higher the value, the farther the prediction is from the correct answer, and the closer to zero, the closer to the correct answer.

Let's return to the example of the pathology slide from Part 2. Suppose the neural network outputs a probability of 0.7 for normal, 0.2 for positive, and 0.1 for malignant for a certain slide, but the actual correct answer is "malignant." How bad is this prediction?

We can use a cost function called Cross-Entropy. In this case, the value will be approximately -log(0.1) ≈ 2.30. If the neural network had assigned a probability of 0.9 to the correct answer, malignant, it would be -log(0.9) ≈ 0.11, which is much lower. If it assigned a probability of 1 to the correct answer, it would be exactly 0.

We calculate this for all examples in the training dataset and take the average. The average cost for the entire training dataset is what we want to minimize. This average cost is called the training loss.

Although we will discuss it in detail in Appendix A.1, let's just grasp the key concepts here. The cost function has two properties:

First, it depends only on the parameters. If you fix the dataset and the neural network structure, the value of the cost function is determined by the values of the parameters (weights and biases). If you change the parameters slightly, the cost will change slightly.

Second, it is differentiable. If the parameters change continuously, the cost also changes continuously. That is, you can answer the question, "If you increase this parameter slightly, how much will the cost increase or decrease?" This partial derivative is the essence of the gradient that will be discussed later.

Analogy: In the enzyme example, the cost you define is "maximum activity value - activity at the current conditions." If you are at the optimal conditions, this value is 0. As you move away from the optimum, it increases. If the values of temperature and pH change slightly, the activity will change slightly, and the cost will also change slightly. The cost function of the neural network is exactly the same structure, except that there are not two, but 33.62 million parameters.

Gradient: Which Direction Should We Move to Go Down?

Let's return to the problem of the ball rolling in the landscape. We said that the ball only needs to sense the slope under its feet. What corresponds to this "gradient" in the parameter space?

The gradient is a vector that collects the partial derivatives of the cost function with respect to each parameter. That is, if there are 33.62 million parameters, the gradient vector also has 33.62 million dimensions.

The meaning of each element of the gradient vector is as follows: "If you increase this parameter slightly, how much will the cost increase (positive) or decrease (negative)?"

• If the value of the element is large and positive: Increasing this parameter will quickly increase the cost → it should be decreased.
• If the value of the element is large and negative: Increasing this parameter will quickly decrease the cost → it should be increased.
• If the value of the element is close to 0: This parameter has little effect on the cost in the vicinity → it can be left as is.

This vector, calculated for all 33.62 million parameters, is the gradient.

The direction of the gradient points to the direction of the fastest increase in cost. Therefore, we must move in the opposite direction. This is the origin of the name Gradient Descent. Descent is the opposite of ascent.

Now, the image of a ball rolling in a multidimensional free energy landscape is more concretely defined as follows:

1. Calculate the gradient at the current parameter location.
2. Move the parameters slightly in the opposite direction of the gradient.
3. Recalculate the gradient at the new location.
4. Repeat.

One move is a step, and the entire process of the parameters gradually converging to better values (lower cost) is called training.

But how do we calculate the gradient? There are 33.62 million parameters, but how do we calculate the partial derivative for each of them? This is where the backpropagation algorithm, the protagonist of Part 4, comes in. We will discuss it in detail in Part 4. For now, let's assume that "we can calculate the gradient, and we will postpone how it is calculated to Part 4," and continue with the story.

Learning Rate: How Much Should We Move?

Now that we know the direction opposite to the gradient, the next question is: How much should we move?

The value that adjusts this magnitude is the learning rate. It is denoted by the Greek letters η (eta) or α (alpha).

The parameter update formula is very simple:

New_parameter = Current_parameter - Learning_rate × Gradient

If the learning rate is large, we move a lot in one step, and if it is small, we move a little.

If the learning rate is too small, the training will be slow. Training that would take a few hours to reach the vicinity of the optimum may take several days.

If the learning rate is too large, we will overshoot the bottom of the valley and end up on the other side of the hill. In the next step, we move in the opposite direction, but we overshoot again. This is called oscillation. In extreme cases, the training may diverge, ending up in a much worse state.

An appropriate learning rate depends on the curvature of the landscape. If the valley is narrow and deep, a small learning rate is appropriate, and if it is wide and shallow, a large learning rate is appropriate. In practice, it is standard to start with a large learning rate in the early stages of training (for rapid approximation) and then gradually decrease the learning rate in the later stages. This is called learning rate scheduling.

Bio-mapping: Let's revisit the directional evolution screening mentioned in Part 1. The learning rate corresponds exactly to the mutation rate. If the mutation rate is too low, it will not escape from the original sequence and evolution will not proceed, and if it is too high, useful combinations will not be maintained and will collapse. The optimal mutation rate depends on the problem and the round, and it is usually lowered in the later stages - the concept is the same as learning rate scheduling in neural network training.

Local Minima: Traps in the Landscape

In the free energy landscape, there are two types of minima where reactions often get trapped.

Global minimum: The point where the free energy is lowest in the entire landscape. The most thermodynamically stable state.

Local minimum: A point that is lower than its surroundings but not the global minimum. If the reaction gets trapped here, it cannot overcome the activation energy barrier and remains in a bad state.

The same problem exists in neural network training. While rolling down the parameter space, we may get trapped in a local minimum. This is a state where the training loss has decreased to some extent and then stops decreasing.

Interestingly, it has been empirically shown that the training loss landscape of large neural networks (hundreds of millions of parameters) has much fewer local minima than theoretically expected. This is because when the number of parameters is very large, if it is a local minimum in one direction, it is usually possible to go down in another direction. That is, saddle points (a maximum in one direction and a minimum in the other) are much more common than true local minima. Fortunately, saddle points can be escaped by shaking them a little.

Even so, in practice, training often encounters the problem of "the training loss stagnating at a certain point." There are several tricks to alleviate this, which will be discussed in the optimizers that will follow.

Bio analogy: The protein folding problem has the same structure. When a polypeptide chain folds, if it gets trapped in a local minimum and folds incorrectly, it forms aggregates - the amyloid beta and alpha-synuclein aggregates in Alzheimer's and Parkinson's diseases, which were mentioned in Part 1. Nature solves this problem with chaperone proteins and the heat shock response. The optimizers in neural network training play the role of these chaperones.

Stochastic Gradient Descent — Insights from Batching

Now, let's consider a practical problem. To calculate the cost function and obtain the gradient for the entire training dataset, you need to pass all training examples through the neural network in each step. If the dataset has hundreds of millions of examples, a single step can take several minutes to hours. At this rate, training becomes impractical.

The solution is to use only a small portion (minibatch) of the dataset in each step. For example, in a dataset of 2 million examples, randomly select 256 examples for each step, calculate the cost and gradient for that batch, and update the parameters.

This is called Stochastic Gradient Descent (SGD). It's called "stochastic" because the gradient in each step is an approximation of the true gradient for the entire dataset. Since the batch is selected randomly, you get a slightly different gradient direction each time.

This noise can actually be a blessing. If you were to perfectly follow the gradient of the entire dataset, you might end up at a saddle point and get stuck. However, with stochastic noise, it's easier to escape the saddle point. This is why SGD often works better in practice than theoretical pure gradient descent.

Trade-offs of batch size:

• Small batch (e.g., 32～256): Larger noise, more unstable training, but better at escaping saddle points and saves GPU memory.
• Large batch (e.g., 4096～32768): Smaller noise, more stable, but risks getting stuck in saddle points and requires more GPU memory.

In practice, you use the largest batch size that fits within your hardware budget, or you tune the learning rate in combination with the batch size.

The Evolution of Optimizers — From Pure SGD to Adam

Pure SGD simply moves in the opposite direction of the gradient by the learning rate in each step. This simplicity can lead to several problems.

Problem 1: Directional oscillation. When it oscillates back and forth on one side of a valley, making it difficult to descend.

Problem 2: Difficulty in tuning the learning rate. The optimal learning rate may be different for each parameter, but we use the same learning rate for all.

Problem 3: Vanishing gradients. Parameters with very small gradients may effectively stop updating after training.

Optimizers have been developed to mitigate these problems.

Momentum. Introduces the concept of inertia from physics. It maintains some of the movement direction from the previous step, suppressing oscillations and accelerating towards the valley. It's like a rolling ball with inertia instead of friction.

RMSprop. Tracks the average of the recent magnitudes of the gradients for each parameter and automatically adjusts the learning rate for each parameter. It automatically reduces the learning rate for parameters with large gradients and maintains it for parameters with small gradients.

Adam (Adaptive Moment Estimation). Combines Momentum and RMSprop. It is now the de facto standard optimizer for training neural networks. Most papers and practical code use Adam (or its variant, AdamW) as the default.

AdamW. An improved version of Adam that separates weight decay. It is now the de facto standard for training large language models.

The optimizer used in Part 1's pre-training and fine-tuning is mostly AdamW. Appendix A.4～A.7 covers the exact formulas and hyperparameters of each optimizer.

The Actual Training Process

Let's put all the pieces together. The entire cycle of neural network training looks like this:

Initialize parameters randomly
Repeat (tens of thousands to millions of times):
    Randomly select a minibatch (e.g., 256 examples)
    Perform forward propagation on the minibatch (Part 2) → Get predictions
    Calculate the cost based on the predictions and the ground truth
    Perform backpropagation (Part 4) → Calculate gradients
    Update parameters using an optimizer (Adam)
Terminate when the training loss reaches a satisfactory level

Each step takes milliseconds to seconds, and you need to repeat these steps tens of thousands to millions of times for large models to converge. This is why training large neural networks (GPT-level) can take weeks to months.

Back to bio mapping. In Part 1, we mapped this entire cycle to a directional evolution cycle.

• Initial parameters = initial random antibody library
• Minibatch = samples to be tested in this round
• Cost calculation = measuring the target binding strength
• Parameter update = reconstituting the library in a preferred direction (mutation + selection)
• Repeating millions of steps = several rounds of screening

Interesting perspective. We know why directional evolution is much faster than natural evolution — because the selection pressure is explicit and strong. Neural network training is also faster for the same reason — because there is an explicit selection pressure (cost function) that tells you the direction to go in each step. Unlike natural evolution, which relies heavily on random drift, neural network training has information in the form of gradients, making it much more directional.

Evolution of Parameters — What Changes During Training?

Let's visualize what actually changes during training. Consider the example of a pathological neural network.

At the beginning of training: Parameters are random. The neural network's predictions are completely random (approximately 33% probability for normal, positive, and malignant). Training loss is very high.

Early training (thousands of steps): Parameters are gradually adjusted, and the neural network begins to capture some basic patterns. For example, "if there are many dark spots in the image, the probability of a tumor increases," which is a very low-level pattern. Training loss drops rapidly.

Mid-training (tens of thousands of steps): Hidden layer neurons start to capture more sophisticated features. Cell boundary and nucleus abnormality detectors, as discussed in Part 2, are formed. Training loss decreases gradually.

Late training (hundreds of thousands of steps): Parameters are fine-tuned, and the accuracy of classifying difficult boundary cases (e.g., early-stage tumor vs. normal tissue) is improved. The rate of decrease in training loss becomes very slow.

Overfitting point: At some point, the training loss continues to decrease, but the validation loss (the loss on data not used for training) starts to increase. This is a sign that the neural network is starting to memorize the specific noise in the training data. You should stop training before this point to get good performance in the real world. This is called early stopping, and monitoring the validation loss is a standard practice.

Bio Application Scenarios

How do the concepts in this part translate to practical applications?

Scenario 1 — Training a Cell Type Classifier

When you actually train the scRNA-seq cell type classifier discussed in Part 2, you use the concepts in this part directly.

• Cost function: Cross-entropy between the predicted probability distribution and the ground truth cell type for each cell.
• Minibatch: 128 to 1024 cells from the training dataset.
• Optimizer: Usually Adam or AdamW.
• Learning rate: Start with 1e-3 and reduce to 1e-4 after 10 epochs.
• Early stopping: Stop if the accuracy on the validation dataset does not improve for 5 epochs.

Libraries like Scanpy and scVI handle this training loop internally, but they provide APIs that allow you to tune the optimizer and learning rate.

Scenario 2 — AlphaFold Training Follows the Same Principles

The training of AlphaFold (protein structure prediction) by DeepMind also follows the same principles as this part.

• Cost function: Coordinate error between the predicted structure and the actual structure (multiple components).
• Minibatch: 128 proteins from the training database (PDB).
• Optimizer: AdamW.
• Learning rate schedule: Warmup followed by cosine decay.
• Training scale: 128 TPUs for several weeks.

The number of parameters is in the tens of millions to hundreds of millions. This matches the scale discussed in Part 2.

Scenario 3 — An Alternative to Bayesian Optimization for Optimizing Experimental Conditions

You can also use it in reverse. To optimize experimental conditions (enzyme activity, cell culture conditions, gene editing efficiency), you can train a neural network to learn the mapping from conditions to results, and then find the optimal conditions using gradient descent within the neural network.

Traditionally, these problems are solved using Bayesian Optimization, but a neural network-based approach can be more efficient if there is enough data. This is especially true when the parameter space is high-dimensional (more than tens of dimensions).

Key Takeaways

• Neural network training is an optimization problem of finding the minimum of a cost function. This is exactly the same concept as the process of minimizing free energy in nature.
• The gradient of the cost function is the direction in which the cost increases most rapidly in parameter space. Moving in the opposite direction reduces the cost.
• The magnitude of the movement is the learning rate. It is conceptually the same as the mutation rate in directional evolution.
• In practice, we use stochastic gradient descent (SGD), which uses only a minibatch of the data. The noise helps escape saddle points.
• More sophisticated optimizers, such as Adam and AdamW, are much more practical than pure SGD. They are now the standard.
• Local minima and saddle points are potential pitfalls, but saddle points are more common in high-dimensional space, and SGD noise usually helps escape them.
• Overfitting starts in the later stages of training, so monitoring the validation loss and using early stopping are standard practices.

📐 Appendix — Mathematical Formulas for Experts

Difficulty: Very Hard Target Audience: Readers familiar with linear algebra, calculus, and optimization theories at the graduate level.

A.1 Cost Function Definitions

Mean Squared Error (MSE) — Standard for regression problems:

L_MSE(θ) = (1/N) · Σ_{i=1}^{N} (y_i - f_θ(x_i))^2

• x_i: The i-th input.
• y_i: The i-th ground truth.
• f_θ(x_i): The neural network prediction.
• θ: The set of all parameters.

Cross-Entropy — Standard for classification problems:

L_CE(θ) = -(1/N) · Σ_{i=1}^{N} Σ_{c=1}^{C} y_{i,c} · log(f_θ(x_i)_c)

• C: The number of classes.
• y_{i,c}: One-hot ground truth (1 for the correct class c, 0 otherwise).
• f_θ(x_i)_c: The probability assigned by the neural network to class c.

This is the same form as defined in Part 1, Appendix A.3.

A.2 Gradient Definition

The gradient of the cost function with respect to the parameter θ is:

∇_θ L = [ ∂L/∂θ_1, ∂L/∂θ_2, ..., ∂L/∂θ_P ]

P is the total number of parameters (millions to billions).

Each partial derivative ∂L/∂θ_i is the limit of how much the cost changes when the parameter θ_i is slightly increased:

∂L/∂θ_i = lim_{ε → 0} [L(θ_1, ..., θ_i + ε, ..., θ_P) - L(θ)] / ε

An algorithm to efficiently calculate these partial derivatives is backpropagation (Part 4).

A.3 Vanilla Gradient Descent Update Rule

A single step of parameter update:

θ_{t+1} = θ_t - η · ∇_θ L(θ_t)

• θ_t: The parameters at step t.
• η: The learning rate.
• ∇_θ L(θ_t): The gradient at step t.

Repeating this will converge to a local minimum (convergence is guaranteed only under certain conditions).

A.4 SGD's Mini-Batch Approximation

Instead of the entire dataset, approximate the gradient using a mini-batch B_t ⊂ D (random sample, size |B|):

∇̂_θ L(θ_t) = (1/|B|) · Σ_{i ∈ B_t} ∇_θ ℓ(θ_t, x_i, y_i)

ℓ is the loss for a single example. This is an unbiased estimator of the overall gradient:

E[∇̂_θ L] = ∇_θ L (overall)

A.5 Momentum

Maintain the direction of movement from the previous step (velocity) like inertia:

v_t = β · v_{t-1} + ∇̂_θ L(θ_t)
θ_{t+1} = θ_t - η · v_t

• β: The momentum coefficient (typically 0.9).

This is similar to the movement of a ball with a friction coefficient of (1 - β) in a physical simulation.

A.6 RMSprop

Automatically adjust the learning rate based on the average of the recent squared magnitudes of the gradients for each parameter:

s_t = ρ · s_{t-1} + (1 - ρ) · [∇̂_θ L(θ_t)]^2
θ_{t+1} = θ_t - η · ∇̂_θ L(θ_t) / (sqrt(s_t) + ε)

• ρ: The exponential moving average coefficient (typically 0.999).
• ε: A small value for numerical stability (e.g., 1e-8).

For parameters with large gradients, sqrt(s_t) becomes large, effectively reducing the learning rate.

A.7 Adam (Adaptive Moment Estimation)

Combines Momentum and RMSprop:

m_t = β_1 · m_{t-1} + (1 - β_1) · ∇̂_θ L(θ_t)          (first moment)
v_t = β_2 · v_{t-1} + (1 - β_2) · [∇̂_θ L(θ_t)]^2      (second moment)
m̂_t = m_t / (1 - β_1^t)                                (bias correction)
v̂_t = v_t / (1 - β_2^t)                                (bias correction)
θ_{t+1} = θ_t - η · m̂_t / (sqrt(v̂_t) + ε)

Typical hyperparameters:

• β_1 = 0.9
• β_2 = 0.999
• ε = 1e-8

Bias correction addresses the issue of the moments being biased towards zero early in training (when t is small).

A.8 AdamW (Adam with Decoupled Weight Decay)

Separate weight decay as a separate term instead of applying it within Adam:

θ_{t+1} = θ_t - η · m̂_t / (sqrt(v̂_t) + ε) - η · λ · θ_t

• λ: The weight decay coefficient (typically 0.01).

This is different from L2 regularization and exactly shrinks the weights. Standard for training large language models.

A.9 Learning Rate Schedule

Cosine annealing:

η_t = η_min + (1/2) · (η_max - η_min) · (1 + cos(π · t / T))

T is the total number of training steps. The learning rate decreases smoothly along a cosine curve.

Warmup: Linearly increase the learning rate from 0 to the target value for the first few steps of training. Mitigates initial instability in large models.

One-cycle policy: warmup → cosine decay → maintain at a very low value for a short period. Widely used in practice.

A.10 Saddle Point Analysis

Classify critical points by the sign of the eigenvalues of the Hessian H = ∇^2 L(θ) in parameter space:

• All eigenvalues > 0: Local minimum (convex in all directions).
• All eigenvalues < 0: Local maximum.
• Mixed signs: Saddle point.
• Some eigenvalues = 0: Flat point (difficult to analyze).

In high-dimensional neural networks, saddle points are much more common than local minima. Theoretically, the probability that a random critical point in P dimensions is a local minimum is approximately 2^{-P}, which is exponentially small. Therefore, in a 33.62 million-dimensional space, virtually all encountered critical points are saddle points.

References

All content, scenarios, analogies, and numbers in this part are developed in-house by BioPlayground. The following are external references that can help with conceptual learning.

• Adam original paper: Kingma & Ba, "Adam: A Method for Stochastic Optimization" (ICLR 2015)
• AdamW original paper: Loshchilov & Hutter, "Decoupled Weight Decay Regularization" (ICLR 2019)
• Standard deep learning textbook: Goodfellow et al., "Deep Learning" Chapter 8 (Optimization)
• Loss landscape visualization: Li et al., "Visualizing the Loss Landscape of Neural Nets" (NeurIPS 2018)
• Gradient descent convergence theory: Nocedal & Wright, "Numerical Optimization" (Springer)
• Deep learning visualization education: 3Blue1Brown "Deep Learning" Ch 2 (YouTube) — for pedagogical reference
• AlphaFold paper: Jumper et al., "Highly accurate protein structure prediction with AlphaFold" (Nature 2021)

This part is the second part of the principles part and the foundation for Part 4 (backpropagation). The next part will delve into how to efficiently obtain the gradients that are "assumed to be calculable" in this part for all 33.62 million parameters.

Next Concepts

• Ep. #4 backpropagation-intuition — How to calculate gradients. The democratic tug-of-war of error and the chain rule.
• Ep. #5 transformer-and-embedding — How this learning principle is extended in Transformers.
• Ep. #11 hallucination-and-alignment — RLHF is also ultimately an application of the gradient descent from this episode.