Stochastic Gradient Descent (SGD’s) Frequency Bias and How Adam Fixes It

Modern language models are trained on data with extremely uneven token distributions. A small number of words appear in almost every sentence, while many rare but meaningful tokens occur only occasionally. This creates a hidden optimization challenge: parameters associated with common tokens receive constant gradient updates, while parameters tied to rare tokens may go hundreds or thousands of steps without receiving any meaningful signal. Under standard Stochastic Gradient Descent (SGD), every parameter uses the same learning rate, so frequently updated weights converge quickly while rare-token weights often remain close to their random initialization.

This is where Adam’s adaptive optimization becomes important. While Adam is commonly described as SGD with momentum, its most impactful feature in practice is variance normalization. Adam tracks the historical gradient statistics for each parameter independently and automatically adjusts update sizes based on how often reliable gradient information has been observed. Parameters that rarely receive updates end up getting proportionally larger effective learning rates, allowing underrepresented features to learn much faster than they would under vanilla SGD.

To demonstrate this behavior concretely, we build a controlled NumPy experiment using a six-token vocabulary spanning four orders of magnitude in frequency — from tokens appearing in nearly every batch to tokens appearing only 0.1% of the time. We train the same linear model twice, once with SGD and once with Adam, while keeping all target weights identical. By comparing final parameter values, non-zero gradient counts, and Adam’s effective learning rates for each token, we can directly observe how adaptive optimization compensates for frequency imbalance in real training dynamics.

Setting up the dependencies

We begin by constructing a deliberately simplified training environment that isolates a single factor: token frequency. The vocabulary contains six tokens ranging from extremely common words like “the” to very rare tokens like “thalweg,” with appearance probabilities spanning four orders of magnitude. Every token is assigned the same ground-truth importance — the correct weight for all tokens is set to 1.0 — so the experiment removes semantic complexity and focuses entirely on how often each parameter receives gradient updates.

Each training sample is represented as a sparse binary vector indicating which tokens are present in the batch. The target value is simply the sum of the active token weights plus a small amount of noise. We then train a small linear model using this synthetic dataset. Because gradients are only computed for tokens that appear in a batch, rare tokens naturally receive far fewer updates than common ones. This setup creates a clean environment for observing how SGD and Adam behave under highly imbalanced gradient exposure.

Copy CodeCopiedUse a different Browser

import numpy as np import matplotlib.pyplot as plt import matplotlib.gridspec as gridspec

np.random.seed(42)

Copy CodeCopiedUse a different Browser

TOKENS = ["the", "model", "embedding", "tokenization", "xenobiotic", "thalweg"]

Appearance probability per batch -- spans 4 orders of magnitude

FREQ = np.array([0.95, 0.60, 0.20, 0.05, 0.005, 0.001]) TRUE_W = np.ones(6) # all weights should reach 1.0

N_STEPS = 3000 LR = 0.05 BATCH_SIZE = 32 # samples per step

def sample_batch(batch_size): """ Each sample is a sparse binary feature vector. Token i appears in the sample with probability FREQ[i]. Target y = x @ TRUE_W + small noise. """ X = (np.random.rand(batch_size, 6) 1e-9).astype(float) history[t] = w.copy()

return history, grad_counts

ADAM

We now train the same model using Adam to observe how adaptive optimization changes the learning dynamics. Alongside the model weights, Adam maintains two additional running statistics for every parameter: a momentum estimate mmm, which tracks the average direction of past gradients, and a variance estimate vvv, which tracks the average magnitude of squared gradients. Before applying updates, both statistics are bias-corrected to account for their initialization at zero.

Copy CodeCopiedUse a different Browser

def train_adam(n_steps, lr, batch_size, beta1=0.9, beta2=0.999, eps=1e-8): w = np.zeros(6) m = np.zeros(6) v = np.zeros(6) history = np.zeros((n_steps, 6)) v_history = np.zeros((n_steps, 6)) # track variance accumulation

for t in range(1, n_steps + 1): X, y = sample_batch(batch_size) error = X @ w - y grad = (X.T @ error) / batch_size

m = beta1 * m + (1 - beta1) * grad v = beta2 * v + (1 - beta2) * grad ** 2

m_hat = m / (1 - beta1 t) v_hat = v / (1 - beta2 t)

w -= lr * m_hat / (np.sqrt(v_hat) + eps)

history[t-1] = w.copy() v_history[t-1] = v_hat.copy()

return history, v_history

Running both

With both optimizers implemented, we train the model twice under identical conditions — once using SGD and once using Adam. Each optimizer sees the same synthetic data distribution, uses the same initialization, learning rate, batch size, and training duration. This ensures that any difference in the final behavior comes entirely from the optimization strategy itself rather than changes in the dataset or model architecture.

Copy CodeCopiedUse a different Browser

print("Training SGD...") sgd_history, sgd_grad_counts = train_sgd(N_STEPS, LR, BATCH_SIZE)

print("Training Adam...") adam_history, adam_v_history = train_adam(N_STEPS, LR, BATCH_SIZE)

print()

Measuring the failure

We now evaluate how well each optimizer learned the token weights after training. Since every token has the same true target weight of 1.0, the ideal outcome is that all learned weights also end close to 1.0 regardless of token frequency. Along with the final weights, we also measure how many training steps each token actually received a non-zero gradient. This helps us directly compare optimization quality against gradient exposure frequency.

The results clearly show the difference between SGD and Adam. For common tokens, both optimizers learn the correct weights successfully because these tokens appear in almost every batch. But for rare tokens, SGD struggles badly. “xenobiotic” only receives gradients in about 15% of training steps and its weight stops around 0.53 instead of 1.0. The rarest token, “thalweg,” receives gradients in only 3.4% of steps and SGD barely learns it at all, ending near 0.15. Adam, however, keeps both rare-token weights close to the correct value despite receiving the same sparse gradient signals.

Copy CodeCopiedUse a different Browser

sgd_final = sgd_history[-1] adam_final = adam_history[-1]

print("=" * 62) print(f"{'Token':6} {'SGD w':>8} {'Adam w':>8} {'SGD grads':>10}") print("-" * 62) for i, token in enumerate(TOKENS): sgd_err = abs(sgd_final[i] - TRUE_W[i]) adam_err = abs(adam_final[i] - TRUE_W[i]) flag = " ← fails" if sgd_err > 0.3 else "" print( f"{token:6.3f} {sgd_final[i]:>8.4f} " f"{adam_final[i]:>8.4f} {int(sgd_grad_counts[i]):>10}{flag}" ) print() print(f"True weight for all tokens: {TRUE_W[0]:.1f}") print()

How many steps did each token get a non-zero gradient?

print("Non-zero gradient steps out of", N_STEPS) for i, token in enumerate(TOKENS): pct = sgd_grad_counts[i] / N_STEPS * 100 bar = "█" * int(pct / 2) print(f" {token:", color="#c0392b", lw=1.2), bbox=dict(boxstyle="round,pad=0.3", facecolor="#fff0f0", edgecolor="#c0392b", alpha=0.85) )

── 2. Final weight error bar chart ───────────────────────────

ax2 = fig.add_subplot(gs[0, 2]) ax2.set_facecolor(BG)

x = np.arange(6) w_sgd = sgd_final w_adam = adam_final width = 0.35

bars_sgd = ax2.bar(x - width/2, np.abs(w_sgd - TRUE_W), width, color="#c0392b", alpha=0.85, label="SGD error") bars_adam = ax2.bar(x + width/2, np.abs(w_adam - TRUE_W), width, color="#2980b9", alpha=0.85, label="Adam error")

ax2.set_xticks(x) ax2.set_xticklabels([t[:8] for t in TOKENS], rotation=30, ha="right", fontsize=8) ax2.set_ylabel("|learned w − true w|", fontsize=9) ax2.set_title("Final Weight Error\n(lower = better)", fontsize=11, color=DARK) ax2.legend(fontsize=8) ax2.spines[["top", "right"]].set_visible(False)

── 3. Adam weight trajectories ───────────────────────────────

ax3 = fig.add_subplot(gs[1, :2]) ax3.set_facecolor(BG) ax3.axhline(1.0, color=DARK, lw=1, ls="--", alpha=0.3, label="True weight = 1.0")

for i, (token, color) in enumerate(zip(TOKENS, TOKEN_COLORS)): ax3.plot(steps, adam_history[:, i], color=color, lw=1.8, label=f"{token} (freq={FREQ[i]:.3f})")

ax3.set_title("Adam -- Weight Trajectories\nRare tokens converge via variance normalization", fontsize=11, color=DARK) ax3.set_xlabel("Training Step", fontsize=9) ax3.set_ylabel("Learned Weight", fontsize=9) ax3.legend(fontsize=8, loc="right") ax3.set_ylim(-0.3, 1.6) ax3.spines[["top", "right"]].set_visible(False)

ax3.annotate( "Rare tokens converge\ndespite sparse gradients", xy=(N_STEPS * 0.95, adam_history[-1, 5]), xytext=(N_STEPS * 0.60, 0.3), fontsize=8.5, color="#27ae60", arrowprops=dict(arrowstyle="->", color="#27ae60", lw=1.2), bbox=dict(boxstyle="round,pad=0.3", facecolor="#f0fff4", edgecolor="#27ae60", alpha=0.85) )

── 4. Effective LR vs frequency ─────────────────────────────

ax4 = fig.add_subplot(gs[1, 2]) ax4.set_facecolor(BG)

ax4.scatter(FREQ, effective_lr, c=TOKEN_COLORS, s=120, zorder=5, edgecolors="white", lw=1.5) for i, token in enumerate(TOKENS): ax4.annotate(token, (FREQ[i], effective_lr[i]), textcoords="offset points", xytext=(6, 4), fontsize=7.5, color=TOKEN_COLORS[i])

ax4.axhline(LR, color=DARK, lw=1, ls="--", alpha=0.4) ax4.text(0.5, LR * 1.05, f"Nominal LR = {LR}", fontsize=8, color=DARK, alpha=0.6)

ax4.set_xscale("log") ax4.set_yscale("log") ax4.set_xlabel("Token Frequency (log scale)", fontsize=9) ax4.set_ylabel("Adam Effective LR lr/√v̂ (log scale)", fontsize=9) ax4.set_title("Adam's Automatic Equalizer\nRare tokens get amplified LR", fontsize=11, color=DARK) ax4.spines[["top", "right"]].set_visible(False)

plt.savefig("sgd_vs_adam.png", dpi=150, bbox_inches="tight", facecolor=BG) plt.show()

Check out the Codes with Notebook. Also, feel free to follow us on Twitter and don’t forget to join our 150k+ ML SubReddit and Subscribe to our Newsletter. Wait! are you on telegram? now you can join us on telegram as well.

Need to partner with us for promoting your GitHub Repo OR Hugging Face Page OR Product Release OR Webinar etc.? Connect with us

The post Stochastic Gradient Descent (SGD’s) Frequency Bias and How Adam Fixes It  appeared first on MarkTechPost.