NeuralPath // AI & ML Learning Hub

The Core Thesis: Backpropagation is only the compass—it tells your neural network which way is downhill by computing the gradient $\nabla L$ . The Optimizer is the engine that actually moves the weights, and the Learning-Rate Schedule is the throttle that controls how big your steps are at every stage of the journey so your model reaches the lowest valley quickly without crashing.

Beginner Foundations: What Does an Optimizer Do? (Batch vs. Mini-Batch SGD)

Once Backpropagation finishes computing the gradient $\mathbf{g}_t = \nabla_{\mathbf{w}} L(\mathbf{w}_t)$ , we need a rule to update our weights $\mathbf{w}_t$ . The simplest rule is Gradient Descent, which steps in the opposite direction of the gradient scaled by a step size $\eta$ (the Learning Rate)[cite: 10]:

\mathbf{w}_{t+1} = \mathbf{w}_t - \eta \cdot \mathbf{g}_t

How many training examples should we look at before taking one step? There are three ways, and modern Deep Learning overwhelmingly uses the third:

1. Full-Batch GD

Entire Dataset (

N

rows)

Computes the exact average gradient over all $10{,}000{,}000$ samples before taking $1$ step. Way too slow and cannot fit in GPU memory!

2. Pure Stochastic GD

1 Single Sample (

B = 1

)

Updates weights after every single data point[cite: 10]. Fast, but the gradient jumps around wildly and wastes GPU parallel cores.

3. Mini-Batch SGD

B = 32

2{,}048

samples

The industry standard! Computes gradients on a small batch of samples—fast on GPUs, smooth enough to steer downhill, and noisy enough to bounce out of bad traps.

Why Plain Mini-Batch SGD Struggles on Deep Neural Networks

Plain SGD has two major flaws: (1) inside a narrow, steep canyon (Ravine), it violently zig-zags back and forth across the walls while barely moving forward along the valley floor, and (2) it uses the exact same learning rate $\eta$ for every weight, even when some weights receive huge gradients and rare word embeddings receive tiny gradients[cite: 10]!

Beginner to Intermediate: Momentum (The Heavy Bowling Ball)

How do we stop SGD from zig-zagging across narrow canyons and getting stuck in flat plateaus? Imagine replacing a lightweight ping-pong ball with a heavy bowling ball rolling downhill!

SGD with Momentum keeps a moving average of past gradients called Velocity ( $\mathbf{v}_t$ ). Instead of stepping strictly along the current step's noisy gradient $\mathbf{g}_t$ , the ball blends $\beta$ (typically $0.9$ , or $90\%$ ) of its previous velocity with the new gradient:

\mathbf{v}_t = \beta \mathbf{v}_{t-1} + \mathbf{g}_t \quad \Longrightarrow \quad \mathbf{w}_{t+1} = \mathbf{w}_t - \eta \mathbf{v}_t

1. Cancels Out Side-to-Side Zig-Zags

When gradients flip-flop between $+10$ and $-10$ across steep canyon walls, the moving average $\mathbf{v}_t$ cancels those opposite signs out to $\approx 0$ , smoothing the trajectory!

2. Builds Speed Down the Valley

Along directions where gradients consistently point the same way, velocity builds up speed and rolls right through flat saddle points and tiny bumps.

⚡ Knowledge Check

Suppose a weight's gradient alternates between $+5$ and $-5$ on every step as it bounces between two steep walls, while another weight's gradient is a steady $+1$ pointing down the valley floor. What does Momentum ( $\beta = 0.9$ ) do to these two weights?

A) It dampens the +5/-5 oscillation and accelerates the steady +1 direction▼

✓ Correct!Adding

+5

and

-5

into the velocity buffer cancels the opposing oscillations out, while repeatedly adding

+1

in the same direction builds up momentum down the valley floor.

B) It amplifies the +5/-5 oscillation to +50/-50▼

✕ Incorrect.Because the signs alternate (

+5

, then

-5

), summing them with past velocity subtracts them from one another rather than amplifying them.

Intermediate: Adaptive Optimizers (RMSProp, Adam, and AdamW)

Momentum fixes the direction of the step, but what about the size of the step for each individual parameter? In a Transformer, an embedding weight for a rare word might only get a non-zero gradient once every $1{,}000$ batches, whereas a final layer weight gets huge gradients on every batch.

Adaptive Optimizers give every single weight its own custom learning rate by dividing the step by the square root of past squared gradients ( $\sqrt{v_t}$ ):

1. RMSProp (Scaling by Gradient Size)

Tracks an exponentially decaying average of squared gradients ( $v_t = \beta_2 v_{t-1} + (1-\beta_2)g_t^2$ ) and divides the learning rate by $\sqrt{v_t} + \epsilon$ . Weights with huge gradients get smaller steps; weights with tiny gradients get larger steps!

w_{t+1} = w_t - \dfrac{\eta}{\sqrt{v_t} + \epsilon} g_t

2. Adam (Momentum + RMSProp Combined!)

Adaptive Moment Estimation (2014): combines 1st Moment $m_t$ (Momentum) with 2nd Moment $v_t$ (RMSProp), plus a bias correction for early steps ( $\hat{m}_t, \hat{v}_t$ ):

w_{t+1} = w_t - \eta \dfrac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon}

The Complete Adam / AdamW Equations

m_t = \beta_1 m_{t-1} + (1 - \beta_1)g_t \qquad \text{and} \qquad v_t = \beta_2 v_{t-1} + (1 - \beta_2)g_t^2

\hat{m}_t = \dfrac{m_t}{1 - \beta_1^t}, \quad \hat{v}_t = \dfrac{v_t}{1 - \beta_2^t} \quad \Longrightarrow \quad w_{t+1} = w_t(1 - \eta \lambda) - \eta \dfrac{\hat{m}_t}{\sqrt{\hat{v}_t} + \epsilon}

Why Modern AI Uses AdamW Instead of Standard Adam (Decoupled Weight Decay $\lambda$ )

In standard Adam, $L_2$ regularization (weight decay) was added directly into the gradient $g_t$ , which meant Adam's $\sqrt{\hat{v}_t}$ denominator accidentally scaled down the regularization penalty for weights with large gradients! AdamW (Loshchilov & Hutter, 2019) fixes this by decoupling weight decay—shrinking the weights directly via $w_t(1 - \eta \lambda)$ outside the adaptive gradient step. Today, AdamW is the universal default for training Transformers and LLMs.

Optimizer	Tracks Direction ( $m_t$ )?	Adapts Step Size ( $v_t$ )?	Extra Memory per Weight	Best Use Case in AI
Plain SGD	No	No	$0$ states	Simple convex models & textbook demos
SGD + Momentum	Yes ( $\mathbf{v}_t$ )	No	$1\times$ ( $4$ bytes/param)	ResNets & classic CNN image classification
RMSProp	No	Yes ( $v_t$ )	$1\times$ ( $4$ bytes/param)	Recurrent nets & some Reinforcement Learning (DQN)
AdamW	Yes ( $m_t$ )	Yes ( $v_t$ )	$2\times$ ( $8$ bytes/param)	Default for Transformers, LLMs, ViTs, and Diffusion!

Intermediate to Advanced: Learning-Rate Schedules (Warmup & Cosine Decay)

Should you keep the learning rate $\eta$ fixed at the exact same number from step $1$ to step $100{,}000$ ? Never! Early in training, you want a larger learning rate to cross the loss landscape quickly, but near the end of training, a large learning rate will bounce around the rim of the valley without ever settling into the minimum[cite: 10].

A Learning-Rate Schedule dynamically adjusts $\eta_t$ after every step or epoch. Modern LLM training combines two stages:

Stage 1: Linear Warmup (First 1%–5% of Steps)

At step $0$ , weights are random and AdamW's variance estimate $v_t$ has seen almost no data—so a full learning rate causes violent gradient spikes! Linear Warmup starts $\eta$ near $0$ and ramps up linearly to $\eta_{\max}$ over $T_{\text{warmup}}$ steps:

\eta_t = \eta_{\max} \cdot \dfrac{t}{T_{\text{warmup}}}

Stage 2: Cosine Annealing Decay (Remaining Steps)

After reaching $\eta_{\max}$ , Cosine Annealing smoothly glides the learning rate down to $\eta_{\min}$ along a gentle cosine curve, allowing the weights to settle into a wide, flat minimum:

\eta_t = \eta_{\min} + \dfrac{1}{2}(\eta_{\max} - \eta_{\min})\left(1 + \cos\left(\dfrac{t_{\text{cur}}}{T_{\text{decay}}}\pi\right)\right)

⚡ Knowledge Check

Why do Transformer models almost always require a Linear Warmup phase during the first few hundred or thousand steps when trained with AdamW?

A) Early gradients are noisy and AdamW's second-moment v_t needs time to calibrate▼

✓ Correct!During the first few steps, random weights produce erratic gradients and AdamW's

\sqrt{\hat{v}_t}

denominator is estimated from only 1–2 batches. Ramping

\eta

up gradually prevents early attention layers from diverging!

B) Because GPUs need a few minutes to warm up their physical temperature▼

✕ Incorrect.Learning-rate warmup is purely a mathematical optimization technique to stabilize early gradient updates; it has nothing to do with GPU hardware temperature!

Advanced Production Optimization: Optimizer VRAM Math, 8-Bit Adam & Linear Scaling

At the Advanced / LLM Systems level, the optimizer is often the largest consumer of GPU memory in your entire cluster! Look at why:

1. The 2x Memory Cost of AdamW States

For every single weight parameter $w$ , AdamW must store two full 32-bit state tensors in GPU VRAM: the 1st moment $m_t$ ( $4$ bytes) and the 2nd moment $v_t$ ( $4$ bytes).

7B Model = 56 GB of VRAM just for FP32 AdamW Optimizer States!

2. 8-Bit AdamW, ZeRO, and WSD Schedules

To fit large models into memory, engineers quantize $m_t$ and $v_t$ to 8-bit integers (bitsandbytes 8-bit AdamW, saving $75\%$ optimizer RAM) or shard states across GPUs (DeepSpeed ZeRO / FSDP).

Modern LLM Schedule: WSD (Warmup - Stable - Decay)

The Batch-Size Scaling Rule (When You Double Batch Size, Adjust Learning Rate!)

When you multiply your mini-batch size by $k$ (for example, scaling from $B = 64$ to $B = 512$ across $8$ GPUs), each step averages $8\times$ more samples and you take $8\times$ fewer steps per epoch. To converge at the same speed, use the Linear Scaling Rule ( $\eta_{\text{new}} = k \cdot \eta$ ) for SGD, or Square-Root Scaling ( $\eta_{\text{new}} = \sqrt{k} \cdot \eta$ ) for AdamW!

⚡ Knowledge Check

A neural network has $1{,}000{,}000{,}000$ ( $1$ Billion) parameters. If AdamW stores its momentum ( $m_t$ ) and variance ( $v_t$ ) states in 32-bit floating point (float32 = $4$ bytes per number), how much GPU memory is required strictly for the two optimizer state tensors?

A) 8 GB (4 GB for m_t + 4 GB for v_t)▼

✓ Correct!Each of the two state tensors (

m_t

and

v_t

) has

10^9

numbers

\times 4\text{ bytes} = 4\text{ GB}

, totaling

8\text{ GB}

of VRAM just for the optimizer states (on top of the weights and gradients)!

B) 0 GB, because AdamW does not store historical states▼

✕ Incorrect.Only plain stateless SGD uses

0\text{ GB}

of state memory. AdamW maintains both first-moment (

m_t

) and second-moment (

v_t

) buffers for every trainable weight.

Visual Explanation: The Modern AI Optimization & Schedule Loop

Look at how gradients from Backpropagation are clipped, fed into AdamW to update 1st and 2nd moments with decoupled weight decay, and stepped using a Warmup + Cosine Annealing Schedule:

Python Implementation: AdamW + Linear Warmup & Cosine Decay in PyTorch

Here is a complete, runnable PyTorch implementation showing how to configure torch.optim.AdamW (separating weight-decayed matrices from non-decayed biases/LayerNorms) paired with a Linear Warmup + Cosine Annealing Schedule:

adamw_cosine_schedule.pyPython 3.11+ · PyTorch 2.x

import torchimport torch.nn as nnfrom torch.optim.lr_scheduler import LinearLR, CosineAnnealingLR, SequentialLR # 1. Build a Sample MLP Modeltorch.manual_seed(42)model = nn.Sequential(nn.Linear(16, 64), nn.GELU(), nn.Linear(64, 2)) # 2. Pro Practice: Apply Weight Decay ONLY to 2D Weight Matrices (NOT 1D Biases!)decayParams = [p for p in model.parameters() if p.dim() >= 2]noDecayParams = [p for p in model.parameters() if p.dim() < 2] optimizer = torch.optim.AdamW(    [        {"params": decayParams, "weight_decay": 0.01},        {"params": noDecayParams, "weight_decay": 0.0}    ],    lr=1e-3,    betas=(0.9, 0.999)) # 3. Build Warmup (5 steps) -> Cosine Decay (15 steps) SchedulewarmupSched = LinearLR(optimizer, start_factor=0.2, end_factor=1.0, total_iters=5)cosineSched = CosineAnnealingLR(optimizer, T_max=15, eta_min=1e-5)scheduler = SequentialLR(optimizer, schedulers=[warmupSched, cosineSched], milestones=[5]) # 4. Run 20 Optimization Steps and Track Learning RateX = torch.randn(8, 16)Y = torch.randint(0, 2, (8,))criterion = nn.CrossEntropyLoss() for step in range(1, 21):    optimizer.zero_grad()    loss = criterion(model(X), Y)    loss.backward()    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)    optimizer.step()    scheduler.step()    if step in (1, 5, 12, 20):        currentLr = optimizer.param_groups[0]["lr"]        print("Step:", step, "| LR:", round(currentLr, 6), "| Loss:", round(loss.item(), 4))

Pro Tip (Why We Exclude 1D Biases and LayerNorm from Weight Decay):

Look at Step 2 in the code above (p.dim() >= 2)! In every production LLM codebase (like nanoGPT, Llama, and HuggingFace), weight decay is applied only to 2D weight matrices and set to $0.0$ for 1D biases and LayerNorm scales. Decaying biases or LayerNorm scales hurts model accuracy without preventing overfitting!

Key Points

✓Mini-Batch Stochastic Gradient Descent (SGD) computes gradients on small batches (

32\text{--}2{,}048

samples), balancing GPU parallelism with fast updates.

✓Momentum (

\mathbf{v}_t = \beta \mathbf{v}_{t-1} + \mathbf{g}_t

) acts like a heavy ball rolling downhill, canceling out side-to-side canyon zig-zags while building speed along consistent gradient directions.

✓RMSProp adapts the learning rate for every individual parameter by dividing by the root-mean-square of past gradients (

\sqrt{v_t}

✓AdamW combines Momentum (

m_t

), RMSProp (

v_t

), bias correction, and decoupled weight decay (

\lambda

), making it the universal default optimizer for Transformers and LLMs.

✓A Linear Warmup + Cosine Annealing schedule prevents early training divergence while smoothly reducing the step size so weights settle into a deep, flat minimum.

✓Because AdamW stores two state tensors (

m_t

and

v_t

) per parameter, optimizer states consume

8\text{ bytes}

per parameter in FP32 (

2\times

the size of FP32 weights).

Common Mistakes

✕ Calling scheduler.step() BEFORE optimizer.step() in PyTorch.

In PyTorch, calling scheduler.step() before optimizer.step() skips the first value of your learning-rate schedule and triggers a user warning. Always run optimizer.step() first, then scheduler.step()!

✕ Using the same learning rate (e.g., lr = 0.1) when switching from SGD to AdamW.

Because AdamW normalizes gradients by $\sqrt{\hat{v}_t}$ , its effective step size is much larger than SGD. While SGD often uses $\eta \approx 0.01\text{--}0.1$ , AdamW typically requires a $100\times$ smaller learning rate around $\eta \approx 10^{-4}\text{ to } 10^{-3}$ .

✕ Applying weight decay to 1D bias vectors and LayerNorm / BatchNorm parameters.

Shrinking biases and normalization scales toward zero restricts the network's ability to shift and scale activations. Separate parameters into decay (2D+ matrices) and no_decay (1D biases/norms) groups.

✕ Stepping a batch-level warmup scheduler only once per epoch instead of once per batch.

If your warmup is configured for $1{,}000$ steps, you must call scheduler.step() inside the inner mini-batch loop (not the outer epoch loop) so the learning rate updates on every step!

The Big Picture

Plain Fixed-Step SGD (Bumpy & Unstable)

Same Step Size for Every Weight → Zig-Zags in Canyons → Overshoots Minimum at the End

AdamW + Warmup & Cosine Decay (Modern AI Standard)

Smooth Momentum Direction + Per-Weight Adaptive Step + Scheduled Throttle → Fast, Stable Convergence

The important conceptual shift is separating computing the gradient (Backpropagation) from using the gradient (Optimization)[cite: 10]. By smoothing gradient directions with Momentum, scaling each weight's step size with AdamW, and shaping the global step size with Warmup and Cosine Annealing, modern deep networks can train billions of parameters smoothly without diverging.

Remember: When starting almost any modern Deep Learning or Transformer project, your golden baseline recipe is AdamW(lr=3e-4, weight_decay=0.01) paired with Gradient Clipping (1.0) and a Linear Warmup + Cosine Decay schedule.

Optimizers and Learning-Rate Schedules

Beginner Foundations: What Does an Optimizer Do? (Batch vs. Mini-Batch SGD)

Beginner to Intermediate: Momentum (The Heavy Bowling Ball)

Intermediate: Adaptive Optimizers (RMSProp, Adam, and AdamW)

Intermediate to Advanced: Learning-Rate Schedules (Warmup & Cosine Decay)

Advanced Production Optimization: Optimizer VRAM Math, 8-Bit Adam & Linear Scaling

Visual Explanation: The Modern AI Optimization & Schedule Loop

Python Implementation: AdamW + Linear Warmup & Cosine Decay in PyTorch

Key Points

Common Mistakes

The Big Picture