LEARN · NEURAL NETWORKS FROM THE GROUND UP
Training a neural network is an optimization problem disguised as an engineering project. We choose an architecture, prepare data, define a loss function, and then repeatedly update millions of parameters until the model learns useful patterns.
For years, stochastic gradient descent (SGD) was the default answer. It is simple, mathematically clean, and still capable of producing excellent models. However, modern deep learning introduced new challenges: enormous parameter counts, noisy gradients, distributed training, and architectures that behave very differently from classic models.
This pushed practitioners toward adaptive optimizers such as Adam and AdamW.
The interesting lesson from modern training practice is that the optimizer is only one part of the recipe. The learning-rate schedule often determines whether an otherwise good optimizer succeeds or fails.
A carefully designed schedule can allow a relatively simple optimizer to outperform a poorly configured adaptive optimizer. Conversely, even a strong optimizer can struggle when the learning rate is too aggressive, too conservative, or poorly timed.
The optimizer answers one question:
How should the gradient be converted into a parameter update?
The learning-rate schedule answers another:
How large should that update be at this point in training?
That second question is frequently underestimated.
This article builds an intuition for SGD, Adam, and AdamW, then implements a controlled comparison using the same neural network, the same synthetic retail-demand dataset, and the same warmup plus cosine decay schedule.
A neural network produces predictions. A loss function measures the difference between those predictions and the target values. Backpropagation calculates gradients, which describe how parameters should change to reduce the loss.
The simplest update rule is SGD:
parameter = parameter - learning_rate × gradient
The learning rate controls the size of each step.
A large learning rate allows fast movement but can overshoot useful solutions. A small learning rate is stable but can make training painfully slow or leave the model stuck.
SGD has an important strength: it directly follows the gradient. Because the behavior is predictable, researchers have spent decades understanding how it interacts with regularization, batch size, and learning-rate schedules.
Its weakness is that every parameter receives the same learning rate.
Imagine a neural network where one parameter has very large gradients and another has tiny gradients. SGD treats both cases identically. Large models with complicated optimization landscapes often benefit from a more adaptive approach.
Adam, short for Adaptive Moment Estimation, combines momentum with adaptive learning rates.
Instead of only considering the current gradient, Adam keeps moving averages of previous gradients.
It tracks:
-
The first moment: a smoothed estimate of gradient direction.
-
The second moment: a smoothed estimate of gradient magnitude.
Conceptually:
First moment = average gradient direction Second moment = average squared gradient size
The update behaves approximately like:
Parameter update = learning rate × first moment / (√second moment + small constant)
This gives each parameter an individual effective learning rate.
Adam became popular because it often trains successfully with less manual tuning. It is especially useful when:
-
Gradients are noisy.
-
Different parameters have different scales.
-
Fast initial convergence matters.
-
The model is large and expensive to experiment with.
The trade-off is that the fastest route to low training loss is not always the route to the best generalization. Some workloads still achieve better final performance with SGD after careful tuning.
AdamW modifies how weight decay is applied.
Weight decay is a regularization technique that discourages unnecessarily large weights. Large weights can make a model memorize training examples instead of learning patterns that transfer to new data.
Traditional implementations often mixed weight decay into Adam’s adaptive gradient calculation. This caused the regularization behavior to depend on Adam’s scaling mechanism.
AdamW separates these concerns:
-
Adam handles gradient-based parameter updates.
-
Weight decay independently shrinks model weights.
This makes the optimizer easier to reason about and improves consistency.
AdamW is now a common choice for transformer-based models and many modern neural networks because it provides reliable training behavior without requiring the same level of optimizer-specific tuning as SGD.
A constant learning rate is rarely the best choice.
Training has different phases.
At the beginning:
-
Parameters are mostly random.
-
The model needs to discover useful directions.
-
Larger updates can accelerate progress.
Near the end:
-
The model is already near a good solution.
-
Large updates can damage learned patterns.
-
Smaller steps help refine the final result.
A learning-rate schedule changes the learning rate during training to match these phases.
A popular modern approach combines two ideas:
-
Warmup.
-
Cosine decay.
Warmup
Warmup starts with a small learning rate and gradually increases it.
The idea is simple: early training is unstable because optimizer statistics have not yet become reliable. A gentle start prevents destructive updates.
Example:
learning rate = maximum learning rate × current step / warmup steps
Warmup is especially common in large language-model training.
Cosine decay
After warmup, cosine decay gradually reduces the learning rate.
The schedule follows a smooth curve:
learning rate = minimum rate + 0.5 × (maximum rate − minimum rate) × (1 + cos(progress))
In practical terms:
-
Train aggressively early.
-
Slow down smoothly.
-
Finish with careful refinement.
A meaningful optimizer comparison must keep the rest of the training recipe controlled.
The previous mistake many experiments make is comparing:
-
SGD with one learning rate.
-
AdamW with another learning rate.
-
Different schedules.
-
Different regularization.
That comparison does not tell us which optimizer is better.
The following experiment uses:
-
The same network.
-
The same generated retail-demand regression data.
-
The same warmup plus cosine schedule.
-
Different optimizer implementations only.
The example uses current PyTorch-style APIs and pins the environment so results can be reproduced. PyTorch 2.7 introduced additional compiler and performance improvements, and current releases continue the PyTorch 2.x API family used below.
Install dependencies:
python -m venv .venv source .venv/bin/activate pip install torch==2.7.1 scikit-learn==1.6.1 matplotlib==3.10.0
Create the experiment:
import math import matplotlib.pyplot as plt import torch from sklearn.datasets import make_regression from sklearn.model_selection import train_test_split from torch import nn from torch.utils.data import DataLoader, TensorDataset torch.manual_seed(42) X, y = make_regression( n_samples=5000, n_features=20, noise=20, random_state=42, ) X_train, X_test, y_train, y_test = train_test_split( X, y, test_size=0.2, random_state=42, ) X_train = torch.tensor(X_train, dtype=torch.float32) y_train = torch.tensor(y_train, dtype=torch.float32).reshape(-1, 1) train_loader = DataLoader( TensorDataset(X_train, y_train), batch_size=64, shuffle=True, ) class Net(nn.Module): def __init__(self): super().__init__() self.model = nn.Sequential( nn.Linear(20, 64), nn.ReLU(), nn.Linear(64, 32), nn.ReLU(), nn.Linear(32, 1), ) def forward(self, x): return self.model(x) def make_optimizer(model, name): if name == "sgd": return torch.optim.SGD( model.parameters(), lr=0.1, momentum=0.9, weight_decay=1e-4, ) return torch.optim.AdamW( model.parameters(), lr=0.001, weight_decay=1e-4, ) def schedule_value( step, total_steps, warmup_steps, max_lr, min_lr, ): if step < warmup_steps: return step / warmup_steps progress = (step - warmup_steps) / ( total_steps - warmup_steps ) cosine = 0.5 * ( 1 + math.cos(math.pi * progress) ) return ( min_lr / max_lr + (1 - min_lr / max_lr) * cosine ) def train(optimizer_name): model = Net() optimizer = make_optimizer( model, optimizer_name, ) loss_function = nn.MSELoss() epochs = 40 total_steps = epochs * len(train_loader) warmup_steps = 200 scheduler = torch.optim.lr_scheduler.LambdaLR( optimizer, lr_lambda=lambda step: schedule_value( step, total_steps, warmup_steps, 0.1, 0.005, ), ) history = [] for _ in range(epochs): total_loss = 0 for features, targets in train_loader: optimizer.zero_grad() predictions = model(features) loss = loss_function( predictions, targets, ) loss.backward() optimizer.step() scheduler.step() total_loss += loss.item() history.append( total_loss / len(train_loader) ) return model, history sgd_model, sgd_history = train("sgd") adamw_model, adamw_history = train("adamw") plt.plot(sgd_history, label="SGD") plt.plot(adamw_history, label="AdamW") plt.xlabel("Epoch") plt.ylabel("Training loss") plt.legend() plt.show()
The important correction in this experiment is that the scheduler scales each optimizer’s intended maximum learning rate rather than accidentally multiplying unrelated base learning rates. A comparison is only useful when the effective learning-rate policy is understood.
The result is not a universal ranking.
Sometimes SGD wins. Sometimes AdamW wins.
The lesson is that the complete training system matters more than the optimizer name alone.
One of the most important trends in recent machine-learning research is that improvements are often caused by training recipes rather than a single optimizer change.
A concrete example is the work behind large language-model scaling research. The Chinchilla study from DeepMind showed that training compute allocation, dataset size, and training duration were critical factors in achieving better model performance. The result challenged the assumption that simply making models larger was the dominant path forward.
Similarly, modern large-model training recipes emphasize:
-
Warmup duration.
-
Learning-rate decay shape.
-
Batch-size scaling.
-
Data ordering.
-
Total number of optimization steps.
The optimizer still matters, but the schedule determines how that optimizer behaves throughout training.
A mediocre optimizer with a well-designed schedule can outperform a better optimizer with poorly chosen learning-rate settings.
Machine-learning engineers often think about models as mathematical objects, but model files are also software inputs.
A surprising security lesson appeared with CVE-2025-32434. The vulnerability affected PyTorch versions up to 2.5.1 and involved torch.load with weights_only=True, a feature many users considered safer than normal pickle-based loading. Under certain conditions, loading a malicious model file could lead to remote code execution. The issue was patched in PyTorch 2.6.0.
The practical lesson:
A model checkpoint is not automatically harmless data.
Before loading external files:
-
Verify where the checkpoint came from.
-
Keep frameworks updated.
-
Avoid loading untrusted model files.
-
Treat serialized models as executable input.
Better optimization cannot compensate for an insecure training pipeline.
Before changing optimizers, check these areas first:
-
Is the learning rate appropriate for the model size?
-
Does training include warmup when instability is possible?
-
Is the decay schedule matched to the number of training steps?
-
Is weight decay configured correctly?
-
Are validation metrics improving?
-
Are you changing one variable at a time during experiments?
A common mistake is switching from SGD to AdamW because training is unstable when the actual problem is an aggressive learning rate.
Another common mistake is spending hours tuning optimizer parameters while ignoring the schedule.
SGD, Adam, and AdamW are not magic buttons. They are different tools for controlling how neural networks learn.
SGD offers simplicity and often excellent final solutions. Adam provides adaptive updates and fast progress. AdamW improves Adam’s regularization behavior and has become a practical default for many modern architectures.
But the larger lesson is this:
The optimizer starts the process. The learning-rate schedule shapes the journey.
If you are building neural networks today, treat warmup and decay strategy as core design decisions. Run controlled experiments, keep data and architecture fixed, and measure the effect of every training change.
A better schedule may improve your model more than a more fashionable optimizer.
Continue to the next lesson by running the SGD versus AdamW experiment on your own model, comparing schedules, and sharing your results with your team or community.
What experiment should you try next: vision models, transformers, or tabular neural networks?