A2 · Hyperparameter Scaling：从一维扫描到联合 scaling

做什么

题面来自官方材料（本地原件在 _local/，官方链接见「材料」）。 编号是原件的编号（notebook 是 cell 序号、PDF 是页码），方便跟官方题面对号。 本项目不提供答案（D53）。

第 1 页

Assignment 2: Hyperparameter Scaling Assignment 2: Hyperparameter Scaling CS 312 Instructions. This assignment is ungraded, but doing the assignment will significantly assist you in obtaining partial credit on your quiz. Submit a single PDF writeup to Gradescope by the deadline in the course schedule. Please limit your PDF submission to 20 pages excluding the bonus question. For each problem, your deliverable will be two parts: an analysis of the provided experiments together with the new experiments requested in each problem, and a short (around one paragraph) synthesis of your interpretation of the experiments. Use both components to check and develop your under- standing. Remember that you will be able to cite your assignment in your quiz to justify your answer and so it is in your best interest to have a detailed and well-organized set of experiments that you fully understand in order to get partial credit on the quiz. In the final, optional problem, you can develop an interesting outcome prediction problem akin to our quizzes, where you will submit a short code diff from our codebase (under 100 lines) and a prediction task, framed as either a single-answer or a multiple-answer multiple-choice question. If we choose to use your submitted quiz question in any way (including in some modified form) on the actual quiz, you will receive bonus points on the quiz equivalent to the value of the quiz question. The example-question answer key is at the end of this handout. Overview Training neural networks requires us to reason about scale: we train larger models, use more data, and increase batch sizes. Hyperparameters that work well at one scale may not work well at another. In this assignment, you will study hyperparameter scaling: how to adjust hyperparameters as the model and training scale change. This assignment addresses the following questions: • How should hyperparameters like learning rate and weight decay scale as we train on more data? • Can initialization and learning-rate scaling let us reuse a single optimal learning rate as we scale? • How does scaling the batch size affect the hyperparameters? The goal is to learn when experiments at one scale can guide training at another, and how to test the limits of those scaling rules. As part of this process, we will carefully study toy examples: simple optimization tasks for which we have solid theory and predictions. Toy examples are important as they form the basis of many theory-based intuitions that researchers have, and understanding both the power and limitations of these models is critical to understanding the extent to which we should trust theory in deep learning optimization. 1

第 2 页

Assignment 2: Hyperparameter Scaling Our model. We use a Llama-style architecture with the following baseline dimensions: 8 layers, hidden size 512, 8 attention heads, a 4,096-token vocabulary, and about 35M parameters. The model uses query/key (QK) normalization, which applies root mean square normalization to queries and keys within each attention head. The baseline context length is 1,024 tokens. RoPE is unscaled and training uses a prefix of the full source dataset shuffled with seed 42. We train on fresh data and evaluate validation loss. At the baseline model size and a batch size of 64 sequences, training on 614.4M tokens takes approximately 15 minutes on a single H100. Provided experiments Use the provided source sweeps for Problems 1 and 2. From the starter repository root, list the run links for a subpart with: uv run python -m experiments.a2.provided_sweeps --part P1a --links Replace P1a with P1b, P1d, P1e, or P2a for the other supplied runs. The command lists each run’s token budget, optimizer, schedule, learning rate, weight decay, and W&B link. Use the command without --links to read the measurements offline. Estimated experiments The suggested grids total about 166 new runs, including 55 five-step tests; supplied runs and optional explorations (including Muon) are excluded. Reuse matching runs across parts. Stay within the 48 GPU- hour budget: plan for at most 36 hours of experiments and reserve 12 for debugging and reruns. Time a representative run at each model size before launching sweeps; shrink the grids as needed and report any untested comparisons. Default training settings Unless a problem specifies otherwise, use the following defaults for language-model experiments. Write η for the peak learning rate (LR) and λ for the weight-decay coefficient (WD). Each problem specifies its token budget. The default model is d8, with the training settings below. Setting Default Optimizer AdamW Peak learning rate η .003 Weight decay λ .1 Adam moment coefficients (β1, β2) (.9, .95) Numerical stabilizer ϵ 10−8 Batch size 64 sequences LR schedule Linear decay after warmup Warmup First 1% of training updates Gradient clipping Global gradient norm capped at 1 Model and data seeds 42 2

第 3 页

Assignment 2: Hyperparameter Scaling AdamW: Adam with decoupled weight decay In this assignment, many of our problems will make use of AdamW. For completeness, we define what we mean by the AdamW optimizer here. AdamW combines an adaptive gradient update with a separate shrinkage of the weights. At update t ≥0, let gt be the gradient supplied to Adam for a weight matrix Wt. With moment coefficients β1, β2 and numerical stabilizer ϵ > 0, define gradient and squared-gradient moments mt, vt, their bias-corrected versions bmt, bvt, and the Adam direction Ut below. Start from m0 = v0 = 0; operations are elementwise. mt+1 = β1mt + (1 −β1)gt, vt+1 = β2vt + (1 −β2)g2 t , bmt+1 = mt+1 1 −βt+1 1 , bvt+1 = vt+1 1 −βt+1 2 , Ut = bmt+1 p bvt+1 + ϵ . With current LR ηt, AdamW updates the matrix as Wt+1 = Wt −ηtUt −ηtλWt. We follow the common practice of excluding embeddings, normalization parameters, and biases from weight decay. Learning-rate scaling with AdamW Optimal learning rates depend on other hyperparameters, such as batch size. How should they change as we scale up the data or model? We might imagine that there are many possible answers to this question. Since the model is the same, and the optimal learning rate is often fairly stable, we might imagine the optimal learning rate stays similar. Or, we might look at the theory of convex optimization and conclude that, based on the convergence theorems that the right thing to do would be to decrease learning rates at rates of O(1/ √ T) or O(1/T) for T steps. This is a simple, but surprisingly complex and consequential topic. Even large scale, serious empirical efforts on this question leads to disagreement, as you can see in this example taken from the StepFun learning rate scaling paper. Examples of various scaling rules for hyperparameters identified across the literature. Note the broad disagreement between credible papers and organizations. We will now study how learning rate scaling works in our d8 system and see if we can build intuition for what types of scaling rules are plausible. Throughout, let D denote the training-token budget. 3

第 4 页

Assignment 2: Hyperparameter Scaling Problem 1: When can training on less data guide large scale runs? (9 new runs) Use the baseline model and default training settings above. Vary LR and the training-token budget, keeping the other settings fixed. For parts (a) and (b), analyze the provided W&B sweeps or the starter’s CSV. Fit the curves and make the plots yourself. (a) Learn from small training budgets (provided runs). At 153.6M, 307.2M, and 614.4M tokens, investigate how final validation loss changes with peak LR using the provided measurements at {.0015, .003, .006}. For each scale, fit a curve mapping peak LR to loss, and use this to estimate the optimal LR. Use these optima to propose a scaling rule in the training-token budget. Plot the fitted optima and your rule. What trend do you observe? (b) Predict before increasing the budget (provided runs). Record your predicted optimal LRs at 1.2288B, 1.8432B, and 2.4576B tokens before inspecting the provided measurements at these budgets. Then fit their loss–LR curves and compare your predictions with the fitted optima. Does the trend from the small budgets continue? (c) Compare and test two scaling rules (5 new runs). Fit one learning-rate power law over all six source budgets from parts (a) and (b), and another over only 1.2288B, 1.8432B, and 2.4576B tokens. Use each to predict the optimal LR at 4.9152B. Which prediction do you expect to give the lower loss, and why? Train at both predicted LRs and at target-grid LRs {.0015, .003, .006}. Compare your predictions with the fitted target optimum and report their measured loss gaps relative to the best sampled LR in the target grid. Did you find gains from using the smaller model fits? Think about bias-variance tradeoffs in fitting scaling laws, and how far you can extrapolate your scaling law fits. (d) Hyperball versus AdamW (recommended: 4 new runs). The shape and predictability of learn- ing rate scaling might vary significantly across optimizers, and one of the motivations for new and exotic optimization methods is that they may scale more predictably than the standard ones. To this end we will explore Hyperball, which is a simple modification of the well-studied Adam optimizer. Hyperball uses an Adam direction but constrains each selected weight matrix to keep its initial norm. For a linear-layer weight matrix Wt, let Ut be Adam update direction. Write ∥·∥F for the Frobenius norm (the square root of the sum of squared entries). For nonzero norms, the Hyperball update is Wt+1 = ∥Wt∥F ∥eWt+1∥F eWt+1 eWt+1 = Wt −ηt ∥Wt∥F ∥Ut∥F Ut. We propose an update eWt+1 by scaling the standard Adam update relative to the weight norm, and enforce that the norm remains the same after the step. Since hyperball already controls the norms, Hyperball uses no weight decay. Embeddings, normalization parameters, and biases use ordinary Adam. The supplied runs use peak LR (.000656/.00630)η ≈.1041η for these Adam parameters, where η is the listed Hyperball peak LR. Apply the same warmup and decay multiplier to both groups, and keep this LR ratio fixed in your target runs. Use the 24 supplied Hyperball runs (P1d) at 153.6M, 307.2M, and 614.4M tokens, with eight LRs at each budget. Fit an optimum at each of these three budgets and fit a power law to those optima. Compare its trend with AdamW fitted to the same three budgets; no larger Hyperball source budgets are supplied. Record a prediction at 1.2288B tokens before running 4

第 5 页

Assignment 2: Hyperparameter Scaling that target. Test the predicted LR and three nearby LRs of your choice (up to four new runs, reusing duplicates). Compare the prediction with the best sampled target loss. Does the optimal LR vary more regularly with budget? Discuss the uncertainty of extrapolating from only three source budgets. (e) Explore schedule variants (optional). Thus far, we have studied optimal LR scaling for a single schedule. Do our findings transfer across learning rate schedules? Can we directly transfer the optimal LR from one schedule to another? Explore different learning-rate schedules, such as cosine decay or a constant learning rate, using the same methodology. 5

第 6 页

Assignment 2: Hyperparameter Scaling Joint learning-rate and weight-decay scaling At this point, you should have a general sense of how learning rates behave as a function of scale, and we’ve seen some intriguing signs in the hyperball experiment that learning rates and the norm of the parameters might be related (at least in our transformer model). In this problem, we will dig deeper into the phenomenon and try to understand how learning rates and weight decay interact with each other. Problem 2: How should learning rate and weight decay scale together? (4 new runs) (a) Let LR and WD vary together (provided runs). Analyze the provided W&B sweeps for the following LR–WD grids. Our runs contain all combinations of peak LR and WDs given in the following table, Tokens (B) Peak LRs WDs .1536 {.00075, .0015, .003} {.4, 1.6, 6.4} .3072 {.0015, .003, .006} {.1, .4, .8} .6144 {.0015, .003, .006} {.1, .2, .4} 1.2288 {.0015, .003, .006} {.1, .2, .4} As we have two hyperparameters, we will have to extend our earlier scaling rules, and consider a bivariate functional form for the loss. Let x = log η and y = log λ. For each scale (token budget), fit a local quadratic bL(x, y) to approximate the final val loss using the coefficients a, b, c, d, e, f: bL(x, y) = a + bx + cy + dx2 + exy + f y2, The cross term exy allows LR and WD to interact. For each budget, draw its contours and mark the fitted optimum at each budget. If we make two plots, one for LR and WD where we plot optimal LR or WD as a function of scale, what do we see? How does the LR trend compare with the one in Problem 1? Can either coordinate be described well by a single power law? Compare the best measured validation losses from joint tuning with those from Problem 1 at the same budgets. Plot both against tokens, along with their difference. How does the benefit of tuning WD depend on the training budget? (b) From coupling to a joint scaling law (no new runs). At each budget, let (η∗, λ∗) be the fitted optimal pair from part (a). On its contour plot, draw your quadratic fit and observed data with the two axes being peak LR and WD (with log-log scale). Can increasing one hyperparameter compensate for decreasing another? Draw a line where ηλ = η∗λ∗on this plot, and see how it relates to the contour plot. Now plot η∗λ∗against tokens and compare it with the separate LR and WD trends. Does the product vary more regularly than either coordinate? Measure the R2 of each of these different scaling fits and figure out which of these variables have the tightest power law scaling. 6

第 7 页

Assignment 2: Hyperparameter Scaling (c) Test the joint rule (4 new runs). Consider a policy with peak LR .003 at every budget (note that this is not asking for a constant LR schedule, we just fix the same peak LR across scales). Our goal will be to see if we can “rescue” this LR policy by adjusting our weight decay. 1. Fit. Use the product law from part (b) fitted on {.1536, .3072, .6144, 1.2288}B tokens. 2. Test. Validate at 2.4576B tokens, excluded from the fit. At peak LR .003, test the WD given by the predicted product divided by .003. Use linear decay for all runs. 3. Compare. Use two baselines, both trained for 2.4576B tokens: (i) reuse the LR–WD pair with the lowest validation loss in the 153.6M sweep, without retuning; (ii) fix peak LR at .003, test WD {.05, .1, .2}, and take the best validation loss. Report the prediction’s loss gap to both baselines. Does the product rule give an effective training recipe? 7

第 8 页

Assignment 2: Hyperparameter Scaling Example question Difficulty: Medium Can we postpone weight decay? Use the default training settings at 153.6M tokens with peak LR .0015 and WD λ = .4, with initialization seeds 42–44 and fixed data order. We modify the standard training process so that weight decay does not happen at every step, but instead applies a stronger weight decay every K steps. Let ηt be the LR actually used at update t (the optimizer’s LR before calling optimizer.step, not the LR after advancing its scheduler). With positive warmup the first update uses zero LR, and the schedule reaches zero only after the final update. Ordinary AdamW scales down the weights by rt = 1 −ηtλ before applying the gradient update. Compare three schedules: • Apply weight decay every update, as usual. • Apply it only at the end of each block of two updates. • Apply it only at the end of each block of four updates. For a block ending at update t with length k, the delayed schedule applies a stronger weight decay every k steps ert = t ∏ j=t−k+1 (1 −ηjλ). The usual schedule gives mean final validation loss approximately 3.209 across the three seeds. Predict the mean final validation loss for each delayed schedule. Round each mean to the nearest multiple of 0.01. • Weight decay every two updates: . • Weight decay every four updates: . 8

第 9 页

Assignment 2: Hyperparameter Scaling How does batch size affect the hyperparameters? We now have had some taste of bivariate hyperparameters scaling (in our learning rate and weight decay tuning problem). We might wonder whether if there is a better way, where we can mathematically derive the right scaling relationships between hyperparameters. To this end, we will study how batch size, learning rate, and scale all interact with each other. Batch size is an ideal variable to study, since the relationship between batch size, learning rates and convergence rates are very well studied and understood as part of the theory of online optimization. We will make use of a well-known toy model called the noisy quadratic model, and see how this gives us learning-rate rules that (might) carry over to language model training in the wild. Background: The noisy quadratic model A noisy quadratic model (NQM) replaces neural-network training with a quadratic loss and noisy gradi- ents. There are two key properties of the noisy quadratic that we will need to vary and understand. The curvature of the quadratic controls the optimization difficulty of the problem, and controls how quickly different directions can be optimized; the gradient noise represents variation among training examples and informs whether larger batches are helpful. We will follow Guodong Zhang et al. (2019), Which Algorithmic Choices Matter at Which Batch Sizes? Insights From a Noisy Quadratic Model, which uses this simple model to explain how learning rate, momentum, and the optimizer affect batch-size scaling. We will investigate two conclusions and the conditions under which they hold: 1. Optimal LR can follow a power law in batch size. Let B be batch size and η∗(B) the LR minimizing expected final loss at a fixed total number of training examples. Write a candidate rule as η∗(B) ≈cBp, with fitted constants c and p. For SGD in the small-step regime, the NQM predicts approximately linear scaling, p = 1: doubling batch size calls for doubling LR. Of course, there will come a point where this prediction breaks down, as an arbitrarily large learning rate will lead to divergence. We will test this scaling rule, and identify the range of parameters where it is useful. 2. Momentum can help more at large batches than at small batches. The classical NQM prediction is that, in the noise-dominated small-batch regime, SGD momentum mainly rescales the effective LR and offers little extra benefit after LR tuning. At larger batches, reduced noise can allow momen- tum to accelerate convergence. We will test this prediction at a finite training budget and perform comparisons with Adam. Our first problem will make use of a simple, two dimensional quadratic model for you to play with. 2D quadratic toy model. Let w = (w1, w2)⊤be the parameters and fix the curvature matrix H = diag(1, 10). The loss and its exact gradient are f (w) = 1 2w⊤Hw = 1 2(w2 1 + 10w2 2), ∇f (w) = Hw. To define the noisy part of this quadratic, we will consider our gradients to be noisy version of the true gradient. For an update with batch size B, we define the corresponding minibatch gradient to be gt = Hwt + ϵt, ϵt ∼N (0, I2/B) independently across updates, where I2 is the two-dimensional identity matrix. This noise model comes from the naive assumption that the gradients are averages of i.i.d. items, and thus we expect variance to scale as O(1/B). 9

第 10 页

Assignment 2: Hyperparameter Scaling Optimizers. In this problem, we will use a simplified optimizer model with constant LR η and no weight decay. For the optimizer itself, we will study Stochastic gradient descent (SGD) (defined as wt+1 = wt −ηgt), Adam with ϵ = 10−8, β1 = 0.9, and β2 = 0.95, and RMSProp, which is Adam with β1 = 0. Problem 3.1: What does this model predict about batch size? (Run locally with CPUs) Consider our 2D toy quadratic model above. You will perform stochastic optimization on this toy model using the noisy gradients and compare how the hyperparameters change the final loss. Concretely, let N = 8, 192 be the total dataset size, B be the batch size (which we specify below), and w0,1 ∼N (0, 1) and w0,2 ∼N (0, .1) as initialization. Compare expected final loss LN(B, η) = E[ f (wN/B)], estimated by averaging over independent initializations and noise (using at least 512 samples). (a) SGD. Sweep LR at batches B ∈{1, 2, 4, 8, 16, 32, 64} and fit a power law to the estimated optimal LRs. Predict the optimal LR at B = 256, then test the prediction with an LR sweep at that batch. Extend the batch range through 512 and plot optimal LR and best final loss after tuning LR. Over what range does the rule work? Think about how reduced gradient noise competes with fewer updates at fixed N. (b) Adaptive updates. Repeat the experiment with RMSProp and Adam, tuning LR separately. Compare the source exponents and the losses at predicted versus tuned LRs at batch 256. Do these optimizers follow the same rule, and does a similar exponent imply similar transfer accu- racy? State your predictions for the language model setting. How should optimal LR change with batch size, and over what range do you expect that prediction to hold? (c) Explore the assumptions. Replace the gradient-noise covariance I2 by σ2I2 for σ ∈{1, 10, 100, 300}. Repeat the RMSProp and Adam source fits and batch-256 tests. Plot the fitted exponent against σ. Does changing the ratio of gradient signal to noise change the scaling rule? Relate your observations to the contributions to the second moment. As a curvature comparison, repeat the experiments with the scalar loss f (w) = w2/2, initial w0 ∼N (0, 2), and gradient-noise variance σ2/B. Which conclusions are sensitive to the curva- ture model? (d) Momentum at small and large batches. Return to H = diag(1, 10) and noise covariance I2. For SGD with momentum, initialize s0 = 0 and use st+1 = µst + gt, wt+1 = wt −ηst+1, where µ is the momentum coefficient. At small batch B = 16 and large batch B = 256, compare µ = 0 and .9, tuning LR separately. For Adam, keep β2 = .95 and jointly sweep β1 and LR at each batch (hint: you should include β1 = 0, 0.9 and other values closer to one). Plot the best loss at each β1 after tuning LR. How does the benefit of momentum depend on batch size? Compare the learning curves of the best (β, LR) pair with the best no-momentum baseline. Do you expect these observations to carry over to language model training? 10

第 11 页

Assignment 2: Hyperparameter Scaling Batch size scaling in a language model Do the scaling rules and momentum effects from the NQM carry over to language-model training? Use the default model and training settings at 614.4M tokens, varying batch size and the specified hyperpa- rameters. Here B is the number of sequences per optimizer update. Use microbatches of min(B, 64) sequences and report the actual number of tokens processed. Use a linear learning-rate schedule with warmup over the first 1% of optimizer updates. The batch-switch example below specifies its own schedule in terms of tokens processed. Problem 3.2: What should scale when batch size changes? (suggested: 48 new runs) (a) Tune LR with weight decay (9 new runs). Set WD to .1 and sweep peak LR at B ∈{8, 16, 32, 64}. Plot the loss–LR curves, fitted optimal LR, and best measured loss against batch size. How do the optimal LR and performance change with batch size? (b) Compare two scaling hypotheses (suggested: 27 new runs). The NQM can also give pre- dictions of how weight decay and batch might jointly scale together. Design experiments at B ∈{8, 16, 32, 64} to test two hypotheses: i. Fix WD at .1 across batches and scale the optimal peak LR. ii. Fix peak LR at .0015 across batches and scale the optimal WD. Use the source results to predict configurations at B = 128 and B = 256, then test which hypothesis transfers better. Discuss whether your findings agree with the NQM. (c) Ablate momentum (suggested: 12 new runs). At B = 8 and B = 256, hold each batch’s best measured LR–WD pair and β2 = .95 fixed. Design experiments including a β1 = 0 control to test how the effect of momentum changes with batch size. Compare with our earlier problem when we retuned LR. (d) What does the NQM explain? (no new runs). Which predictions from your quadratic experi- ments agree with your language-model experiments, and which do not? Give a short prediction–measurement table for LR–batch-size scaling, weight decay, and momen- tum. Include the measured target losses supporting your conclusions. Distinguish a prediction that fails from an effect the NQM does not model. What would you change in the NQM to explain one of the disagreements? Further reading. Power Lines: WD–LR–batch coupling; Step Law: optimal LR and batch-size scaling; Hyperball: control of weight and update norms. 11

第 12 页

Assignment 2: Hyperparameter Scaling Example question Difficulty: Hard Double batchsize during training Train the d8 language model 614.4M tokens with AdamW, initial batch size 64, peak LR .0015, and WD 1.6. Let ηbase(D) denote the learning rate after processing D tokens: warm up over the first 1% of tokens, then decay linearly to zero at 614.4M tokens. After processing exactly 307.2M tokens, double the batch size to 128. Reach the midpoint with 4,687 batches of 64 sequences followed by one batch of 32; this gives 4,688 updates before the switch. Then use 2,343 batches of 128 sequences and a final batch of 96, giving 2,344 updates after the switch. Compute each update’s mean loss using its actual batch size. Keep Adam’s accumulated moments and both β values unchanged. For the remaining half of training, consider five policies: Policy Change at the midpoint LR after the change WD after the change A Neither LR nor WD ηbase(D) 1.6 B Multiply LR by √ 2 √ 2 ηbase(D) 1.6 C Multiply LR by 2 2 ηbase(D) 1.6 D Multiply WD by √ 2 ηbase(D) 1.6 √ 2 E Multiply WD by 2 ηbase(D) 3.2 All runs process the same ordered tokens. Evaluate ηbase(D) at the number of tokens consumed before each update, including the partial batches. Changing the batch size must not restart warmup. Which policy do you predict will give the lowest mean final validation loss across 3 seeds (42,43,44)? Choose one and explain how changing LR or WD affects training when there are ap- proximately half as many optimizer updates per token after the switch. 12

第 13 页

Assignment 2: Hyperparameter Scaling Part 4: Initialization, feature learning, and transfer So far, we have relied on a scaling-law based approach to finding good hyperparameters at large scale. However, even a theory-guided approach of this kind may seem unsatisfying. Why does the optimal hyperparameter have to change at all? If we design the architecture and optimizer right, perhaps we could tune a hyperparameter once, and use that same optimal hyperparameter everywhere. This is the promise of various hyperparameter transfer schemes, and we will go in-depth into a popular kind of transfer scheme known as µP. The goal and the abc parameterization We want to tune hyperparameters on a small model and reuse them as we scale width and depth. It will turn out that we can patch up existing architectures and optimizers to have reasonable hyperparameter transfer by adjusting some small details in how we parametrize initialization, output scale, and learning rates. Following Everett et al., consider a one linear layer with fan-in n, p outputs, and input feature x ∈Rn. We adjust the standard form of this layer with three exponents (a, b, c) which modify the forward multiplier, initialization scale, and learning rate. Let σ and base LR η be scale independent constants (i.e. our hyperparameters that are stable with scale), z = n−aWx, W ∈Rp×n, (W0)ij iid ∼N (0, σ2n−2b), ηW = ηn−c. Two principles: stability and feature learning We compare per-coordinate scales using RMS norms: ∥v∥2 RMS = 1 n n ∑ j=1 v2 j , ∥A∥2 RMS = 1 pn ∑ i,j A2 ij (A ∈Rp×n). For output vectors in Rp, replace n by p in the vector RMS definition. Take width n →∞at a fixed training step t. For a fixed probe input x with ∥x∥RMS = Θ(1), let zt = n−aWtx. Here O(1) means bounded with width and Θ(1) means bounded and nonvanishing. Stability requires usable initial hidden features and bounded changes: ∥z0∥RMS = Θ(1), ∥zt −z0∥RMS = O(1). Initial readout logits may instead vanish, so their initial scale need only be O(1). The Θ(1) initialization requirement above applies to hidden features. Feature learning requires a nonvanishing change at some fixed step: ∥zt −z0∥RMS = ∥n−a(Wt −W0)x∥RMS = Θ(1). If this change scales as n−r, then r < 0 is unstable, r = 0 gives feature learning, and r > 0 gives a vanishing change. Deriving Kaiming initialization from stability We first use stability at initialization to derive Kaiming initialization. Since (W0)ij iid ∼N (0, σ2n−2b) and W0 is independent of the input x, the cross terms vanish: E[z2 0,i | x] = n−2a n ∑ j=1 E[(W0)2 ij]x2 j = σ2n−2(a+b)∥x∥2 2. 13

第 14 页

Assignment 2: Hyperparameter Scaling For stability, we want all hidden features, including x and z, to have Θ(1) RMS norm at initialization. This motivates choosing a + b = 1/2: a + b = 1/2 =⇒ E[∥z0∥2 RMS | x] = σ2∥x∥2 RMS. Kaiming initialization for ReLU chooses a = 0, b = 1/2, and σ2 = 2. See He et al. (2015), Section 2.2 for the derivation of the constant factor 2. Kaiming initialization does not address feature learning during training, which motivates µP theory. To understand it, we introduce notation for alignment between the weight update and the input feature. Alignment between the weight update and the input feature Let ∆W := Wt −W0. The training pair of interest is (∆W, x). By definition, a weight update gives ∆z = n−a∆Wx. For nonzero ∆W and x, define R(∆W, x) = ∥∆Wx∥RMS ∥∆W∥RMS∥x∥RMS , ∥∆z∥RMS = n−aR(∆W, x)∥∆W∥RMS∥x∥RMS. The ratio R represents the alignment between the parameter update and the inputs. Note how the alignment allows us to simplify the change in features in terms of their constituent parts (change in parameters and change in inputs). Now let us get some intuitions for the order of magnitude scale for the feature updates in two special cases. Assume ∆W ∈Rp×n, where n is the fan-in. No alignment: ∆W has iid centered Gaussian entries and is independent of x, so R(∆W, x) ≈√n. Suppose the update entries ∆Wij are independent N (0, τ2) variables, independent of x. For j̸ = r, E[∆Wij∆Wir] = 0, so E[(∆Wx)2 i | x] = ∑ j,r xjxrE[∆Wij∆Wir] = τ2 ∑ j x2 j = nτ2∥x∥2 RMS. Thus the RMS scale of ∆Wx is √n τ∥x∥RMS, while that of ∆W is τ. Gaussian rotational symmetry gives the exact RMS scale of the ratio: E[R(∆W, x)2 | x]
1/2 = √ n. The equality holds in RMS over random updates; an individual update need not have exactly this ratio. Full alignment: ∆W has rank one and x lies along its right singular vector, so R(∆W, x) = n. A gradient built from its input has this form. For t = 1, take one SGD step from W0 on a single-example loss ℓ(z), where z = n−aW0x and δ = ∂ℓ/∂z is the loss gradient with respect to the output. With learning rate ηW, the update is ∆W = −ηWn−aδx⊤, ∆Wx = −ηWn−aδ∥x∥2 2. For a nonzero update, ∆W has rank one and x/∥x∥2 is its right singular vector with nonzero singular value. Each row is proportional to x, so we have ∥∆W∥RMS = ηWn−a∥δ∥RMS∥x∥RMS, R(∆W, x) = ∥x∥2 2 ∥x∥2 RMS = n. This attains the Cauchy–Schwarz bound R(∆W, x) ≤n. We see that √n and n are both reasonable scalings. We will see later which of these is empirically a better fit. 14

第 15 页

Assignment 2: Hyperparameter Scaling Alignment between the readout head and its input feature change Let V0 ∈Rn×q be the initial readout matrix in the transposed convention, where the number of classes q does not scale with n. For a last-layer feature change ∆x ∈Rn, the initial readout contributes ∆zfeature = n−aV⊤ 0 ∆x. For nonzero V0 and ∆x, define S(V0, ∆x) = ∥V⊤ 0 ∆x∥RMS ∥V0∥RMS∥∆x∥RMS , ∥∆zfeature∥RMS = n−aS(V0, ∆x)∥V0∥RMS∥∆x∥RMS. There are two reference cases: • No alignment: iid centered Gaussian V0 independent of ∆x gives E[S(V0, ∆x)2 | ∆x]
1/2 = √n. • Full alignment: the columns of V0 are proportional to ∆x, giving S(V0, ∆x) = n. Deriving µP by layer type For this derivation, we assume that both full-alignment conditions hold simultaneously: 1. Full alignment between the weight update and the input feature. 2. Full alignment between the initial readout head and the input feature change. Problem 4.0 asks you to derive the scaling rules for all four combinations of these alignment assumptions. Hidden and readout share the weight-update constraint. Maximal-update parameterization (µP) keeps updates nonvanishing and signals bounded. Write an Adam update as ∆W = −ηn−cU, where U = bm/( √ bv + ϵ) is the normalized direction. For Adam, we can assume ∥U∥RMS = Θ(1).1 Define α by R(U, x) = Θ(nα). For an order-one input, ∥n−a∆Wx∥RMS = ηn−a−cR(U, x)∥U∥RMS∥x∥RMS = Θ(ηnα−a−c). For feature learning, we want ∥n−a∆Wx∥RMS to stay at a constant, nonzero scale as n →∞. Under full alignment (α = 1), this gives our first equation: a + c = 1. The feature-change constraint depends on the layer’s shape. After one update from (W0, x0), the changes ∆W and ∆x give ∆z = n−a∆Wx0 | {z } weight update + W0∆x | {z } feature change + ∆W∆x | {z } interaction
. We have already shown that the direct weight-update term satisfies ∥n−a∆Wx0∥RMS = Θ(1). For ∥∆x∥RMS = Θ(1), the interaction obeys ∥n−a∆W∆x∥RMS ≤n1−a∥∆W∥RMS∥∆x∥RMS = O(ηn1−a−c) = O(η). It remains to control W0∆x. Assume W0 has p output channels. To express this term’s dependence on n, we use the Gaussian matrix norm bound. Let ∥A∥op = supu̸=0 ∥Au∥2/∥u∥2 be the operator norm. The Gaussian matrix norm bound gives ∥W0∥op = 1Adam can be viewed as signSGD with momentum, using sign( bm), together with variance adaptation; see Balles and Hennig (2018), Dissecting Adam. 15

第 16 页

Assignment 2: Hyperparameter Scaling O(n−b(√p + √n)) with high probability. Converting to RMS norms bounds every ∆x with RMS Θ(1), including changes that depend on W0: ∥n−aW0∆x∥RMS ≤n−a rn p ∥W0∥op∥∆x∥RMS = O  n−a−b √ n + n √p  . 1. Hidden matrices: both dimensions scale together, p ∝n. For p ∝n, the feature-change bound is O(n1/2−a−b). The initialization requirement a + b = 1/2 already makes this O(1), even when ∆x depends on W0. No stronger initialization scaling is needed. 2. Readout: full readout–feature-change alignment gives a + b = 1. For the readout, p = q is fixed. Write V0 = W⊤ 0 ∈Rn×q for the initial readout in the transposed convention. Full alignment with the learned feature change means S(V0, ∆x) = Θ(n). With ∥V0∥RMS = Θ(n−b) and ∥∆x∥RMS = Θ(1), this gives ∥n−aV⊤ 0 ∆x∥RMS = n−aS(V0, ∆x)∥V0∥RMS∥∆x∥RMS = Θ(n1−a−b). Stability requires a + b ≥1. Requiring a nonvanishing, order-one response to the feature change gives a + b = 1. At initialization, the readout and its input feature are independent, so the initial logits have RMS scale Θ(n1/2−a−b) = Θ(n−1/2). 3. Input matrices: fan-in stays fixed. With fixed fan-in d and unchanged external inputs, initialization and direct update scales are n−a−b and ηn−a−c, assuming nonvanishing update alignment. Keeping both order one gives a + b = a + c = 0. Our derivations show that different layers require different choices of (a, b, c) for stable, nonvanishing feature updates. We summarize one such choice below. Here m = n/n0, n0 is the reference width, and k is each matrix’s actual fan-in; the scalings match the baseline at n0. input hidden readout (a, b, c) (0, 0, 0) (0, 1/2, 1) (1, 0, 0) forward multiplier 1 1 1/m initialization variance 1/d 1/k 1/n0 LR η η/m η 16

第 17 页

Assignment 2: Hyperparameter Scaling Problem 4.0: Derive scaling rules for update and feature-change alignment Consider all four combinations of two assumptions, each full or no alignment: • Alignment of the weight-update direction U with the pre-update input x0. • Alignment of the initial readout head V0 ∈Rn×q with its input feature change ∆x = x1 −x0. Here V0 uses the transposed convention, so its response to the feature change is n−aV⊤ 0 ∆x. Write ∆W = −ηn−cU, with fixed η > 0 and ∥U∥RMS = Θ(1). For hidden and readout matrices, assume x0 and ∆x have order-one RMS, and write R(U, x0) = Θ(nα), S(V0, ∆x) = Θ(nωmove). For the interaction term, also assume R(U, ∆x) = O(nα). Use a = 0 for hidden and input matrices and c = 0 for the readout. (a) For each combination, identify α and ωmove and derive the constraints on (a, b, c). Require order-one hidden features at initialization, order-one direct updates, and an order-one response n−aV⊤ 0 ∆x from the initial readout. Summarize (a, b, c) for hidden, readout, and fixed-fan-in input matrices in a table. (b) Using m = n/n0, give the forward multiplier, initialization variance, and LR for each layer type, matching the preceding table at n0. How do the scaling rules change with each alignment assumption? Check that the interaction term stays bounded and identify the combination that recovers the preceding µP rules. 17

第 18 页

Assignment 2: Hyperparameter Scaling 4.1: Toy Example: A five-step Transformer stress test Test how µP scaling transfers hyperparameters across extreme widths and depths. Training protocol. Train for five updates with constant LR from the first update, no warmup, no weight decay, and no gradient clipping. Use unscaled rotary position embeddings (RoPE). Use the same five training batches and 64 fixed validation sequences across parameterizations; evaluate in microbatches and tune LR using validation loss after update five. Width settings. Let m = n/n0 and let k be each matrix’s fan-in. Initialize weight matrices with inde- pendent centered Gaussian entries using the variances below. Setting Kaiming µP Depth 2 Attention-head dimension 64 Reference width n0 512 Weights, activations, optimizer states FP32 Embedding variance 1 Hidden-matrix variance 1/k Initial normalization gains 1 Readout variance 1/n 1/n0 Readout forward multiplier 1 1/m Hidden-matrix LR η η/m Embedding, readout, norm LR η η Depth settings. Fix width and reference width at 64, with one attention head and reference depth two. Keep parameters, gradients, Adam states, and residual additions in FP32; use BF16 autocast for branch computations. Let L count Transformer layers, r = L/2, and cL multiply each attention and feedforward branch before residual addition. Apply these depth multipliers on top of the µP width prescription: Residual cL Block LR multiplier Block Adam ϵ Standard µP 1 1 10−8 Depth-µP r−1/2 r−1/2 10−8r−1/2 CompleteP r−1 1 10−8r−1 Block multipliers apply to hidden matrices and within-block normalization. Embedding, readout, and final-normalization LR and stabilizer remain unchanged. 18

第 19 页

Assignment 2: Hyperparameter Scaling Problem 4.1: Test width and depth scaling over five updates (about 55 runs) Measure feature movement, update alignment, and readout–feature-change alignment. Let ht be a hidden feature vector of width n, evaluated on N fixed probes x1, . . . , xN. For a matrix At with fan-in k, let ∆At = At −At−1 be its actual update and ut−1 its pre-update input. Measure feature movement from initialization and update alignment by Mt = " 1 Nn N ∑ i=1 ∥ht(xi) −h0(xi)∥2 2

1/2

, α(t) upd = logk R(∆At, ut−1). Pool feature RMS over sequences, token positions, and coordinates; compute weight RMS over entries. Inspect normalized features and the residual stream separately. For readout probes, take ht to be the final RMSNorm output before the readout multiplier, and let V0 ∈Rn×q be the initial readout weights in the transposed convention. Measure ω(t) move = logn S(V0, ht −h0). This probes the initial-readout/feature-change term in the derivation. Use the same fixed probes and pool RMS as above; omit the ratio when feature movement is zero. (a) Transfer across widths (about 30 five-step runs). Predict how well Kaiming and standard µP will transfer their best base LR across widths 640, 2,560, and 5,120. Choose an LR sweep for each width and parameterization. Fit each loss–LR curve with a suitable model. Plot fitted optimal base LR and fitted minimum validation loss against width. Which parameterization transfers its tuned base LR better? Does it also achieve lower loss? (b) Probe width scaling (reuse part (a)’s configurations). At each width and parameterization’s best sampled LR, track logit RMS and the RMS and movement Mt of the final RMSNorm output on fixed probes through the five updates. Plot update alignment α(t) upd for each block’s query, key, value, attention-output, and feedforward gate/up/down matrices, and the final readout. Use each matrix’s pre-update input and its actual fan-in k. Also plot ω(t) move, comparing with the no-alignment (1/2) and full-alignment (1) scales. Compare feature movement and alignment across layers and widths. Which assumptions and predictions of the preceding derivation agree with your measurements? (c) Transfer across depths (about 25 five-step runs). Predict which depth prescription should transfer its base LR most reliably. Compare standard µP, Depth-µP, and CompleteP at depths 2, 100, and 1,000, choosing an LR sweep for each depth and prescription. Plot fitted optimal base LR and fitted minimum validation loss against depth. Does better LR transfer imply lower loss after five updates? (d) Probe depth scaling (reuse part (c)’s configurations). At each depth and prescription’s best sampled LR, track residual-stream RMS before normalization and branch-output RMS before multiplication by cL. Measure movement of the normalized residual-stream features at matched relative positions, such as the midpoint and final layer. Also track ω(t) move through all five updates and compare it across depths and prescriptions. Can similar normalized activations hide a growing residual stream? Relate these probes to the LR-transfer results in part (c), distinguishing bounded signals from nonvanishing feature changes. 19

第 20 页

Assignment 2: Hyperparameter Scaling 4.2: Do the scaling rules predict longer training? Do the scaling ideas from Problem 4.1 help over longer training? Use the default model and training settings at 153.6M tokens. The course baseline uses the default architecture and initialization. For µP, change only the four settings in the table below, using reference width n0 = 512. Let n be hidden width, m = n/n0, and η the peak base LR. Let G be a shared matrix of standard Gaussian draws truncated at ±3, used for both embedding and readout initialization and shared across prescriptions. Setting Course baseline (default) µP Embedding initialization G/n G/n0 Readout initialization G/√n G/√n0 Readout forward multiplier 1 1/m Hidden-matrix peak LR η η/m Keep all other training settings fixed, including RMSNorm ϵ = 10−5 and gradient clipping at 1. Normal- ization and clipping can affect widths differently: the baseline embedding scale 1/n may be comparable to √ϵ. Use feature logs and gradients.jsonl to report pre-clipping gradient norms, the embedding’s fraction of squared gradient norm, and the clipping coefficient. Account for these effects when interpret- ing transfer. Reuse the supplied width-512 diagnostics; optional normalization or clipping ablations must fit within your compute budget. Problem 4.2: From five updates to a longer training budget (recommended: 50 new runs) (a) Transfer the reference LR (20 initial runs). Predict whether Problem 1’s best sampled LR at 153.6M tokens transfers to widths 128 and 256 under both prescriptions. Choose an LR sweep for each configuration, including the transferred source LR. Hold data and other settings fixed across new runs; compare with Problem 1’s supplied width-512 curve. At the source LR, plot α(t) upd throughout training for both prescriptions and all three widths. Use fixed inputs and each matrix’s fan-in for every block’s attention and feedforward matrices and the readout. Also plot ω(t) move from Problem 4.1. How do update alignment and readout–feature- change alignment evolve beyond the first five updates? (b) Compare transfer and tuned performance (no new runs). Fit each loss–LR curve. Plot fitted optimal base LR and fitted minimum validation loss against width for both prescriptions. Also plot the measured loss at the transferred LR. Does the prescription with better LR transfer also achieve lower loss after tuning? How does the comparison differ from the five-step test? (c) Predict a held-out width (about 12 new runs). Hold out width 1,024. For each prescription, fit a power law to the optimal base LRs at widths 128, 256, and 512 from part (b), and predict its optimal LR at 1,024. Train at the predicted LR and at the directly transferred source LR. Then run a local LR sweep at 1,024. Compare both predicted LRs with the fitted local optimum and report their measured loss gaps relative to the best sampled target LR. Does fitting the width dependence improve on direct transfer? (d) Transfer across depth (18 initial runs). Fix width 512 and compare depths 4, 8, and 16. Test standard µP, Depth-µP, and CompleteP using the rules from Problem 4.1, now with reference depth eight: for layer count L, use relative depth r = L/8. At depths 4 and 16, compare the 20

第 21 页

Assignment 2: Hyperparameter Scaling transferred source LR with an LR sweep of your choice. Plot fitted optimal base LR and fitted minimum validation loss against depth, alongside the measured losses at the transferred LR. Does a depth correction improve transfer, tuned performance, or both? (e) Identify the training regime (diagnostics from the same runs). At the directly transferred LR and the best sampled LR for each configuration, track logit RMS and normalized hidden- feature movement on fixed inputs during the first five updates and at the end of training, as in Problem 4.1. Include both the held-out width and the depth comparison. Include ω(t) move at the same checkpoints. Which alignment assumptions are supported at different widths, depths, and stages of training? Can a configuration with larger early transients still achieve a better final loss? Does a shift in optimal LR alone show that the model is outside the stable µP regime? Use your diagnostics to explain what the experiments establish. Use the LR–WD coupling from Problem 2 to discuss how holding WD fixed might affect the comparison. (f) Muon and µP (optional). Muon is a momentum-based optimizer that approximately orthog- onalizes the update direction for each hidden weight matrix before applying it. Explore how Muon interacts with µP, following the same approach as in parts (a)–(c). This exploration is not part of the required workload; attempt it only if you have GPU budget remaining after the required experiments. 21

第 22 页

Assignment 2: Hyperparameter Scaling Example question Difficulty: Easy Widening a Kaiming model by copying. Use the five-step width-stress-test setting from Problem 4.1 on matching training batches. Assume Adam’s numerical stabilizer is negligible for these updates. Start with the width-640 Kaiming baseline from the assignment, including its embedding, normal- ization, and readout initialization. Its LR is .001 for every parameter group. Instead of independently initializing a wider model, expand this model by a factor k, where k = 4 or 8. Copy the embedding coordinates and width-dependent normalization gains k times. Copy com- plete attention heads, keeping each head’s dimension at 64; leave the shared 64-dimensional QK- normalization gains unchanged. For each attention or feedforward matrix W (including the rectan- gular gate, up, and down projections), replace it by the k-by-k block matrix whose every block is W/k. If V is the readout matrix, concatenate k copies along its input dimension and divide by k: W′ = 1 k   W · · · W ... ... ... W · · · W  , V′ = 1 k[V · · · V]. There is no additional readout multiplier. The copies are separate trainable parameters, not tied weights. Which learning rates should most closely preserve the width-640 model’s logits on the same inputs after each of the five updates? Select one option for each width. The entries below are actual parameter-group LRs, with no additional width multipliers. “Other” means embedding and all normalization parameters. (a) Width 2560. Option Hidden matrices Readout Other □1 .001 .001 .001 □2 .00025 .00025 .00025 □3 .00025 .00025 .001 □4 .00025 .001 .001 (b) Width 5120. Option Hidden matrices Readout Other □1 .00025 .00025 .001 □2 .000125 .000125 .001 □3 .000125 .000125 .000125 □4 .000125 .001 .001 22

第 23 页

Assignment 2: Hyperparameter Scaling Optional exploration To explore the topics in this assignment more broadly, construct your own prediction problem. This optional exercise can lead to bonus points on the quiz, as described in the Instructions. Problem 5: Your own prediction problem (optional) You may have found an unexpected interaction between schedules, LR and WD, batch size, or model size. Use one such observation to design a question that a classmate could reason about from the assignment, but whose answer goes beyond the supplied examples. Include the following in your exploration: (a) A question writeup in the style of the examples. Provide answer options and state whether to select one or all that apply. Record your prediction and reasoning before testing the target intervention. Make the outcome objectively checkable from an experiment with specified model, data, training budget, and evaluation metric. (b) A compact code diff (under 100 lines), executable code, and the measured evidence that de- termines the answer. Include the answer key and explain any difference from your prediction. Keep the diff focused on the intervention and reuse the shared training interface. If variation across runs could change the correct option, check it and revise the question or option ranges accordingly. Aim for a fair, novel question with plausible alternatives. Start from an experiment you already understand and use the smallest informative comparison. Course staff may use interesting questions in a future quiz. This exploration is additional to the starting grids in the four parts. 23

第 24 页

Assignment 2: Hyperparameter Scaling Example question answer key Try the example questions before reading the answers below. Postponing weight decay (p. 8). Mean loss: 3.21 for both two-update and four-update delays. Match- ing cumulative weight retention does not preserve the training trajectory, but both delays give similar mean loss to decay at every update. Widening a Kaiming model by copying (p. 22). (a) Option 3; (b) Option 2. For expansion factor k, use LR .001/k for hidden matrices and readout, and .001 for embeddings and normalization parameters. With negligible Adam epsilon, scaling a gradient history by a positive constant leaves its normalized update unchanged. Hidden and readout weights must stay 1/k of the original weights, so their LRs scale by 1/k; embeddings and normalization gains remain unscaled copies. Finite precision and nonzero epsilon allow small numerical differences. Doubling batch size halfway through training (p. 12). Policy A: change neither LR nor WD. Its mean loss is 2.93867, compared with 2.94538 for the next-best policy B. Compensating for fewer weight-decay applications per token does not necessarily improve final loss. 24

材料

• 题面 PDF
• 代码模板（GitHub a2_release）

自查

• 交付物 = 提供的实验的分析 + 你自己新跑的实验 + 每题的短解读（不是只贴 loss 数字）。
• 官方限 20 页（不含 bonus）—— 页面限制本身是“取舍”的训练：哪些图必须留、哪些放附录。

这份作业和书的关系

配第 3 讲《优化器（一）：不变量与超参数》、第 4 讲《参数化与预条件》。 书里讲的“学习率随宽度怎么变”“batch size 与 LR 的联合缩放”，这里要你自己扫出来。