Batch Size And Learning Rate A Quick Reference Gui

You’re training a neural network. It’s humming along, loss is dropping, and then you think: “I’ve got all this GPU memory sitting idle. Let me double the batch size to speed things up.” You change one number in your config file, restart training, and… the loss curve spikes into the stratosphere. Or maybe it doesn’t explode, but your model’s accuracy plateaus at a level that’s just… meh.

What happened? You turned one knob without touching the other. Batch size and learning rate are a coupled pair. They’re not independent dials you can tune in isolation. A larger batch gives you a more accurate estimate of the true gradient—imagine averaging the opinions of 256 people instead of 32. Because that average is more trustworthy, you can take a bigger step in the direction it points. But if you don’t actually increase your step size (the learning rate), you’re creeping along with a high-quality map but moving at a snail’s pace. Conversely, if you increase the batch size and don’t adjust the learning rate correctly, you might be taking steps that are way too large for the now-smoother landscape, and you’ll overshoot good solutions entirely.

This article is your quick-reference card for that relationship. We’ll start with the most famous heuristic—the Linear Scaling Rule—as our default starting point. Then we’ll look at the warning signs when it breaks: the generalization gap that can make large-batch models look great on paper but fail in the real world. We’ll cover the standard safety net (learning rate warmup) and a principled way to find your maximum useful batch size using the gradient noise scale. Finally, we’ll translate all this theory into a practical PyTorch scheduler cheat sheet and a decision flowchart you can screenshot.

Here’s a preview of the research landscape we’re navigating: there’s a productive tension in the field. On one side, Goyal et al. (2017) showed you can train ResNet-50 on ImageNet in under an hour by scaling batch size and learning rate together. On the other, Masters & Luschi (2018) argued that for many practical settings, small batches are often best and far less finicky. This isn’t settled science—it’s a set of heuristics, each with known failure modes. By the end of this guide, you’ll know which heuristic to reach for and when to get suspicious.

The Linear Scaling Rule: Your Default Starting Point

Let’s start with the rule you’ll see cited everywhere. In their 2017 paper “Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour,” Goyal et al. state it plainly: “When the minibatch size is multiplied by k, multiply the learning rate by k.” That’s it. If you go from a batch size of 32 to 256, you’ve multiplied by 8, so you multiply your learning rate by 8—say, from 0.01 to 0.08.

Why does this work? Think about what a batch is. Each batch gives you a noisy estimate of the true gradient—the direction the optimizer should move to reduce the loss averaged over the entire dataset. With a small batch, that estimate is noisy. One batch might point mostly uphill, the next mostly downhill. A small learning rate keeps you from careening off course because of that noise. When you increase the batch size by a factor of k, you’re averaging over k times as many examples. The variance of your gradient estimate shrinks by roughly a factor of k. The estimate is now much more reliable. You can afford to take a step that’s k times larger without the same risk of noise-induced chaos.

Let’s see what happens when you forget this rule, and what happens when you apply it correctly.

# This block is fully self-contained: it generates its own synthetic data,
# defines a simple model, and runs two training loops for comparison.

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
import matplotlib.pyplot as plt

# For reproducibility
torch.manual_seed(42)

# Generate a simple synthetic binary classification dataset
# 1000 samples, 20 features, some noise
X = torch.randn(1000, 20)
y = (X[:, 0] + X[:, 1] - X[:, 2] + torch.randn(1000) * 0.5 > 0).float()

dataset = TensorDataset(X, y)

# A tiny linear model
class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(20, 1)
    
    def forward(self, x):
        return self.linear(x).squeeze()

def train_one_config(batch_size, lr, epochs=20, label=""):
    """Train the model with given batch size and LR, return loss history."""
    loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
    model = SimpleModel()
    criterion = nn.BCEWithLogitsLoss()
    optimizer = optim.SGD(model.parameters(), lr=lr)
    
    losses = []
    for epoch in range(epochs):
        epoch_loss = 0.0
        for batch_X, batch_y in loader:
            optimizer.zero_grad()
            preds = model(batch_X)
            loss = criterion(preds, batch_y)
            loss.backward()
            optimizer.step()
            epoch_loss += loss.item()
        losses.append(epoch_loss / len(loader))
    return losses

# Baseline: batch_size=32, lr=0.01
losses_baseline = train_one_config(batch_size=32, lr=0.01, label="Baseline (bs=32, lr=0.01)")

# Mistake: scale batch size 8x but forget to scale LR
losses_no_scale = train_one_config(batch_size=256, lr=0.01, label="No scaling (bs=256, lr=0.01)")

# Correct: apply Linear Scaling Rule
losses_scaled = train_one_config(batch_size=256, lr=0.08, label="Scaled (bs=256, lr=0.08)")

# Plot the comparison
plt.figure(figsize=(10, 6))
plt.plot(losses_baseline, label="Baseline (bs=32, lr=0.01)", linewidth=2)
plt.plot(losses_no_scale, label="No scaling (bs=256, lr=0.01)", linewidth=2)
plt.plot(losses_scaled, label="Scaled (bs=256, lr=0.08)", linewidth=2)
plt.xlabel("Epoch")
plt.ylabel("Training Loss")
plt.title("Effect of the Linear Scaling Rule")
plt.legend()
plt.grid(True, alpha=0.3)
plt.show()

Look at the three curves. The baseline (batch size 32, LR 0.01) drops steadily. When we scale the batch size to 256 but forget to scale the learning rate, the loss crawls down—the optimizer is taking tiny, timid steps with a very precise map. It’s not diverging, but it’s wasting computation. Now look at the curve where we applied the Linear Scaling Rule (LR 0.08). It tracks the baseline closely. The larger steps compensate for the fewer updates per epoch, and training progresses at a comparable pace.

Now, two important caveats before you apply this everywhere. First, Goyal et al. developed this rule for large-batch distributed ImageNet training with a ResNet-50 architecture. It’s a starting point, not a guarantee for every model or dataset. Second, at very large batch sizes—think 512, 1024, or beyond—you’ll need an additional ingredient: learning rate warmup for the first few epochs. We’ll get to that in Section 4.

There’s also a complementary way to think about this, from Smith et al. (2017). They frame it as a substitution: instead of decaying the learning rate over time, you can increase the batch size. Both strategies reduce the variance of the weight updates. The intuition is the same—smoother gradients let you take more confident steps, whether that confidence comes from a smaller learning rate or a larger batch.

When Scaling Up Backfires: The Generalization Gap

So the Linear Scaling Rule keeps training stable. Great. But here’s the catch: stable training doesn’t guarantee a good model. In fact, large-batch training has a hidden cost that shows up not during training, but when you evaluate on unseen data. Researchers call this the generalization gap.

Here’s what it looks like in practice. You train one model with a small batch (say, 64) and another with a large batch (say, 1024), using the Linear Scaling Rule so both converge smoothly. At the end, the large-batch model often matches or even beats the small-batch model on the training set. But on the test set, it underperforms. Sometimes by a lot. What’s going on?

Keskar et al. (2016) provided the canonical explanation in their paper “On Large-Batch Training for Deep Learning: Generalization Gap and Sharp Minima.” They showed that large-batch methods “tend to converge to sharp minimizers.” Let’s unpack that in plain English. A “minimum” is a valley in the loss landscape—a set of weights where the loss is low. A “sharp” minimum is a narrow, steep-sided valley. If you nudge the weights just a little, the loss shoots up dramatically. A “flat” minimum is a wide, gentle basin. Small perturbations barely change the loss.

Why does this matter? Because your training data is only a sample of the real world. The loss landscape you see during training is a noisy approximation of the true loss landscape over all possible data. A sharp minimum might be perfectly placed for your training set, but it’s brittle—it doesn’t generalize to slightly different examples. A flat minimum, by contrast, is robust. Small changes in the data don’t throw it off.

Small-batch training has a secret weapon: noise. Each small batch gives a noisy gradient estimate, and that noise acts like a dithering signal. It keeps the optimizer bouncing around, making it harder to settle into those narrow, brittle valleys. The noise is a feature, not a bug—it’s an implicit regularizer that pushes the optimizer toward flat, robust solutions. When you use large batches, you get cleaner gradients, but you lose that beneficial noise. The optimizer sees a smoother, simpler landscape and happily marches into the nearest sharp minimum.

Masters & Luschi (2018) added a practical warning to this picture. In their paper “Revisiting Small Batch Training for Deep Neural Networks,” they found that “as batch size increases, the range of viable learning rates narrows, making training less stable.” Think of it this way: with small batches, you have a wide margin for error in setting your learning rate. The noise smooths over your mistakes. With large batches, you’re walking a tightrope. Set the LR a little too high, and you overshoot. A little too low, and you stall. The window of “just right” shrinks.

This is the hardest part of batch-size tuning. There’s no simple formula that says “batch size X is too large for dataset Y.” The answer depends on your architecture, your data, and even where you are in training. The generalization gap is a diagnostic signal: if your training accuracy keeps climbing while your test accuracy plateaus or degrades, your batch size might be too large.

Let’s check that against the data with a concrete example.

# This block is fully self-contained. It generates a synthetic dataset,
# trains a small model with a small batch and a large batch, and compares
# the train/test accuracy gap.

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
import numpy as np

# For reproducibility
torch.manual_seed(42)
np.random.seed(42)

# Generate a synthetic dataset with a pattern that's easy to overfit
# 2000 samples, 50 features, but only 5 are truly predictive
X = np.random.randn(2000, 50)
true_weights = np.zeros(50)
true_weights[:5] = 1.0
y = (X @ true_weights + np.random.randn(2000) * 0.8 > 0).astype(np.float32)

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42)

# Convert to PyTorch tensors
X_train_t = torch.tensor(X_train, dtype=torch.float32)
y_train_t = torch.tensor(y_train, dtype=torch.float32)
X_test_t = torch.tensor(X_test, dtype=torch.float32)
y_test_t = torch.tensor(y_test, dtype=torch.float32)

train_dataset = TensorDataset(X_train_t, y_train_t)
test_dataset = TensorDataset(X_test_t, y_test_t)

class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(50, 128),
            nn.ReLU(),
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Linear(64, 1)
        )
    
    def forward(self, x):
        return self.net(x).squeeze()

def evaluate(model, loader):
    """Return accuracy on a given data loader."""
    model.eval()
    correct, total = 0, 0
    with torch.no_grad():
        for batch_X, batch_y in loader:
            preds = (torch.sigmoid(model(batch_X)) > 0.5).float()
            correct += (preds == batch_y).sum().item()
            total += batch_y.size(0)
    return correct / total

def train_and_track(batch_size, lr, epochs=100):
    """Train a model and return train/test accuracy histories."""
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    test_loader = DataLoader(test_dataset, batch_size=256, shuffle=False)
    
    model = MLP()
    criterion = nn.BCEWithLogitsLoss()
    optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9)
    
    train_accs, test_accs = [], []
    for epoch in range(epochs):
        model.train()
        for batch_X, batch_y in train_loader:
            optimizer.zero_grad()
            loss = criterion(model(batch_X), batch_y)
            loss.backward()
            optimizer.step()
        
        train_accs.append(evaluate(model, train_loader))
        test_accs.append(evaluate(model, test_loader))
    
    return train_accs, test_accs

# Small-batch training
small_train_acc, small_test_acc = train_and_track(batch_size=32, lr=0.01, epochs=100)

# Large-batch training with Linear Scaling Rule
large_train_acc, large_test_acc = train_and_track(batch_size=512, lr=0.16, epochs=100)

# Plot the comparison
fig, axes = plt.subplots(1, 2, figsize=(14, 5))

axes[0].plot(small_train_acc, label="Train (bs=32)", linewidth=2)
axes[0].plot(small_test_acc, label="Test (bs=32)", linewidth=2)
axes[0].set_title("Small Batch (32)")
axes[0].set_xlabel("Epoch")
axes[0].set_ylabel("Accuracy")
axes[0].legend()
axes[0].grid(True, alpha=0.3)

axes[1].plot(large_train_acc, label="Train (bs=512)", linewidth=2)
axes[1].plot(large_test_acc, label="Test (bs=512)", linewidth=2)
axes[1].set_title("Large Batch (512)")
axes[1].set_xlabel("Epoch")
axes[1].set_ylabel("Accuracy")
axes[1].legend()
axes[1].grid(True, alpha=0.3)

plt.suptitle("Generalization Gap: Small vs. Large Batch")
plt.tight_layout()
plt.show()

# Print the final gap
small_gap = small_train_acc[-1] - small_test_acc[-1]
large_gap = large_train_acc[-1] - large_test_acc[-1]
print(f"Small batch (32)  - Final train acc: {small_train_acc[-1]:.3f}, test acc: {small_test_acc[-1]:.3f}, gap: {small_gap:.3f}")
print(f"Large batch (512) - Final train acc: {large_train_acc[-1]:.3f}, test acc: {large_test_acc[-1]:.3f}, gap: {large_gap:.3f}")

Look at the numbers printed at the end. The large-batch model often squeezes out a higher training accuracy—it’s fitting the training set extremely well. But the test accuracy tells a different story. The gap between train and test accuracy is wider for the large-batch model. That’s the generalization gap in action. The large-batch optimizer found a solution that’s highly specific to the training data, while the small-batch noise pushed the optimizer toward a solution that transfers better to unseen examples.

The Warmup Safety Net (And Why It’s Not the Only Answer)

If you’ve scaled up your batch size and applied the Linear Scaling Rule, you might be watching the first few epochs with your coffee in hand. Suddenly, the loss spikes. The model diverges before it even has a chance to learn. What now?

This is where learning rate warmup comes in. At the very start of training, your weights are random. The gradients are chaotic—different layers are learning at wildly different rates, and the loss surface is a mess. If you hit that chaos with a large learning rate (which you need for your large batch), the optimizer can get flung off to a part of the parameter space from which it never recovers. Warmup is the standard band-aid: start the learning rate near zero and linearly increase it to your target value over the first few epochs. This gives the optimizer time to settle into a reasonable region before you hand it the full step size.

Here’s the standard PyTorch recipe. We combine LinearLR for the warmup phase with CosineAnnealingLR for the main training phase, using SequentialLR to chain them together.

# This block is fully self-contained. It demonstrates a warmup + cosine decay
# scheduler chain on a toy model and dataset.

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
from torch.optim.lr_scheduler import LinearLR, CosineAnnealingLR, SequentialLR
import matplotlib.pyplot as plt

# Generate toy data
torch.manual_seed(42)
X = torch.randn(1000, 20)
y = (X[:, 0] + X[:, 1] - X[:, 2] > 0).float()
dataset = TensorDataset(X, y)
loader = DataLoader(dataset, batch_size=256, shuffle=True)

# Define model, optimizer, and target LR
model = nn.Linear(20, 1)
optimizer = optim.SGD(model.parameters(), lr=0.0)  # Start at 0; schedulers will control it
criterion = nn.BCEWithLogitsLoss()

# Scheduler setup
warmup_epochs = 5
total_epochs = 50
peak_lr = 0.08  # Scaled for batch_size=256 from a baseline of 0.01 at bs=32

# Warmup: linearly increase LR from 0 to peak_lr over warmup_epochs
warmup_scheduler = LinearLR(optimizer, start_factor=0.01, end_factor=1.0, total_iters=warmup_epochs)

# Main scheduler: cosine decay from peak_lr to near-zero over remaining epochs
main_scheduler = CosineAnnealingLR(optimizer, T_max=total_epochs - warmup_epochs, eta_min=1e-6)

# Chain them: warmup first, then cosine decay
scheduler = SequentialLR(optimizer, schedulers=[warmup_scheduler, main_scheduler], milestones=[warmup_epochs])

# Track LR and loss for visualization
lrs = []
losses = []

for epoch in range(total_epochs):
    epoch_loss = 0.0
    for batch_X, batch_y in loader:
        optimizer.zero_grad()
        preds = model(batch_X).squeeze()
        loss = criterion(preds, batch_y)
        loss.backward()
        optimizer.step()
        epoch_loss += loss.item()
    
    scheduler.step()  # Update LR after each epoch
    lrs.append(optimizer.param_groups[0]['lr'])
    losses.append(epoch_loss / len(loader))

# Plot LR schedule and loss
fig, ax1 = plt.subplots(figsize=(10, 5))

color = 'tab:blue'
ax1.set_xlabel('Epoch')
ax1.set_ylabel('Learning Rate', color=color)
ax1.plot(lrs, color=color, linewidth=2)
ax1.tick_params(axis='y', labelcolor=color)

ax2 = ax1.twinx()
color = 'tab:red'
ax2.set_ylabel('Training Loss', color=color)
ax2.plot(losses, color=color, linewidth=2, alpha=0.7)
ax2.tick_params(axis='y', labelcolor=color)

plt.title('Learning Rate Warmup + Cosine Decay')
plt.grid(True, alpha=0.3)
plt.show()

Look at the blue curve. For the first 5 epochs, the learning rate climbs from near-zero to 0.08. This is the warmup phase—the safety net. The loss drops rapidly during this period because the model is learning the easy, large-scale patterns without the risk of overshooting. After epoch 5, the cosine decay takes over, smoothly reducing the learning rate to near zero. This lets the optimizer settle into a good minimum without bouncing around.

Now, here’s the interesting part. Warmup is the standard answer, but it’s not the only answer. You et al. (2020) proposed an alternative called CLARS (Complete Layer-wise Adaptive Rate Scaling) in their paper “Large Batch Optimization for Deep Learning: Training BERT in 76 minutes.” Their method scales the learning rate differently for each layer based on the ratio of the layer’s weight norm to its gradient norm. They claim it “outperforms gradual warmup technique by a large margin” for extremely large-batch training. Different layers have different gradient scales, so a single global warmup schedule is a blunt instrument. Layer-wise adaptation is more surgical.

The takeaway for a practitioner? Warmup is a reliable default. If you’re scaling to batch sizes above ~256, add a 5-10 epoch linear warmup and you’ll avoid most early-training catastrophes. But treat it as a starting heuristic, not settled science. The field is still actively debating the best way to stabilize large-batch training, and layer-wise methods like CLARS (and LARS before it) are gaining ground.

The Noise Scale: A Principled Way to Find Your Max Batch Size

Up to this point, we’ve been working with heuristics. The Linear Scaling Rule says “multiply LR by k.” The generalization gap warns “don’t go too large.” Warmup says “ease into it.” But none of these tell you the largest useful batch size for your specific model, dataset, and training stage. For that, we need a more principled approach.

McCandlish et al. (2018) gave us one in their paper “An Empirical Model of Large-Batch Training.” They introduced a measurable quantity called the gradient noise scale. In plain English, it’s a number that tells you how noisy your gradient estimates are at a given batch size. More importantly, it predicts the batch size beyond which you’re just wasting computation—bigger batches give you the same gradient quality, not better.

Think of it this way. Imagine you’re trying to estimate the average height of people in a city by measuring random samples. If you measure 10 people, your estimate is noisy—measure another 10 and you might get a different answer. If you measure 1,000 people, your estimate is much more stable. But at some point—say, 10,000 people—adding more people to your sample barely changes the estimate. You’ve hit the point of diminishing returns. The gradient noise scale is the batch-size equivalent of that point. It’s the batch size at which the variance of your gradient estimate stops meaningfully decreasing.

Here’s what makes this powerful: the noise scale changes over the course of training. It tends to increase, which means the optimal batch size grows over time. This connects beautifully to an idea from Smith et al. (2017): instead of decaying the learning rate, you can increase the batch size during training. Both strategies reduce the noise in the weight updates, but increasing the batch size has the added benefit of better computational efficiency.

The practical heuristic: if doubling your batch size doesn’t meaningfully reduce the variance of your gradient estimates, you’ve hit the noise scale ceiling. Stop scaling up. You’re just burning GPU cycles for no statistical benefit.

Estimating the noise scale requires measuring gradient statistics during training, which adds implementation complexity. The original paper provides an open-source implementation, but the core idea is conceptually simple: compare the variance of gradients computed on small batches to the variance of gradients computed on larger batches, and find the point where the variance stops shrinking proportionally.

# This block is fully self-contained. It demonstrates the conceptual idea
# behind gradient noise scale estimation on a toy model.
# NOTE: This is a simplified illustration, not the full McCandlish et al. method.

import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset
import numpy as np
import matplotlib.pyplot as plt

torch.manual_seed(42)

# Generate toy data
X = torch.randn(2000, 10)
y = (X[:, 0] * 2 + X[:, 1] - X[:, 2] + torch.randn(2000) * 0.5 > 0).float()
dataset = TensorDataset(X, y)

model = nn.Linear(10, 1)
criterion = nn.BCEWithLogitsLoss()

# We'll estimate gradient variance at different batch sizes
batch_sizes = [8, 16, 32, 64, 128, 256, 512]
grad_variances = []

for bs in batch_sizes:
    loader = DataLoader(dataset, batch_size=bs, shuffle=True)
    
    # Collect gradient samples from multiple batches
    grad_samples = []
    for i, (batch_X, batch_y) in enumerate(loader):
        if i >= 50:  # Collect 50 gradient samples per batch size
            break
        model.zero_grad()
        preds = model(batch_X).squeeze()
        loss = criterion(preds, batch_y)
        loss.backward()
        
        # Extract the gradient of the weight parameter as a flat vector
        grad = model.weight.grad.clone().detach().flatten().numpy()
        grad_samples.append(grad)
    
    # Compute the variance of the gradient across samples
    grad_samples = np.array(grad_samples)
    # Average variance across all parameter dimensions
    avg_variance = np.mean(np.var(grad_samples, axis=0))
    grad_variances.append(avg_variance)

# Plot: how does gradient variance change with batch size?
plt.figure(figsize=(10, 6))
plt.plot(batch_sizes, grad_variances, 'o-', linewidth=2, markersize=8)
plt.xscale('log')
plt.yscale('log')
plt.xlabel('Batch Size (log scale)')
plt.ylabel('Gradient Variance (log scale)')
plt.title('Gradient Variance vs. Batch Size\n(Conceptual Noise Scale Illustration)')
plt.grid(True, alpha=0.3, which='both')

# Annotate the diminishing returns region
plt.axvline(x=128, color='red', linestyle='--', alpha=0.7, label='Diminishing returns region')
plt.legend()
plt.show()

print("Gradient variance by batch size:")
for bs, var in zip(batch_sizes, grad_variances):
    print(f"  Batch size {bs:3d}: variance = {var:.6f}")

Look at the printed variances. As batch size increases, the gradient variance drops—that’s expected. But notice the shape of the curve. At small batch sizes, doubling the batch size cuts the variance dramatically. At larger batch sizes, the benefit shrinks. The point where the curve starts to flatten is your noise scale ceiling. Beyond that batch size, you’re spending more computation for negligible improvement in gradient quality. This is a more principled answer to “how big can I go?” than the Linear Scaling Rule alone.

Quick-Reference Card: PyTorch Schedulers at a Glance

All this theory is useful, but when you’re staring at your training script, you need to know which scheduler class to reach for. Here’s your quick-reference card—a skimmable table of the PyTorch schedulers you’ll actually use, with one-line descriptions and when to deploy each.

Here’s a single, annotated code block that puts it all together: a model, an optimizer, and a scheduler chain (warmup + cosine decay) with comments mapping each choice back to the concepts we’ve covered.

# This block is fully self-contained. It defines a complete training setup
# with warmup + cosine decay, annotated with the reasoning behind each choice.

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
from torch.optim.lr_scheduler import LinearLR, CosineAnnealingLR, SequentialLR

# ------------------------- DATA -------------------------
# Toy dataset — replace with your own
torch.manual_seed(42)
X = torch.randn(5000, 32)
y = (X[:, 0] + X[:, 1] - X[:, 2] > 0).float()
dataset = TensorDataset(X, y)
loader = DataLoader(dataset, batch_size=256, shuffle=True)

# ------------------------- MODEL -------------------------
model = nn.Sequential(
    nn.Linear(32, 64),
    nn.ReLU(),
    nn.Linear(64, 1)
)

# ------------------------- OPTIMIZER -------------------------
# Baseline LR for bs=32 would be 0.01. We scaled bs by 8x, so LR is 0.08.
# (Linear Scaling Rule — Goyal et al. 2017)
base_lr = 0.01
batch_size_multiplier = 8  # 32 -> 256
scaled_lr = base_lr * batch_size_multiplier  # 0.08

optimizer = optim.SGD(model.parameters(), lr=scaled_lr, momentum=0.9, weight_decay=1e-4)

# ------------------------- SCHEDULER CHAIN -------------------------
warmup_epochs = 5
total_epochs = 90

# Warmup: ease into the large LR to avoid early instability
# (Standard heuristic for batch sizes > ~256; see also CLARS — You et al. 2020)
warmup = LinearLR(
    optimizer,
    start_factor=0.01,   # Start at 1% of target LR
    end_factor=1.0,       # Reach full target LR
    total_iters=warmup_epochs
)

# Main schedule: cosine decay — smooth, works well with large batches
# (Modern default; replaces step decay in many vision recipes)
main = CosineAnnealingLR(
    optimizer,
    T_max=total_epochs - warmup_epochs,  # Remaining epochs after warmup
    eta_min=1e-6                          # Don't decay all the way to zero
)

# Chain: warmup first, then cosine decay
scheduler = SequentialLR(
    optimizer,
    schedulers=[warmup, main],
    milestones=[warmup_epochs]  # Switch from warmup to main after epoch 5
)

# ------------------------- TRAINING LOOP -------------------------
criterion = nn.BCEWithLogitsLoss()

for epoch in range(total_epochs):
    model.train()
    epoch_loss = 0.0
    for batch_X, batch_y in loader:
        optimizer.zero_grad()
        preds = model(batch_X).squeeze()
        loss = criterion(preds, batch_y)
        loss.backward()
        optimizer.step()
        epoch_loss += loss.item()
    
    scheduler.step()  # Update LR at the end of each epoch
    
    if epoch % 10 == 0 or epoch == total_epochs - 1:
        current_lr = optimizer.param_groups[0]['lr']
        print(f"Epoch {epoch:3d} | Loss: {epoch_loss/len(loader):.4f} | LR: {current_lr:.6f}")

This is the “I just need to get something working” block. Copy it, replace the toy data with your own, adjust base_lr and batch_size_multiplier to match your setup, and you’ve got a training loop that respects the Linear Scaling Rule, protects against early instability with warmup, and uses a modern cosine decay schedule.

Putting It All Together: A Decision Flowchart

You’ve got a training job to run. You have constraints: GPU memory, dataset size, a deadline. What batch size should you start with, and how should you set the learning rate? Here’s the synthesis of everything we’ve covered, in decision-tree form.

START
  │
  ├─► Step 1: Pick a starting batch size
  │      └─ Does it fit in GPU memory and give reasonable throughput?
  │         ├─ Yes → Use 32 or 64 as a safe default (Masters & Luschi 2018)
  │         └─ No → Reduce until it fits
  │
  ├─► Step 2: Do you NEED to scale up?
  │      (Distributed training? Wall-clock time constraints?)
  │      │
  │      ├─ No → Keep your batch size. Tune LR with a simple cosine schedule.
  │      │
  │      └─ Yes → Apply the Linear Scaling Rule:
  │            Multiply LR by the same factor as batch size (Goyal et al. 2017)
  │            │
  │            ├─► Step 3: Is your new batch size > ~256?
  │            │      │
  │            │      ├─ No → Skip warmup. Go to Step 4.
  │            │      │
  │            │      └─ Yes → Add LR warmup for 5-10 epochs
  │            │            (Standard heuristic; CLARS is an alternative — You et al. 2020)
  │            │
  │            ├─► Step 4: Monitor the generalization gap
  │            │      Train/test accuracy diverging?
  │            │      │
  │            │      ├─ No → Your batch size is probably fine. Go to Step 6.
  │            │      │
  │            │      └─ Yes → Your batch size may be too large
  │            │            (Sharp minima — Keskar et al. 2016)
  │            │            Consider reducing batch size or adding explicit regularization
  │            │
  │            ├─► Step 5 (Optional): Estimate the gradient noise scale
  │            │      Have engineering bandwidth?
  │            │      │
  │            │      ├─ No → Skip. Use heuristics from Steps 1-4.
  │            │      │
  │            │      └─ Yes → Find your true max useful batch size
  │            │            (McCandlish et al. 2018)
  │            │            Stop scaling when gradient variance stops shrinking
  │            │
  │            └─► Step 6: Pick a scheduler
  │                  ├─ CosineAnnealingLR → Strong default
  │                  ├─ OneCycleLR → Worth trying for faster convergence
  │                  └─ StepLR → Simple baseline for comparison
  │
  END

This flowchart isn’t a rigid algorithm—it’s a structured way to apply the heuristics we’ve discussed. Start with a reasonable batch size. If you must scale up, apply the Linear Scaling Rule. Add warmup if you’re going large. Watch for the generalization gap. If you have the resources, measure the noise scale to find your ceiling. And pick a modern scheduler as your default.

Recap and What’s Next

Let’s land the plane with the key takeaways from this guide:

1. Batch size and learning rate are a coupled pair. Scale them together using the Linear Scaling Rule (Goyal et al., 2017) as your starting heuristic: multiply the learning rate by the same factor as the batch size.

2. Large batches can hurt generalization. They tend to converge to sharp minima that don’t transfer well to unseen data (Keskar et al., 2016). Small-batch noise is an implicit regularizer—it’s a feature, not a bug.

3. LR warmup is a standard safety net for large-batch training. Start the learning rate near zero and linearly increase it over the first 5-10 epochs. Active research (You et al., 2020) is exploring layer-wise alternatives like CLARS.

4. The gradient noise scale gives a principled upper bound on useful batch size (McCandlish et al., 2018). When doubling your batch size stops meaningfully reducing gradient variance, you’ve hit the ceiling—stop scaling up.

5. PyTorch’s scheduler ecosystem gives you all the tools to implement these ideas in a few lines of code. CosineAnnealingLR is a strong default; SequentialLR lets you chain warmup and decay cleanly.

In the next part of this series, we’ll explore how optimizer choice—SGD vs. Adam vs. LAMB—interacts with batch size scaling, and when the Linear Scaling Rule breaks down entirely. Different optimizers have different scaling properties, and what works for SGD with momentum might not work for adaptive methods.

Check Your Understanding

Test yourself against these six questions, moving from simple recall to creative application.

Remember: State the Linear Scaling Rule in one sentence.

Understand: Explain why a larger batch size allows a larger learning rate. What property of the gradient estimate changes?

Apply: You increase your batch size from 64 to 512. Your current learning rate is 0.1. What learning rate does the Linear Scaling Rule suggest? What should you add to the first few epochs of training?

Analyze: Your large-batch model achieves 99% training accuracy but only 92% test accuracy, while your small-batch model achieves 97% training and 94% test. What phenomenon is likely at play, and what paper first characterized it?

Evaluate: A colleague claims, “The Linear Scaling Rule means you should always use the largest batch size your GPU can handle.” Provide two counterarguments from the research literature.

Create: Design a learning rate schedule for a ResNet-50 trained on ImageNet with a batch size of 2048. Specify the scheduler classes, the warmup duration, the peak learning rate, and the decay strategy. Justify each choice with a reference to a paper or established heuristic.

Apply What You Learned

# This block is a minimal CIFAR-10 training loop. It is fully self-contained.
# Modify the lines marked with *** to apply the concepts from this guide.

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from torch.optim.lr_scheduler import CosineAnnealingLR, LinearLR, SequentialLR

# ------------------------- DATA -------------------------
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

train_dataset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
test_dataset = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)

# *** CHANGE THESE TWO VALUES ***
BATCH_SIZE = 64       # <-- Double this (e.g., 64 -> 128, or 128 -> 256)
LEARNING_RATE = 0.01  # <-- Apply Linear Scaling Rule here

# *** ADD WARMUP IF BATCH_SIZE > ~256 ***
# If your new batch size is large, add a LinearLR warmup phase before the main scheduler.

train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=2)
test_loader = DataLoader(test_dataset, batch_size=256, shuffle=False, num_workers=2)

# ------------------------- MODEL -------------------------
class SimpleCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(3, 32, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(32, 64, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )
        self.fc = nn.Sequential(
            nn.Flatten(),
            nn.Linear(64 * 8 * 8, 256),
            nn.ReLU(),
            nn.Linear(256, 10)
        )
    
    def forward(self, x):
        x = self.conv(x)
        return self.fc(x)

# ------------------------- TRAINING SETUP -------------------------
model = SimpleCNN()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=LEARNING_RATE, momentum=0.9)

# *** MODIFY THE SCHEDULER ***
# Add warmup if needed, then use CosineAnnealingLR or your chosen schedule.
scheduler = CosineAnnealingLR(optimizer, T_max=50, eta_min=1e-6)

# ------------------------- TRAINING LOOP -------------------------
EPOCHS = 50

def evaluate(model, loader):
    model.eval()
    correct, total = 0, 0
    with torch.no_grad():
        for images, labels in loader:
            outputs = model(images)
            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    return correct / total

train_losses = []
test_accuracies = []

for epoch in range(EPOCHS):
    model.train()
    epoch_loss = 0.0
    for images, labels in train_loader:
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        epoch_loss += loss.item()
    
    scheduler.step()
    
    avg_loss = epoch_loss / len(train_loader)
    test_acc = evaluate(model, test_loader)
    train_losses.append(avg_loss)
    test_accuracies.append(test_acc)
    
    if epoch % 10 == 0 or epoch == EPOCHS - 1:
        print(f"Epoch {epoch:2d} | Loss: {avg_loss:.4f} | Test Acc: {test_acc:.4f} | LR: {optimizer.param_groups[0]['lr']:.6f}")

print(f"\nFinal Test Accuracy: {test_accuracies[-1]:.4f}")