Machine Learning Fundamentals: Optimizers, Scaling Laws, and Clustering

Scenario

As a Machine Learning Engineer working on foundational model pre-training and downstream clustering pipelines, you are tasked with making principled architectural, optimization, and resource-allocation decisions. Walk through the following three technical problems in detail.

Part 1: Optimization Mechanics (Adam vs. AdamW)

Consider training a large transformer model. In standard implementations of Adam with L2 regularization, the weight decay penalty is implemented by adding lambda * theta_t directly to the gradient vector g_t before computing the first (m_t) and second (v_t) moment estimates:

1. Write out the update equations for standard Adam with L2 regularization versus decoupled weight decay (AdamW).
2. Explain mathematically why adding L2 regularization to g_t fails to act as true weight decay for parameters that have large historical gradient magnitudes versus small historical gradient magnitudes.
3. What empirical issue does this cause when training deep models with varied parameter update frequencies (e.g., rare embedding tokens vs. frequent attention projection weights)?

Part 2: Compute-Optimal Scaling Laws

You are allocated a fixed compute budget of C = 6 * 10^23 FLOPs for pre-training a dense decoder-only transformer. Assume standard Chinchilla scaling rules where FLOPs are approximated by C approx 6 * N * D, where N is the number of non-embedding parameters and D is the number of training tokens.

1. Under the Chinchilla optimal compute allocation (where N and D scale approximately in equal proportion, i.e., N proportional to C^0.5 and D proportional to C^0.5), calculate the optimal parameter count N and token count D given that a balanced model with N = 10^9 parameters requires roughly D = 20 * 10^9 tokens (a 1:20 ratio).
2. If an engineering constraint requires serving the model at low latency with a strict memory footprint limit of N <= 1.5 * 10^9 parameters, how would you reallocate the remaining compute budget? What trade-offs in downstream loss and training stability occur when scaling tokens far past the compute-optimal frontier (inference-optimal regime)?

Part 3: Clustering Mechanics & Derivations

Analyze the relationship between the Expectation-Maximization (EM) algorithm for Gaussian Mixture Models (GMMs) and the classic k-means clustering algorithm.

1. Formulate the GMM posterior probability gamma_ik = P(z_i = k | x_i, theta) where all covariance matrices are isotropic and identical: Sigma_k = sigma^2 * I for all clusters k, with uniform cluster priors pi_k = 1 / K.
2. Show what happens to the posterior responsibility gamma_ik in the limit as sigma^2 -> 0. Prove how the EM updates for the cluster means reduce exactly to Lloyd's algorithm for k-means.