AdamW › Visualizing AdamW's Regularization Behavior on Embeddings

# Simulate embedding training with heterogeneous update frequencies
np.random.seed(42)

# Simulate 100 "embedding" parameters with varying update frequencies
n_embeddings = 100
n_steps = 1000

# Simulate gradient frequencies: some embeddings updated every step,
# some only occasionally
update_probs = np.concatenate(
    [
        np.ones(20),  # 20 common tokens: updated every step
        np.ones(30) * 0.3,  # 30 medium tokens: updated 30% of steps
        np.ones(50) * 0.05,  # 50 rare tokens: updated 5% of steps
    ]
)

# Simulate gradient magnitudes proportional to update frequency
gradient_scale = np.sqrt(
    update_probs
)  # larger gradients for more frequent tokens

# Track v_hat (second moment estimate) for each embedding
beta2 = 0.999
v = np.zeros(n_embeddings)
v_hat_history = np.zeros((n_steps, n_embeddings))

for t in range(1, n_steps + 1):
    # Simulate which embeddings get gradients this step
    active = np.random.rand(n_embeddings) < update_probs
    g = np.where(
        active, gradient_scale + np.random.randn(n_embeddings) * 0.1, 0.0
    )
    v = beta2 * v + (1 - beta2) * g**2
    v_hat = v / (1 - beta2**t)
    v_hat_history[t - 1] = v_hat

# Final v_hat values
final_v_hat = v_hat_history[-1]

# Compute effective regularization for each embedding
lam = 0.01
eps = 1e-8

effective_reg_l2 = lam / (np.sqrt(final_v_hat) + eps)
effective_reg_adamw = np.full(n_embeddings, lam)

# Sort by update probability for cleaner visualization
sort_idx = np.argsort(update_probs)[::-1]