<!DOCTYPE html><html lang="en">
<head>
<meta http-equiv="content-type" content="text/html; charset=UTF-8">
<title>[2407.05872] Untitled Document</title>
<meta name="viewport" content="width=device-width, initial-scale=1">
<meta name="twitter:card" content="summary">
<meta name="twitter:title" content="Untitled Document">
<meta name="twitter:image:src" content="https://ar5iv.labs.arxiv.org/assets/ar5iv_card.png">
<meta name="twitter:image:alt" content="ar5iv logo">
<meta property="og:title" content="Untitled Document">
<meta property="og:site_name" content="ar5iv">
<meta property="og:image" content="https://ar5iv.labs.arxiv.org/assets/ar5iv_card.png">
<meta property="og:type" content="article">
<meta property="og:url" content="https://ar5iv.labs.arxiv.org/html/2407.05872">

<!--Generated on Mon Aug  5 17:19:15 2024 by LaTeXML (version 0.8.8) http://dlmf.nist.gov/LaTeXML/.-->
<meta name="viewport" content="width=device-width, initial-scale=1, shrink-to-fit=no">

<script>
  function detectColorScheme(){
    var theme="light";
    var current_theme = localStorage.getItem("ar5iv_theme");
    if(current_theme){
      if(current_theme == "dark"){
        theme = "dark";
      } }
    else if(!window.matchMedia) { return false; }
    else if(window.matchMedia("(prefers-color-scheme: dark)").matches) {
      theme = "dark"; }
    if (theme=="dark") {
      document.documentElement.setAttribute("data-theme", "dark");
    } else {
      document.documentElement.setAttribute("data-theme", "light"); } }

  detectColorScheme();

  function toggleColorScheme(){
    var current_theme = localStorage.getItem("ar5iv_theme");
    if (current_theme) {
      if (current_theme == "light") {
        localStorage.setItem("ar5iv_theme", "dark"); }
      else {
        localStorage.setItem("ar5iv_theme", "light"); } }
    else {
        localStorage.setItem("ar5iv_theme", "dark"); }
    detectColorScheme(); }
</script>
<link media="all" rel="stylesheet" href="/assets/ar5iv-fonts.0.7.9.min.css"><link media="all" rel="stylesheet" href="/assets/ar5iv.0.7.9.min.css"><link media="all" rel="stylesheet" href="/assets/ar5iv-site.0.2.2.css">
</head>
<body>
<div class="ltx_page_main">
<div class="ltx_page_content">
</div>

<div class="ltx_document"><div class="ltx_para"><div class="ltx_p"><span class="ltx_ERROR">
Conversion to HTML had a Fatal error and exited abruptly. This document may be truncated or damaged.
</span></div></div></div>
</article>
<div class="ar5iv-footer"><a href="/html/2407.05871" class="ar5iv-nav-button ar5iv-nav-button-prev">◄</a>
    <a class="ar5iv-home-button" href="/"><img height="40" alt="ar5iv homepage" src="/assets/ar5iv.png"></a>
    <a href="/feeling_lucky" class="ar5iv-text-button">Feeling<br>lucky?</a>
    <a href="/log/2407.05872" class="ar5iv-text-button ar5iv-severity-fatal">Conversion<br>report</a>
    <a class="ar5iv-text-button" target="_blank" href="https://github.com/dginev/ar5iv/issues/new?template=improve-article--arxiv-id-.md&title=Improve+article+2407.05872">Report<br>an issue</a>
    <a href="https://arxiv.org/abs/2407.05872" class="ar5iv-text-button arxiv-ui-theme">View&nbsp;original<br>on&nbsp;arXiv</a><a href="/html/2407.05873" class="ar5iv-nav-button ar5iv-nav-button-next">►</a>
</div><footer class="ltx_page_footer">
<a class="ar5iv-toggle-color-scheme" href="javascript:toggleColorScheme()" title="Toggle ar5iv color scheme"><span class="color-scheme-icon"></span></a>
<a class="ar5iv-footer-button" href="https://arxiv.org/help/license" target="_blank">Copyright</a>
<a class="ar5iv-footer-button" href="https://arxiv.org/help/policies/privacy_policy" target="_blank">Privacy Policy</a>

<div class="ltx_page_logo">Generated  on Mon Aug  5 17:19:15 2024 by <a target="_blank" href="http://dlmf.nist.gov/LaTeXML/" class="ltx_LaTeXML_logo"><span style="letter-spacing:-0.2em; margin-right:0.1em;">L<span class="ltx_font_smallcaps" style="position:relative; bottom:2.2pt;">a</span>T<span class="ltx_font_smallcaps" style="font-size:120%;position:relative; bottom:-0.2ex;">e</span></span><span style="font-size:90%; position:relative; bottom:-0.2ex;">XML</span><img src="data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAsAAAAOCAYAAAD5YeaVAAAAAXNSR0IArs4c6QAAAAZiS0dEAP8A/wD/oL2nkwAAAAlwSFlzAAALEwAACxMBAJqcGAAAAAd0SU1FB9wKExQZLWTEaOUAAAAddEVYdENvbW1lbnQAQ3JlYXRlZCB3aXRoIFRoZSBHSU1Q72QlbgAAAdpJREFUKM9tkL+L2nAARz9fPZNCKFapUn8kyI0e4iRHSR1Kb8ng0lJw6FYHFwv2LwhOpcWxTjeUunYqOmqd6hEoRDhtDWdA8ApRYsSUCDHNt5ul13vz4w0vWCgUnnEc975arX6ORqN3VqtVZbfbTQC4uEHANM3jSqXymFI6yWazP2KxWAXAL9zCUa1Wy2tXVxheKA9YNoR8Pt+aTqe4FVVVvz05O6MBhqUIBGk8Hn8HAOVy+T+XLJfLS4ZhTiRJgqIoVBRFIoric47jPnmeB1mW/9rr9ZpSSn3Lsmir1fJZlqWlUonKsvwWwD8ymc/nXwVBeLjf7xEKhdBut9Hr9WgmkyGEkJwsy5eHG5vN5g0AKIoCAEgkEkin0wQAfN9/cXPdheu6P33fBwB4ngcAcByHJpPJl+fn54mD3Gg0NrquXxeLRQAAwzAYj8cwTZPwPH9/sVg8PXweDAauqqr2cDjEer1GJBLBZDJBs9mE4zjwfZ85lAGg2+06hmGgXq+j3+/DsixYlgVN03a9Xu8jgCNCyIegIAgx13Vfd7vdu+FweG8YRkjXdWy329+dTgeSJD3ieZ7RNO0VAXAPwDEAO5VKndi2fWrb9jWl9Esul6PZbDY9Go1OZ7PZ9z/lyuD3OozU2wAAAABJRU5ErkJggg==" alt="Mascot Sammy"></a>
</div></footer>
</div>

    <script>
      var canMathML = typeof(MathMLElement) == "function";
      if (!canMathML) {
        var body = document.querySelector("body");
        body.firstElementChild.setAttribute('style', 'opacity: 0;');
        var loading = document.createElement("div");
        loading.setAttribute("id", "mathjax-loading-spinner");
        var message = document.createElement("div");
        message.setAttribute("id", "mathjax-loading-message");
        message.innerText = "Typesetting Equations...";
        body.prepend(loading);
        body.prepend(message);

        var el = document.createElement("script");
        el.src = "https://cdn.jsdelivr.net/npm/mathjax@3/es5/tex-mml-chtml.js";
        document.querySelector("head").appendChild(el);

        window.MathJax = {
          startup: {
            pageReady: () => {
              return MathJax.startup.defaultPageReady().then(() => {
                body.removeChild(loading);
                body.removeChild(message);
                body.firstElementChild.removeAttribute('style');
              }); } } };
      }
    </script>
    <script>
    // Auxiliary function, building the preview feature when
    // an inline citation is clicked
    function clicked_cite(e) {
      e.preventDefault();
      let cite = this.closest('.ltx_cite');
      let next = cite.nextSibling;
      if (next && next.nodeType == Node.ELEMENT_NODE && next.getAttribute('class') == "ar5iv-bibitem-preview") {
        next.remove();
        return; }
      // Before adding a preview modal,
      // cleanup older previews, in case they're still open
      document.querySelectorAll('span.ar5iv-bibitem-preview').forEach(function(node) {
        node.remove();
      })

      // Create the preview
      preview = document.createElement('span');
      preview.setAttribute('class','ar5iv-bibitem-preview');
      let target = document.getElementById(this.getAttribute('href').slice(1));
      target.childNodes.forEach(function (child) {
        preview.append(child.cloneNode(true));
      });
      let close_x = document.createElement('button');
      close_x.setAttribute("aria-label","Close modal for bibliography item preview");
      close_x.textContent = "×";
      close_x.setAttribute('class', 'ar5iv-button-close-preview');
      close_x.setAttribute('onclick','this.parentNode.remove()');
      preview.append(close_x);
      preview.querySelectorAll('.ltx_tag_bibitem').forEach(function(node) {
        node.remove();
      });
      cite.parentNode.insertBefore(preview, cite.nextSibling);
      return;
    }
    // Global Document initialization:
    // - assign the preview feature to all inline citation links
    document.querySelectorAll(".ltx_cite .ltx_ref").forEach(function (link) {
      link.addEventListener("click", clicked_cite);
    });
    </script>
    </body>
</html>

2407.05872
 Robust and effective scaling of models from small to large width typically requires the precise adjustment of many algorithmic and architectural details, such as parameterization and optimizer choices. In this work, we propose a new perspective on parameterization by investigating a key assumption in prior work about the alignment between parameters and data and derive new theoretical results under weaker assumptions and a broader set of optimizers. Our extensive empirical investigation includes *tens of thousands* of models trained with *all combinations of* three optimizers, four parameterizations, several alignment assumptions, more than a dozen learning rates, and fourteen model sizes up to 26.8B parameters.  We find that the best learning rate scaling prescription would often have been excluded by the assumptions in prior work. Our results show that all parameterizations, not just maximal update parameterization (muP), can achieve hyperparameter transfer; moreover, our novel per-layer learning rate prescription for standard parameterization outperforms muP. Finally, we demonstrate that an overlooked aspect of parameterization, the epsilon parameter in Adam, must be scaled correctly to avoid gradient underflow and propose *Adamatan2*, a new numerically stable, scale-invariant version of Adam that eliminates the epsilon hyperparameter entirely. 
abstract
Abstract
2407.05872
 A neural network parameterization is a prescription for scaling a set of important quantities with respect to a set of scaling dimensions. Most often, the parameterized quantities include the initialization scale, parameter multipliers and learning rate, and scaling dimensions may include model 1Google DeepMind 2MIT 3Work done at Google DeepMind.  Correspondence to: Katie Everett <everettk@google.com>.  width, model depth, context length, batch size and training horizon. Parameterizations with well-understood scaling behavior can prescribe the exponents for these quantities to ensure that the training dynamics behave in a stable and predictable manner as the model increases in scale.  When exponents are not carefully selected, models can have scaling mismatches that are not obvious from experiments at small scale. One such mismatch can occur in the learning rate exponents between different layers in a neural network: while most models currently train with a global learning rate, where all parameters in the model are updated with the same learning rate, this constrains all layers to use the same exponent when differing exponents may be required.  For example, we will see this phenomenon in standard parameterization models trained with stochastic gradient descent (SGD). Under common assumptions, the hidden layer learning rate would ideally scale like O(1/ √n) whereas the readout layer learning rate would scale like O(1/n) where n is the model width. With a single global learning rate, we are forced to make one of two bad choices: either we choose a learning rate that scales like O(1/ √n), which makes nontrivial updates to the hidden layer but causes the readout layer to explode with scale, or we use a learning rate that scales like O(1/n), which preserves the readout layer stability but causes "vanishing" updates to the earlier activations in the sense that the updates approach zero as the width goes to infinity. This mismatch may persist silently until the model is scaled past a threshold where the difference in exponents dominates over other factors: a risky situation in the current "train once" era of very large models. This strongly motivates a principled understanding of parameterization and well-defined scaling limits rather than relying solely on extrapolation of empirical results.  In addition, parameterizations that are implemented with a particular functional form can enable hyperparameter transfer across scales, where relatively cheap hyperparameter search using small models can be used to select hyperparameters for larger, more expensive models (Yang et al., 2022). This functional form specifies each parameterized quantity using a constant multiplicative factor that is typically determined empirically along with a scaling exponent that is motivated theoretically. This recipe for parameterized quantities, where  parameterized quantity =empirical * [15] M. C. Gonzalez-Garcia, M. C. Gonzalez-Garcia, M.   scaling dimension theoretical exponent 2 allows for the reuse of the constant factors: when the scaling behavior is fully encapsulated by the theoretical exponent, the optimal empirical constant is the same across scales.  We can therefore determine the optimal constant factors via hyperparameter search on small models and reuse them on large models where hyperparameter search would be expensive or impossible.  Among the scaling dimensions, the existing literature for width scaling has the most extensive theoretical results but important open questions remain. Yang & Hu (2021) and Yang & Littwin (2023) define a space of width-scaling parameterizations for SGD and Adam respectively. Based on the goals of ensuring the activations remain at constant scale and the logits do not exceed constant scale with respect to width, they derive constraints for stability and non-triviality of a parameterization that predict how the learning rate should scale with width. The derivations of these constraints make a key assumption: the updates to the parameters are assumed to be sufficiently correlated with the data distribution to impact the scaling exponent of the activations.  We refer to this correlation as "alignment" because, when correlated, the parameter and data vectors point in similar directions. However, the literature is lacking in extensive empirical measurements of when, where, and how much alignment accumulates between the parameters and activations during training and its impact on the scaling exponents of the learning rate.  In this paper, we take a broad perspective on the theory and practice of width scaling across parameterizations and optimizers. Our theoretical contributions consider a more general space of parameterizations which explicitly quantifies the contribution of several distinct alignment terms; under specific alignment assumptions we recover prior work as a special case. We propose a metric for alignment that we use in our empirical investigation. In addition, we develop new parameterization theory for a family of adaptive optimizers with parameter scaling, including Adafactor.  In our experiments, we measure alignment throughout training across optimizers, parameterizations and model sizes.  Our measurements suggest that existing theory may be overly conservative, thereby excluding interesting parameterizations. We show that for all parameterizations and optimizers, there are theoretically motivated per-layer learning rate exponents that improve performance over global learning rate baselines, and that in numerous settings the best performing exponents would have been excluded by the alignment assumptions of prior work.  In particular, a novel per-layer learning rate prescription for standard parameterization is shown to outperform muP: we scale the learning rate for the embedding layer as O(1) and the learning rate for hidden and readout layers as O(1/n) for standard parameterization, and this outperforms all other combinations of parameterizations and learning rate prescriptions for Adam. In addition, while prior work emphasizes the hyperparameter transfer properties of muP specifically (Yang et al., 2022), we show that all parameterizations can perform hyperparameter transfer. We introduce constant multiplicative factors to the per-layer learning rate prescriptions and show that tuning these factors is both essential and practical: the constants can be tuned at relatively small scale and successfully reused across model sizes up to 26.8 billion parameters, inducing substantial performance gains.  Finally, we consider the epsilon hyperparameter in Adam and similar adaptive optimizers. Our theoretical prediction that the mean-field parameterization will be most sensitive to epsilon underflow is validated in our experiments. We see significant performance improvements from three strategies to mitigate epsilon underflow, including our proposal Adam-atan2 that eliminates epsilon entirely. After addressing epsilon underflow, parameterizations that are theoretically equivalent give very similar performance, illustrating that finite precision plays a practical role in the study of parameterization. 
introduction
Introduction
2407.05872
 We define a width-scaling parameterization as in Yang & Hu (2021) as the prescription of the scaling exponents for three quantities on each layer: (1) the initialization variance for the parameters, (2) a parameter multiplier1 by which the trainable parameter weights are multiplied during the forward pass, and (3) the learning rate. It is typical for different layer types (embedding, hidden, and readout) to use different parameterizations within the same network.  We will consider all combinations of four common parameterizations and three common optimizer families. Our parameterizations, shown in Table 1, include standard parameterization (Neal, 1996; Glorot & Bengio, 2010; He et al., 2015), Neural Tangent Kernel (NTK) parameterization (Jacot et al., 2018), Maximal Update parameterization (muP) (Yang & Hu, 2021), and Mean-Field parameterization (MFP) (Mei et al., 2018; Bordelon & Pehlevan, 2022).  Following convention, the *names* of parameterizations will refer to the initialization scale and parameter multipliers, 1The parameter multiplier is a constant that multiplies the output of the matrix multiplication in the layer during the forward pass. The trainable parameters are updated during training but the parameter multiplier is not. However, the backpropagated gradients for the parameters will include this multiplier as a term.  | Initialization Parameter                                                                                                        | Gradient   | SGD LR,    | Adam LR,   | Adafactor LR,   | SGD LR,   | Adam LR,   | Adafactor LR,   |        |    | |---------------------------------------------------------------------------------------------------------------------------------|------------|------------|------------|-----------------|-----------|------------|-----------------|--------|----| | Variance                                                                                                                        | Multiplier | Full Align | Full Align | Full Align      | No Align  | No Align   | No Align        |        |    | | √ n                                                                                                                             | √ n        | 1          | 1          | √ n             | 1         | 1          |                 |        |    | | Embedding                                                                                                                       | 1          | 1          | 1/         |                 |           |            |                 |        |    | | Standard                                                                                                                        | √ n        | 1/ √ n     | 1/n        | 1/ √ n          | 1         | 1/ √ n     | 1               |        |    | | Hidden                                                                                                                          | 1/n        | 1          | 1/         | √ n             | 1/ √ n    | 1/ √ n     | 1               |        |    | | Readout                                                                                                                         | 1/n        | 1          | 1          | 1/n             | 1/n       | 1/         |                 |        |    | | √ n                                                                                                                             | √ n        | 1          | 1          | √ n             | 1         | 1          |                 |        |    | | Embedding                                                                                                                       | 1          | 1          | 1/         |                 |           |            |                 |        |    | | NTK                                                                                                                             | √ n        | 1/n        | √ n        | 1/ √ n          | 1/ √ n    | n          | 1               | 1      |    | | Hidden                                                                                                                          | 1          | 1/ √ n     | 1/ √ n     | 1               | 1/ √ n    | 1/ √ n     | √ n             | 1      | 1  | | Readout                                                                                                                         | 1          | 1/         |            |                 |           |            |                 |        |    | | Embedding                                                                                                                       | 1/n        | √ n        | 1/ √ n     | 1               | 1/ √ n    | 1          | 1               | 1/ √ n | 1  | | muP                                                                                                                             | √ n        | √ n        | 1/ √ n     | 1               |           |            |                 |        |    | | Hidden                                                                                                                          | 1/n        | 1          | 1/n        | 1               | 1/n       | 1/         |                 |        |    | | √ n                                                                                                                             | 1/ √ n     | 1          | 1/ √ n     | 1               | 1         | 1          | 1               |        |    | | Readout                                                                                                                         | 1/n        | 1/         |            |                 |           |            |                 |        |    | | Embedding                                                                                                                       | 1          | 1          | 1/n        | n               | 1         | 1          | n               | 1      | 1  | | √ n                                                                                                                             | 1/n1.5     | n          | 1/ √ n     | 1/ √ n          | n         |            |                 |        |    | | Mean Field                                                                                                                      | Hidden     | 1          | 1/         | 1.5             | 1         | 1          |                 |        |    | | Readout                                                                                                                         | 1          | 1/n        | 1/n        | n               | 1         | 1          | n               | √ n    | 1  | | Table 1. Summary of parameterizations, their gradients and the learning rates derived in Section 3. Left: Parameterizations and |            |            |            |                 |           |            |                 |        |    |  although formally speaking, the learning rate prescription is an essential element of a parameterization.  pass quantities like gradients or updates are indirect consequences.  We select optimizers that represent three distinct widthscaling regimes: Stochastic Gradient Descent (SGD), Adam (Kingma & Ba, 2014), and Adafactor (Shazeer & Stern, 2018) due to their varying relationships between the parameter, gradient, and update scales. Our theoretical perspective focuses on the width-scaling relationships between these elements and will omit more specific optimizer features like momentum and gradient or update clipping. SGD represents optimizers where the scale of the update matches the scale of the learning rate times the scale of the gradients. Adam represents adaptive optimizers where, due to the normalization by the gradient scale, the scale of the update matches the scale of the learning rate regardless of the gradient. Finally, Adafactor represents adaptive optimizers with parameter scaling that normalize the gradient similarly to Adam but then multiply by the parameter scale.  This results in an update scale that matches the learning rate scale × the parameter scale; under constant learning rates, the Adafactor updates match the RMS (or Frobenius) norm of the parameters. Note that parameter scaling is the key distinction between the last two regimes: Adam plus parameter scaling falls into the Adafactor family; Adafactor with parameter scaling removed falls into the Adam family. 
body
Background 2.1. Parameterizations And Optimizers
2407.05872
 Yang & Hu (2021) and Yang & Littwin (2023) derive a system of linear constraints on the exponents of a parameterization from the two following concepts: *stability*, where the activations are exactly constant scale and the logits are no more than constant scale, and *nontriviality*, where the change in logits after the initialization is at least constant scale. Note these constraints are defined solely in terms of the activations and logits, so that only the forward pass is directly constrained and any constraints on backward In addition, they define *feature learning*, where the latent representations change and adapt to the data during training, as a constant scale change after initialization in the activations directly before the readout layer. These activations are the *features* or latent representations learned by the model, which, for example, could be reused with another classifier if we were to remove the readout layer. When the change in these activations is exactly constant scale, the change is meaningful in the infinite-width limit rather than becoming infinitesimally small as the width becomes large. Finally, they fully characterize the infinite-width limits of the space of width-scaling parameterizations as a dichotomy between a feature learning regime and a kernel regime. 
body
Stability, Nontriviality And Feature Learning
2407.05872
 As we will see in Section 3, the conditions for stability differ between the first and subsequent forward passes because of the learning process itself. While the initial random parameters are independent from the data distribution, correlations can develop over time because the updates carry information about the data. Such correlations cause "alignment" between the parameters and activations, in the sense that the vectors may point in similar directions. As such, the norm of the activations after a given layer depends on three quantities: the scale of the input to the layer, the scale of the parameters in the layer, and the alignment between the parameters and the input or "data". In a matrix multiplication, when we sum over the interior dimension n, this alignment contributes a scaling term that is O( √n) when there is no alignment and O(n) when there is significant alignment.  The intuition for this calculation can be seen by considering the simpler case of the scaling of the inner product of two random vectors, because the entries in the product of a matrix multiplication are each vector-vector inner products themselves. As the length of the two random vectors becomes large, by straightforward application of the Central Limit Theorem the inner product is a sample from a normal distribution, so its norm has two terms coming from the mean and the variance of this distribution. The takeaway is that the mean term contributes an O(n) term to the norm of the inner product and the variance term contributes O( √n).  When the mean term is zero because the vectors are zeromean and independent, then the variance term dominates and the inner product scales like O( √n). However, when the vectors are *correlated* or in the worst case are identical, then the inner product scales like O(n) because the coefficient to the mean term is a constant. Yang & Hu (2021) refers to this idea as "Central Limit Scaling" versus "Law of Large Numbers" scaling owing to the idea that the Law of Large Numbers governs how the mean converges and the Central Limit Theorem governs how the variance converges.  This scaling affects the norm of the outgoing activations from a layer that multiplies its parameter matrix by the input to that layer. On the first forward pass, due to the random initialization the parameters cannot be aligned with the data, so the O( √n) scaling holds. However, during the first backward pass, the updates to the parameters are a function of the first batch of data, so the parameters may become correlated to the data distribution. As a result, during subsequent forward passes, we can no longer assume perfect independence between the parameters and data and instead the activations might scale with up to an O(n) term.  Similar to this alignment between parameter updates and the data distribution, it is also possible to develop alignment between the parameters in two different layers in the network. For example, the backpropagated gradients used to update an earlier layer are a function of the parameters in later layers, which can introduce correlation between the earlier layer updates and the later layer parameters. The consequence of either type of alignment is that the learning rate needs to *counteract* the O( √n) or O(n) term, so the maximal stable learning rates can be *smaller* by a factor of O( √n) when there is significant alignment than when there is none. 
body
Alignment
2407.05872
 In this section we make four theoretical contributions. First, we define a general space of width-scaling parameterizations that explicitly quantifies the contribution of three alignment terms. Rather than making specific assumptions about alignment and then deriving which parameterizations are stable and nontrivial under those assumptions, as in Yang & Hu (2021) and Yang & Littwin (2023), we propose general stability and nontriviality constraints as a function of those alignment variables. Second, we propose theory for Adafactor or other adaptive optimizers using parameter scaling.  Third, for all parameterizations × optimizers, we find the maximum stable learning rate for each layer type as a function of the alignment terms, and compute the learning rate exponents under two specific alignment assumptions, which we refer to as "full alignment" and "no alignment". While the alignment assumptions in Yang & Hu (2021) prevent standard and NTK parameterizations from feature learning regardless of the per-layer learning rate prescription, under our assumptions of either full alignment or no alignment, all parameterizations have per-layer learning rates in the feature learning limit. Fourth, we propose the *alignment ratio* metric that we will use for empirical investigation, which measures the contribution of alignment to the activations during training. 
body
Theoretical Contributions
2407.05872
 Following a similar model and notation as Yang & Hu (2021), we consider a multilayer perceptron with L hidden layers, input and output dimensionality d, hidden layer dimensionality n, and nonlinearity ϕ : R → R. The weight matrices are denoted:  - W1 ∈ R n×dfor the embedding layer - W2*, . . . W*L ∈ R n×n for the hidden layers, and - WL+1 ∈ R d×n for the readout layer. The parameterization for each layer l is specified by three values {al, bl, cl}, where: - the parameter multiplier is n −al,  - the parameter initialization is Wl ∼ N (0, n−2bl ), and - the learning rate ηl ∝ n −cl with width-independent constant of proportionality that we omit here. For an input x ∈ R d, the model has activations z1*, . . . z*L and outputs logits zL+1:  $$\begin{array}{l l}{{z_{1}=\phi(n^{-a_{1}}W_{1}\cdot x)}}\\ {{z_{l}=\phi(n^{-a_{l}}W_{l}\cdot z_{l-1}),\quad l\in[2,L]}}\\ {{z_{L+1}=n^{-a_{L+1}}W_{L+1}\cdot z_{L}}}\end{array}$$ 4 In addition, we define ∆Wt l and ∆z t l to be the change in parameters and activations, respectively, in layer l between initialization and step t. We omit the time superscript throughout this section when it is clear from context or the statement holds for any O(1) value of t.  We are interested in the scaling behavior of various quantities as we increase the *width* or hidden layer dimensionality n, while other dimensions are held constant. In particular, we assume input and output dimensionality d, the depth L, and the number of training steps T are fixed and constants  ![4_image_0.png](4_image_0.png)  ![4_image_1.png](4_image_1.png)  a with respect to width. We omit batch size dimensions, assuming the batch size is one. We assume the nonlinearity ϕ has bounded derivative so its contribution is negligible in the large-width limit, and as such treat it as the identity function in our derivations. Throughout this section, we refer to the "scale" of quantities, meaning their exponents with respect to width in the large-width limit. For formal definitions, additional assumptions, and a full derivation see Appendix B. 
body
Model And Notation
2407.05872
 These parameterizations occupy equivalence classes because in any layer we can "factor out" a constant term from the parameter initialization into the parameter multiplier, which exactly preserves the output of the forward pass while multiplying the gradients by this constant. This change in the gradients can then be "corrected for" by modifying the learning rate in an optimizer-specific manner.  In this one-dimensional symmetry group parameterized by θ, to preserve the forward pass, regardless of the optimizer, apply  $$\begin{array}{l}{{a_{l}\gets a_{l}+\theta}}\\ {{b_{l}\gets b_{l}-\theta.}}\end{array}$$  Then specific to the optimizer, to preserve the effect of the backwards pass, correct the learning rate according to  $$\begin{array}{r l}{\mathrm{SGD:}}&{{}c_{l}\gets c_{l}-2\theta}\\ {\mathrm{Adam:}}&{{}c_{l}\gets c_{l}-\theta}\end{array}$$ $$\begin{array}{r l}{\mathrm{Adafactor:}}&{{}c_{l}\gets c_{l}.}\end{array}$$  In particular, under the right learning rates, our four parameterizations occupy two equivalence classes: standard and NTK are equivalent and muP and mean-field parameterization are equivalent. In Figure 1, we visualize the subspaces spanned by the equivalence classes for each layer type: each parameterization is plotted at (al, bl) for each layer l, and the dashed lines span the equivalence class (al + *θ, b*l − θ) for all values of θ. We see that for both the embedding layer and the hidden layers, all four parameterizations have equivalent initializations. However, the readout layer has two distinct equivalence classes: standard and NTK parameterizations have aL+1 + bL+1 = 1/2, corresponding to constant scale logits at initialization, whereas muP and MFP have aL+1 + bL+1 = 1, resulting in logits that scale like 1/ √n at initialization. This difference, specifically the shift by a factor of √n in the readout layer, is the key difference between the standard + NTK equivalence class and the muP + MFP equivalence class.  In this paper, we will consider all four parameterizations separately, as these equivalences hold only under infinite precision, while neural networks regularly encounter finiteprecision effects. These equivalences were observed for SGD and Adam in Yang & Hu (2021) and Yang & Littwin (2023) respectively, and we propose this equivalence for Adafactor. 
body
Equivalence Classes
2407.05872
 We now propose a general space of parameterizations where we define three alignment variables and derive the space of parameterizations that are stable and nontrivial as a function of these variables. We define a parameterization to be *stable* if the activations have exactly constant scale and the logits have at most constant scale throughout training, and to be nontrivial if the change in logits after initialization has at least constant scale.  | Scaling Exponents Across Parameterizations and Optimizers   |                                                |                                            |                           |                     |     |           |                           |     | |-------------------------------------------------------------|------------------------------------------------|--------------------------------------------|---------------------------|---------------------|-----|-----------|---------------------------|-----| | SGD                                                         | Adam                                           | Adafactor                                  |                           |                     |     |           |                           |     | | initialization                                              | a1 + b1 = 0                                    |                                            |                           |                     |     |           |                           |     | | Stability at                                                | al + bl = 1/2 for l ∈ [2, L] aL+1 + bL+1 ≥ 1/2 |                                            |                           |                     |     |           |                           |     | | r1 := g1 + a1 + c1 ≥ 0                                      | r1 := a1 + c1 ≥ 0                              | r1 := c1 ≥ 0                               |                           |                     |     |           |                           |     | | Stable activations during training                          |  gl + al + cl − αl                            |  al + cl − αl                             |  1/2 + cl − αl           |                     |     |           |                           |     | |                                                           |                                              |                                          |                           |                     |     |           |                           |     | | rl := min                                                   | gl + al + cl + rl−1 − ul                       | ≥ 0                                        | rl := min                 | al + cl + rl−1 − ul | ≥ 0 | rl := min | 1/2 + cl + rl−1 − ul      | ≥ 0 | |  1/2 + rl−1 − ωl                                          |  1/2 + rl−1 − ωl                             |  1/2 + rl−1 − ωl                         |                           |                     |     |           |                           |     | | where gi :=                                                 | cl ≥ 0                                         |                                            |                           |                     |     |           |                           |     | | max(aL+1 + bL+1, 2aL+1 + cL+1) + ai                         |                                                |                                            |                           |                     |     |           |                           |     | | during training                                             |  aL+1 + bL+1 + rL − ωL+1                      |  aL+1 + bL+1 + rL − ωL+1                  |  aL+1 + bL+1 + rL − ωL+1 |                     |     |           |                           |     | | Stable logits                                               |                                              |                                          |                         |                     |     |           |                           |     | | min                                                         | 2aL+1 + cL+1 − αL+1                            | ≥ 0                                        | min                       | aL+1 + cL+1 − αL+1  | ≥ 0 | min       | aL+1 + bL+1 + cL+1 − αL+1 | ≥ 0 | |  2aL+1 + cL+1 + rL − uL+1                                 |  aL+1 + cL+1 + rL − uL+1                     |  aL+1 + bL+1 + cL+1 + rL − uL+1 cL+1 ≥ 0 |                           |                     |     |           |                           |     |  $$\begin{array}{l}{{+\Delta z_{l-1})}}\\ {{\cdot_{1}+W_{l}\Delta z_{l-1}+\Delta W}}\end{array}$$  The activations change after initialization due to both changes in the parameters in the previous layer and changes in the activations immediately prior to that layer. Specifically, starting from the second forward pass, the activations for layer l are computed as  $$(W_{l}z)$$ 
body
Alignment-General Space Of Parameterizations
2407.05872
 The first term Wlzl−1 in the expanded sum contains the initial random parameters and initial activations, which are not aligned. The remaining three terms in the sum may have alignment as they result from updates that might be aligned to the data distribution or other parameters in the model.  We will define αl, ωl, and ulto be the exponents of the alignment contributions from these three terms, where the alignment exponent quantifies how the norm of the product scales compared to the norm of the factors. We define  - αlto be the alignment exponent for ∆Wlzl−1, - ωlto be the alignment exponent for Wl∆zl−1, and - ulto be the alignment exponent for ∆Wl∆zl−1. so that, for example, ∥∆Wlzl−1∥ scales like n αl∥∆Wl∥∥zl−1∥. With this definition, αl = 1/2 corresponds to no alignment in the ∆Wlzl−1 term and αl = 1 corresponds to high alignment.  We next define a feature learning residual quantity rlthat measures how far the parameterization is from the feature learning regime. For each layer l in [1, L], we define rl as the negative exponent of the scale of ∆zl, where ∆zlis the change in activations following layer l during training.  To preserve stability, this change cannot exceed constant scale, so rl, as the negative exponent, cannot be less than zero. Feature learning, where the change in activations immediately prior to the readout layer has constant scale, then corresponds to rL = 0 exactly. Conceptually, feature learning occurs if at least one of the embedding or hidden layers contributes at least one constant scale term to the activations. Lastly, for convenience, we define glto denote the negative exponent of the gradient scale. 
body
Zl = N Al(Wl + ∆Wl)(Zl−1 + ∆Zl−1) = N Al(Wlzl−1 + ∆Wlzl−1 + Wl∆Zl−1 + ∆Wl∆Zl−1).
2407.05872
 To derive the stability constraints, which are shown in Table 2, we will first derive constraints on the initialization scale and parameter multipliers so that the parameterization is stable during the first forward pass. In the next subsection, we will derive constraints for stability during training that ensure the *change* in activations after every layer has at most constant scale and, similarly, we will bound the change in logits to ensure they do not exceed constant scale. For a more detailed derivation of the stability constraints, see Appendix B. For nontriviality, we include the derivation and constraints in Appendix B.14 and B.15.  During the first forward pass, the conditions for stability depend only on the parameter initialization and the parameter multipliers, and not on the optimizer, learning rates or alignment variables because no updates have occurred yet.  For all optimizers, for input data x with constant scale, the constraints for stability at initialization are:  a1 + b1 = 0 $a_{l}+b_{l}=1/2$, $l\in[2,\ldots,L]$, $a_{L+1}+b_{L+1}\geq1/2$ 7 3.3.2. STABILITY DURING TRAINING During the second and subsequent forward passes, the constraints for stability are specific to the optimizer. In Appendix B, we derive the constraints on the parameter multipliers, parameter initialization, and learning rates under each optimizer that ensure stability during training, that is, constant activations and at most constant logits. In our notation, recall that rlis the negative exponent of the scale of ∆zl and glto is the negative exponent of the gradients with respect to parameters Wl.  For intuition, this derivation first considers the second forward pass, and assuming the constraints for stability at initialization hold, iteratively adds constraints starting with the embedding layer and working up to the readout layer, where each constraint ensures stability of that layer assuming stability of all earlier layers. We compute the scale of the change in activations in each layer, and bound its exponent to be at most zero. Note that in our constraints, this is written to bound the *negative* exponent to be *at least* zero. Similarly, we constrain the logits to stay at most constant scale in the second forward pass. The constraints derived from the second forward pass are included in Appendix B.12.  Next, we consider the third forward pass, because the readout parameters may change in scale between the first update and second update. This slightly modifies the constraints to use the maximum possible scale of the readout parameters, which may come from the readout initialization or the readout update depending on the parameterization. This produces the final set of stability constraints, as after the third forward pass the scale of each quantity remains the same over a constant number of training steps.  We present these constraints for stability during training in Table 2. The set of constraints has some similarity across optimizers. The change in activations following a hidden layer l is computed as the sum of three terms, ∆Wlzl−1, Wl∆zl−1 and ∆Wl∆zl−1. As the term with the maximum scale dominates the exponent, the value of rlis the minimum over three expressions each coming from one of these terms.  The constraint rl ≥ 0 then preserves stability by preventing any of these terms from exceeding constant scale, where when rL = 0 exactly, we are in the feature learning regime.  To highlight the differences across optimizers, we note some of the SGD constraints include the term gl, as the SGD update depends on the scale of the gradient. Compared to SGD, the Adam constraints result from removing the contribution of gl due to the normalization of the gradient in Adam, even in the readout layer where gL+1 = aL+1.  Then, compared to Adam, the constraints for Adafactor have additional appearances of bl as a result of the parameter scaling; in some places this simplifies due to substituting a1 + b1 = 0 or al + bl = 1/2. Lastly, there is an additional constraint cl ≥ 0 required for Adafactor to prevent the parameters from growing exponentially with the number of training steps as a result of parameter scaling. 
body
1. Stability At Initialization
2407.05872
 We exactly recover the stability and nontriviality constraints in Yang & Hu (2021) for SGD and Yang & Littwin (2023) for Adam2if and only if - αl = 1 ∀l ∈ [2, L + 1],  - $\omega_l=1/2\quad\forall l\in[2,L],\quad$ and...  - ωL+1 = 1.  The choice of αl = 1 on all layers depends on the assumption that updates to the parameters ∆Wl are aligned with the activations zl−1 due to alignment between the parameter updates and the data distribution. We will investigate settings with and without this assumption.  We note that the assumption ωL+1 = 1, which we will relax, is at the very core of the theoretical motivation for muP.  Recall that ωL+1 is the alignment exponent for WL+1∆zL, where WL+1 is the readout layer initialization and ∆zL is the change in activations immediately prior to the readout layer resulting from changes in the parameters in earlier layers. In theory, this alignment could develop when the gradients propagated to earlier layers contain information from the initialization parameters in later layers. However, we also note the assumption of ωl = 1/2 on all layers except the readout layer assumes that the analogous alignment does not develop in any earlier layers.  A key consequence of the ωL+1 = 1 assumption is that neither standard nor NTK parameterizations are able to achieve feature learning *regardless of what learning rate prescription is used*. Recall that feature learning corresponds to rL = 0, where the change in activations ∆zL is exactly constant scale, and that standard and NTK parameterizations both have aL+1 + bL+1 = 1/2. Therefore, due to the constraint on the logits that aL+1 +bL+1 +rL −ωL+1 ≥ 0, it is not possible when ωL+1 = 1 for standard and NTK parameterizations to have rL = 0 regardless of how it might be induced. However, the shift in the readout layer in muP and MFP so that aL+1 + bL+1 = 1 allows these parameteri-2Yang & Littwin (2023) carefully considers the role of the epsilon hyperparameter in Adam. This correspondence holds if we assume what they call *faithfulness*, which is equivalent to setting epsilon following the per-layer epsilon prescription we implement in §4.3. It also holds if we disregard epsilon (i.e. consider epsilon to be zero) and view Adam as perfectly scale-invariant.  zations to attain feature learning when ωL+1 = 1. 
body
3. Prior Work As A Special Case
2407.05872
 3.3.4. MAXIMUM STABLE LEARNING RATE EXPONENTS We will next define two specific sets of alignment assumptions in terms of αl, ωl, and ul and then derive the maximum stable learning rates shown in Table 1 for each layer and parameterization under these assumptions. To select these assumptions, we first consider the types of alignment measured by each variable and how the alignment assumptions dictate the maximum stable learning rates.  The ∆Wlzl−1 term, whose alignment is measured by αl, may have alignment between the parameter updates and data distribution. If we unroll the layers in the term ∆Wlzl−1 = ∆Wl·Wl−1 ·Wl−2 *. . . W*1 ·x, since the parameters W1*, . . . , W*l−1 are from the random initialization, any alignment comes from the matrix multiplication between the parameter updates ∆Wl and the data x.  In contrast, for the Wl∆zl−1 term, whose alignment is measured by ωl, the alignment occurs not between parameters and data, but between the updates to parameters in earlier layers and the initialization parameters in the later layer l. Recall that this alignment could develop as the backpropagated gradients, which are used to update the earlier layers, depend on the initialization parameters in the later layers. The final term ∆Wl∆zl−1, whose alignment is measured by ul, may contain both kinds of alignment, between the parameters and data or between parameters in different layers.  In our alignment settings, we will assume that the alignment between parameters in different layers stays small, and focus on exploring the range of possible degrees of parameter-todata alignment. Specifically, we assume that ωl = 1/2 on all layers including the readout layer ωL+1 = 1/2. As a consequence, we will be able to consider feature learning versions of all four parameterizations.  For the assumptions on αl and ulthat measure alignment between parameters and data, we note that in the stability constraints, αl and ul always appear in tandem in pairs of constraints, where in fact the constraints could be rewritten entirely in terms of the single quantity max(αl, ul − rl−1).  When using the maximum stable learning rates on earlier layers, rl−1 will be zero, so the learning rate exponents are constrained by the maximum of αl and ul. We will therefore consider the two extremes, and make two choices of alignment assumptions: "full alignment" where αl = ul = 1 and ωl = 1/2 and "no alignment" where αl = ul = ωl = 1/2. Intermediate choices where αl = ultakes a value between 1/2 and 1 would also be interesting for future work, but scaling studies at practical sizes may not have sufficient resolution to distinguish smaller variations In Table 1, we compute the maximum stable per-layer learning rate exponents under these two assumptions of full alignment and no alignment. Our full alignment per-layer learning rate prescriptions for standard and NTK differ from prior work and are in feature learning limits under our assumptions. For muP and MFP, our full alignment assumptions result in learning rate exponents that coincide with prior work as relaxing the ωL+1 constraint does not impact these parameterizations. Our no alignment prescriptions are also in feature learning limits for all parameterizations, and again differ from prior work that assumed αl = 1. 
body
In The Exponents.
2407.05872
 The alignment variables αl, ωl, and ul quantify the alignment contributions to the activations zl from the individual terms in the expanded sum (Wl + ∆Wl)(zl−1 + ∆zl−1) = Wlzl−1 + ∆Wlzl−1 + Wl∆zl−1 + ∆Wl∆zl−1. However, we note that in practice the alignment in the terms in this sum may interfere constructively or destructively, and the single alignment quantity that actually governs the scale of the activations zlis the alignment between (Wl +∆Wl) and (zl−1 + ∆zl−1).  We therefore propose a metric that measures this alignment, between (Wl + ∆Wl) and (zl−1 + ∆zl−1), in order to understand empirically how alignment that accumulates during training is contributing to the activation scales throughout the model. This metric is defined on each dense layer and quantifies the contribution from alignment between the current parameters and the current pre-layer activations on the scaling exponent of the post-layer activations.  Consider a neural network layer l at training step t with parameters Wt l ∈ R fan-out×fan-in, and pre-layer activations z t l−1 ∈ R fan-in. The alignment between the pre-layer activations and the rows of the parameter matrix dictates the scaling term contributed when summing over the fan-in dimension in the matrix multiplication, giving an alignment term with the fan-in as the base of the exponent.    **Definition 3.1**.: We define the _log alignment ratio_ as  $$A_{l}^{t}=\log_{\text{fan-in}}\frac{\left\|W_{l}^{t}z_{l-1}^{t}\right\|_{RMS}}{\left\|W_{l}^{t}\right\|_{RMS}\left\|z_{l-1}^{t}\right\|_{RMS}}\in\mathbb{R},$$  where the norm is the RMS norm${}^{3}$.   We note that while our theoretical derivations consider exponents in the infinite-width limit, this metric is intentionally defined at finite size so that we can measure it in practice.  We will see in Section 4.1 that this metric shows convergence within our model sizes suggesting that in practical 3An equivalent definition for the log alignment ratio in terms  ![8_image_0.png](8_image_0.png)  settings this finite size approximation is indicative of largewidth behavior. 
body
Alignment Ratio
2407.05872
 We investigate the role of alignment, per-layer learning rate exponents and constant factors, and the epsilon hyperparameter by running *tens of thousands* of experiments in a Transformer language model across all combinations of the three optimizers, four parameterizations, learning rate sweeps with a granularity of 2 0.25 or 2 0.5, and fourteen model widths ranging up to 26.8 billion parameters.  We use the NanoDO decoder-only Transformer architecture (Liu et al., 2024) employing learned positional embeddings, pre-layer norm (Ba et al., 2016; Xiong et al., 2020), and GeLU nonlinearity (Hendrycks & Gimpel, 2016). All models are trained on the C4 dataset (Raffel et al., 2020). To scale the width, we fix the attention head dimension h = 128 and co-scale the model dimension D and the number of attention heads H so that D = H × h. All experiments use a fixed batch size of 256, context length of 512 and depth of 8 Transformer blocks. Our fourteen model widths range from D = 128 to D = 16, 384 corresponding to a range of 9.9 million to 26.8 billion parameters.  All experiments in the main text of the paper train for 50,000 training steps and do not use weight decay. The learning rate schedule for all experiments uses linear warmup of 1,000 steps followed by a cosine decay schedule, with initial and final learning rates of 0.0. We include additional experiments for compute-optimal training horizons in in Appendix I that show that moving from the fixed step setting to the compute optimal setting likely requires sharper decay in the learning rate exponents. In addition, we include experiments in Appendix G that suggest our conclusions should transfer to settings using small amounts of weight decay. See Appendix C for additional experiment and implementation details. 
body
Experiments
2407.05872
 We measure the log alignment ratio throughout training for three model sizes for each parameterization × optimizer, using a global learning rate that is close to optimal. For Adam, alignment values for the readout and MLP layers are shown for model dimension D = 4096 in Figure 2 and full results including the other optimizers and the dense layers within the attention block (query, key, value, and output projection dense layers) are included in Appendix D. All experiments use batch size B = 256. As expected, the measured alignment values start at 0.5 due to the independence between the random initialization and the data, and change during training as parameters become aligned with the data or parameters in earlier layers.  The alignment values vary significantly across the training horizon and the trajectories depend heavily on the parameterization and layer type. We see similar values across three model sizes, suggesting that these model sizes are sufficiently large for our measurements to be indicative of large-width behavior. For SGD, the results show high instability and are difficult to interpret; one consistent pattern is that NTK and MFP have almost no alignment and STP and muP have low amounts. Adam and Adafactor show matching trends: the readout layer has the highest peak among the layers, with high peak readout alignment above 0.9 in muP and mean-field parameterization and only moderate peak readout alignment between 0.7 and 0.8 in standard and NTK parameterizations. In the MLP layers, alignment in all parameterizations is moderate and does not exceed 0.8.  The high readout alignment early in training that is specific to muP and MFP parameterizations may result from the relationship between the readout updates and readout initialization. In standard and NTK parameterizations, the scale of the readout update matches the scale of the readout initialization, whereas in muP and MFP, the scale of the update is larger than the scale of the initialization. At each step, the parameter updates may be highly aligned with the data distribution, resulting in high readout alignment early in training because these highly aligned updates dominate the  ![9_image_0.png](9_image_0.png)   readout parameters in muP and MFP. However, as more of these updates accumulate, the alignment actually decreases to a moderate level, as the sum of the accumulated updates may be less aligned than the individual updates. In contrast, in standard and NTK parameterizations, the individual updates may still be highly aligned with the data distribution without resulting in high readout alignment as they don't dominate over the readout initialization.  Overall, these results indicate that the alignment assumptions in Yang & Hu (2021); Yang & Littwin (2023) may be overly conservative; if in practice alignment contributes less than one to the activation exponents, then the learning rate exponents can be larger. However, even given these measurements for alignment, it is not obvious how exactly the learning rate exponents that control the base or peak learning rate should be adjusted. First, we see that alignment is a dynamical quantity that varies widely throughout training: even when the readout alignment is close to maximum early in training for muP and MFP, it is much lower for the vast majority of training steps. In addition, alignment varies across layers within the same layer type, which typically use the same learning rate or at least the same learning rate exponent. Further, the interaction between the alignment measurement and learning rate schedule is complex.  The learning rate schedule likely influences the alignment measurements, and alignment likely influences the optimal learning rate schedule: one possible role of learning rate schedules is that the decay counteracts alignment that develops later in training. Even if we used alignment measurements from experiments under one set of learning rates to inform adjustments to the learning rate exponents, this would induce an iterative loop where the adjusted learning rates would then affect the alignment. We will therefore take an empirical approach to determine how alignment should influence the learning rate exponents, and consider several choices of alignment assumptions for the per-layer learning rate experiments in the following section. 
body
Alignment Experiments
2407.05872
 All parameterizations have per-layer learning rate exponents specific to the optimizer and the choice of alignment assumption, as shown in Table 1. In this section, we empirically validate these theoretical learning rate exponent prescriptions and investigate the impact of tuning the per-layer constant factors on all combinations of parameterizations × optimizers. We compare against global learning rates as a baseline: while global learning rates are in most cases not theoretically principled, they are the overwhelmingly dominant paradigm in practice. In most settings, the theoretically ideal learning rate exponents differ across layers, implying a mismatch in at least some layers when using global learning rates. However, there are certain cases, such as muP + SGD + full alignment, or Adafactor + any parameterization + no alignment, where the theoretically motivated per-layer exponents happen to be the same in all layers so that global learning rate actually coincides with the theoretical prescription. We include experiments using the theoretical prescriptions from both the full alignment and no alignment settings since our empirical alignment measurements show intermediate and highly dynamic values. We first present the series of experiments for Adam, followed by the results for SGD+momentum and Adam+parameter scaling.  Additional results are included in the appendix including additional ablations (Appendix E), eval losses for all settings for the six largest model sizes (Table E1) and learning rate sweeps for all settings for SGD (§J), Adam (§K) and Adam + parameter scaling (§L). 
body
Per-Layer Learning Rates
2407.05872
 Our first experiment compares per-layer learning rates when assuming full alignment against the baseline of global learning rates. For all experiments, we select a base model dim of b = 1024 and define the learning rate in layer l as ηl = βn · n b −cl. We perform a one-dimensional sweep of βn at each model dim n to determine the best value and  ![10_image_0.png](10_image_0.png)   1  report the eval loss from this best βn. For global learning rates, cl = 0 for all l, and for per-layer learning rates, cl follows Table 1. Since there is no contribution from the learning rate exponents at the base model size, the global and per-layer learning rate settings coincide exactly at the base model dim and differ only when scaling away from the base model size. For Adam, the per-layer results in Figure 3(b), compared with the global learning rate results in Figure 3(a), show that standard, muP and mean-field parameterizations all improve significantly with per-layer learning rates. In contrast, NTK actually performs worse with the prescribed per-layer exponents: we note that these exponents assume full alignment and return to this result later in this section.  We next introduce constant multiplicative factors to the perlayer learning rates and propose a hyperparameter transfer strategy where we tune these constants at small scale and reuse them across model sizes. Again using the base model dim of b = 1024, we now define the learning rate in layer l as ηl = βn · γl· n b −cl where γlis this constant factor that should be determined empirically. We use the same learning rate within each layer type, so we need to determine γ1 for the embedding layer, γh for the hidden layers, and γL+1 for the readout layer. The previous experiments correspond to setting these constants equal to one by default. We continue to sweep one dimension at each model size to determine βn, so we note the choice of (γ1, γh, γL+1) really defines two ratios as a common factor can be absorbed by βn. We tune (γ1, γh, γL+1) at the base model dim using a threedimensional hyperparameter search described in §C.4 and reuse the best set of values across all model sizes to give the eval losses in Figure 3(c). We include the values found for these estimated optimal constant factors for all optimizers and parameterizations in §C.4.  Using these optimal constants and the per-layer learning rate exponents that assume full alignment, we see that all parameterizations can perform hyperparameter transfer across width. First, the eval loss improves across all scales when using the optimal constants instead of the default constants, indicating successful transfer of these constant multipliers. For the base model dim b = 1024, tuning these constants can only improve the performance, but comparing Figure 3(c) to 3(b) shows that for Adam these constants improve performance substantially for all model sizes across all parameterizations. If instead, these constants improved performance only at or near the model sizes where they were tuned, it would be less practical to include them for large model training because the three-dimensional sweep would be prohibitively expensive at large scale. The only exception is that these constants may not transfer well to the mean-field parameterization models above 2B parameters: we will show in Section 4.3 that this is likely due to epsilon underflow and is addressed with our epsilon mitigations. These results indicate that our recipe for these per-layer constant multiplicative factors is both essential and practical: the performance gains are substantial and the hyperparameter transfer makes them feasible.  Second, the optimal base learning rate is consistent across scale in all parameterizations. In Figure 4, we first find the  SGD  STP 3.657 3.510 3.385 **3.318** NTK 3.732 3.627 **3.313** 3.324  muP 4.224 4.184 **3.809** 4.222  MFP 4.092 4.319 4.092 **3.898**  | Global LR +       | Per-layer LR                     |                                               |       |       | |-------------------|----------------------------------|-----------------------------------------------|-------|-------| | default constants | (full align) + default constants | Per-layer LR (full align) + optimal constants |       |       | | STP               | 3.657                            | 3.510                                         | 3.385 | 3.318 | | NTK               | 3.732                            | 3.627                                         | 3.313 | 3.324 | | muP               | 4.224                            | 4.184                                         | 3.809 | 4.222 | | MFP               | 4.092                            | 4.319                                         | 4.092 | 3.898 | | STP               | 2.625                            | 2.576                                         | 2.572 | 2.575 | | NTK               | 2.592                            | 2.615                                         | 2.591 | 2.581 | | muP               | 2.727                            | 2.646                                         | 2.599 | 2.590 | | MFP               | 2.638                            | 2.630                                         | 2.633 | 2.638 | | STP               | 2.580                            | 2.613                                         | 2.675 | 2.577 | | NTK               | 2.570                            | 2.623                                         | 2.667 | 2.566 | | muP               | 2.574                            | 2.655                                         | 2.606 | 2.575 | | MFP               | 2.624                            | 2.640                                         | 2.772 | 2.623 |  Adam  STP 2.625 2.576 **2.572** 2.575  NTK 2.592 2.615 2.591 **2.581**  muP 2.727 2.646 2.599 **2.590** MFP 2.638 **2.630** 2.633 2.638  Adam+PS  STP 2.580 2.613 2.675 **2.577**  NTK 2.570 2.623 2.667 **2.566** muP **2.574** 2.655 2.606 2.575  MFP 2.624 2.640 2.772 **2.623**  optimal learning rate at each model dim n with a sweep of the base learning rate βn, then fit a power law to the optimal base learning rate vs model dim. If the base learning rate were exactly the same at all scales, we would see a power law exponent of zero. We find exponents for standard, NTK, and muP of -0.05, -0.02, and -0.06, respectively, where these exponents close to zero indicate that the base learning rate is almost perfectly scale-invariant. For mean-field parameterization there is slight deviation from zero with an exponent of 0.26. This illustrates hyperparameter transfer for all parameterizations: when the constant factors are consistent across scale, the scaling prescription has correctly encapsulated the dependence on the scaling dimension into the prescribed exponents. We can therefore find the optimal base learning rate on a smaller model and reuse it on all model sizes.  Finally, we investigate the impact of the alignment assumptions, still using per-layer learning rates and optimal constants. In Figure 3(d), we use the learning rate exponents from the right side of Table 1 derived from assuming no alignment rather than full alignment. For NTK, muP and MFP, we see slight improvements in the eval losses across model sizes when using the no alignment exponents. On the surface, this improved performance would indicate that the no alignment exponents are preferable. However, when we look at the learning rate sweeps in Figure E1, the power law exponents fit to the optimal base learning rates for standard, NTK and muP are in the range of -0.58 to -0.69. Since these exponents are not close to zero, despite the modest performance improvements using the no alignment exponents over the full alignment exponents, the no alignment exponents do not appear to capture the learning rate scaling behavior as well as the full alignment exponents. For mean-field  | Per-layer LR (no align) + optimal constants   | |-----------------------------------------------|  parameterization, it is less clear what the optimal learning rate exponents are as both the full alignment and no alignment settings have power law exponents slightly above zero.  This indicates that in practice, there may be some nuance in selecting the alignment assumptions that are optimal for a particular use case and that there may be multiple interesting choices of exponents for a given parameterization within our more general space of parameterizations.  We note these empirical results across alignment assumptions could occur due to a mechanism other than the activation or logit growth with respect to width that is predicted by alignment in dense layers. For example, prior work notes training instabilities in Transformers due to attention logit growth (Dehghani et al., 2023; Zhai et al., 2023) or output logit divergence (Chowdhery et al., 2023) and that these instabilities are sensitive to the learning rate (Wortsman et al., 2023). We may also miss slight undercorrection of the alignment in our finite size models: if the true alignment is slightly larger than the learning rate exponents account for, we could have slow growth of activations or logits with respect to width that does not harm performance in our models but would eventually induce instability at sufficient width.  However, our largest model with 26.8B parameters has a model dimension of 16,384 which encompasses the width of many even larger models. As such, even if additional training instabilities may occur, we expect these per-layer learning rate and constant factor prescriptions to be relevant in practice to width scaling in large Transformers.  While muTransfer (Yang et al., 2022) emphasizes that muP is the unique parameterization that allows hyperparameter transfer across width, our results show that all parameterizations can perform hyperparameter transfer across width when each parameterization uses theoretically motivated per-layer learning rate exponents. In addition, our eval loss results for the full alignment exponents contrast with the empirical results in Yang et al. (2022) where muP outperforms standard parameterization in a 6.7B parameter GPT-3 model (Brown et al., 2020). We note, however, that their comparison across parameterizations has several elements that favor muP. First, the muP experiments use per-layer learning rate scaling while the standard parameterization experiments use a global learning rate. Second, the muP results tune a handful of constant factors at small scale and transfer them to large scale. While the specific constant factors differ from our setting, this is comparable to using our optimal constant learning rate factors instead of the default constant factors. As such, their comparison of muP and standard parameterization is analogous to comparing muP in our third experiment (Figure 3(c), green curve) against standard parameterization in our first experiment (Figure 3(a), blue curve), which indeed shows a benefit for muP. Instead, we argue that the fair comparison across parameterizations would use per-layer learning rates and optimal constants for both, that is, to consider both parameterizations in our third experiment Figure 3(c). There, we see that standard parameterization outperforms muP and in fact substantially so: the second largest standard parameterization model with 15.3B parameters outperforms the *largest* muP model with 26.8B parameters despite having 57% as many parameters. We therefore recommend our full alignment per-layer learning rate prescription for standard parameterization for Adam where the embedding layer learning rate scales like O(1) and the hidden and readout layer learning rates scale like O(1/n). 
body
1. Adam
2407.05872
 We next present the results for the same series of experiments for SGD and Adam + parameter scaling in Figure 5, where we use (a) global learning rates + default constants, (b) per-layer learning rates (full alignment) + default constants, (c) per-layer learning rates (full alignment) + optimal constants, and (d) per-layer learning rates (no alignment) + optimal constants. As we might expect, the results for SGD have significantly worse performance than the other optimizers, and the eval loss values are quite noisy. The impact of the different learning rate exponents is difficult to distinguish from noise for SGD, but there are visible performance improvements when using the optimal constants over the default constants.  For the per-layer learning rate experiments in the Adafactor family of optimizers, we use Adam + parameter scaling as the optimizer. The Adafactor optimizer was proposed in Shazeer & Stern (2018) and includes three types of changes to Adam. First, it reduces the memory requirements for the optimizer state by using a low-rank approximation of the gradient second moment and eliminating the exponential moving average on the first moment. Second, it introduces a parameter scaling term in the update rule. Third, it perform update clipping rather than gradient clipping, which stabilizes the update when the second moment estimate is out-of-date due to the moving average.  From the perspective of width-scaling, the most important of these differences is the parameter scaling, which multiplies the normalized gradients by the norm of the existing parameters. To focus our investigation on the impact of this parameter scaling term, we use Adam + parameter scaling as the optimizer in this section, so that the parameter scaling term is the sole change from our implementation of Adam.  Additionally, the low-rank approximation in Adafactor introduces changes in the tensor shapes, which caused issues out-of-the-box with our implementation of fully sharded data parallelism (FSDP) (Rajbhandari et al., 2020). The choice of Adam + parameter scaling avoids these issues. As a cross-check, in Appendix H we compare Adafactor and Adam + parameter scaling and show that the performance differences are small and that overall these two optimizers occupy the same width-scaling regime.  Overall, we see very good empirical performance from Adam + parameter scaling. When comparing the best setting for each parameterization across Adam with or without parameter scaling, the addition of parameter scaling slightly improves the eval loss for all parameterizations except standard. In particular, the best performing model across all parameterizations and optimizers is quite unexpected: it is Adam + param scaling + NTK + no alignment! This result further validates that it is critical to consider our more general space of parameterizations rather than making specific alignment assumptions upfront.  For the experiments in Figure 5 across different settings with global, full alignment per-layer and no alignment perlayer learning rate exponents, we first note that for Adam + parameter scaling, the no alignment setting is equivalent to the global learning rate setting. Concretely, when we compute the maximum stable learning rate in the no alignment setting in the rightmost column of Table 1, for all parameterizations we get O(1) scaling in all layers. This means the per-layer exponents in the no alignment setting do not modify the base learning rate, which is equivalent to the global learning rate setting that does not apply these perlayer exponents. Conceptually, this illustrates an interesting theoretical property of parameter scaling: in a sense, it accomplishes with the optimizer alone what the width-scaling theory of parameterization intends when carefully selecting the learning rates to preserve the scale of the activations and logits. Recall that the stability constraints ensure that, at initialization, the parameters contribute to constant activations and at-most-constant logits. If we neglect for a moment the contribution of alignment, parameter scaling makes the  ![13_image_0.png](13_image_0.png)   updates to each layer exactly the "right" scale by matching the existing parameter norms, allowing every layer to use a constant scale learning rate regardless of the parameterization. If we do account for alignment, rather than using constant learning rates, we reduce the hidden and/or readout layer learning rates to 1/ √n to account for the extra √n contribution from alignment. The parameter scaling therefore makes all parameterizations behave similarly, which we see in Figures L1 and L2 where all parameterizations use similar values for the base learning rate and have similar scale-dependence of the optimal learning rate under each set of learning rate exponents.  With this equivalence of the no alignment and global learning rate settings in mind, we can now compare the no alignment and full alignment settings for Adam + parameter scaling. For three reasons, the no alignment exponents appear preferable to the full alignment exponents. First, the eval losses from the no alignment (or global) exponents in Figure 5(a) and (d) outperform the full alignment exponents in Figure 5(b) and (c). Moreover, the performance gap increases with scale suggesting the no alignment exponents are truly superior. Second, the full alignment experiments have very high learning rate sensitivity: this is seen in the higher curvature in the eval loss vs learning rate curves in Appendix L where a small deviation from the optimal learning rate causes a larger loss in performance. In addition, the optimal learning rate is quite close to the maximum stable learning rate, which is a risky scenario for training stability.  In general, high learning rate sensitivity is an undesirable property: when all else is equal, it is preferable to have a model that can train well with a larger range of learning rates than one that requires a very narrow range of learning rates for success. (Wortsman et al., 2023) Third, the optimal constant multipliers transfer across scales well for the no alignment exponents but actually harm performance compared to the default constants for full alignment. Recall that the optimal constants should transfer well across scale when the learning rate scaling is well-encapsulated in the per-layer learning rate exponents. The positive transfer of these constants under the no alignment exponents compared to the negative transfer under the full alignment exponents indicates that the no alignment exponents better encapsulate the learning rate scaling into the exponents.  However, despite these reasons to prefer the no alignment exponents for Adam + parameter scaling, we see an interesting phenomenon when we look at the power law exponents on the optimal learning rate in the full alignment and no alignment settings. Across all combinations of the parameterization, default or optimal constants, and full or no alignment, in Figures L1 and L2 we see similar power law exponents in the range of −0.29 to 0.05. In particular, both the full alignment and no alignment settings are close to scale-invariant, but neither is clearly more scale-invariant than the other. This contrasts with the Adam setting, where the full alignment setting had optimal learning rate power law exponents close to zero, and the no alignment setting had power law exponents close to −0.5. In that setting, the exponent difference of approximately 0.5 corresponded with the √n difference in the prescribed learning rate exponents between the two alignment settings. This showed that the full alignment setting was encapsulating the scaling behavior into the prescribed exponents whereas the no alignment setting required the base learning rate to change across scale to absorb what wasn't captured by the prescribed learning rate exponents. In contrast, for Adam + parameter scaling, the power law exponents that are slightly less than zero in both the full alignment and no alignment settings indicate that parameter scaling may introduce factors that influence the optimal learning rate scaling that are more complex than the width-scaling dependence we analyze here.  Although Adafactor has been much less widely adopted than Adam, it has been used successfully in large model training (Raffel et al., 2020; Chowdhery et al., 2023; Du et al., 2022; Fedus et al., 2022; Zoph et al., 2022). However, several papers have reported difficulties with training stability and brittleness to the learning rate or other hyperparameters (Rae et al., 2021b; Zhai et al., 2021). Both the empirical success and this brittleness are consistent with our results. In our experiments with parameter scaling, we observe eval losses that are similar or better than for Adam, and in particular the best eval loss across all optimizers and parameterizations uses Adam + parameter scaling and NTK. In addition, we note two features that may contribute to the empirical success of parameter scaling. First, the global learning rate setting coincides with the theoretically motivated no alignment exponents. Second, the default constants equal to one are not far from the optimal constants (see §C3), which range in value from 0.5 to 2.6. This contrasts with the larger range in the optimal constants for Adam that span 0.15 to 11.7. These two properties may help adaptive optimizers with parameter scaling perform well in the typical setting in practice that uses global learning rates and no perlayer learning rate constants. At the same time, we observe high learning rate sensitivity that may contribute to training stability difficulties or brittleness to optimizer hyperparameter choices. In particular, in our largest models, the optimal learning rate is close to the maximum stable learning rate which may cause issues with training stability.  In conclusion, for adaptive optimizers with parameter scaling, we see similar width-scaling behavior from all parameterizations. We recommend using the no alignment, or equivalently, global learning rate exponents, for improved performance, and close to scale-invariance in the base learning rate. Unlike Adam, tuning the per-layer constant multipliers has minor performance impact and should not be considered essential. In addition, future work is needed to more clearly understand the impact of parameter scaling on the training stability and hyperparameter sensitivity. 
body
2. Sgd And Adam + Parameter Scaling
2407.05872
 In adaptive optimizers like Adam, the denominator of the update rule adds a small epsilon parameter to the gradient second-order moment, originally intended to regularize against division by zero when the gradients are very small (Duchi et al., 2011; Hinton et al., 2012). More recently, Choi et al. (2019) shows that epsilon is a hyperparameter that requires tuning. Despite its small value, typically around 1e-8 (Paszke et al., 2019; Babuschkin et al., 2020), the epsilon parameter prevents Adam from being perfectly scale-invariant in that multiplying the gradients by a constant would not alter the resulting update. In particular, if the gradient scale drops below the size of epsilon then epsilon dominates the gradients instead of acting as a negligible additive constant. Since gradients decrease in scale with model width as in Table 1, in theory, for any constant value of epsilon there exists a sufficiently large model that will encounter this scenario.  Two recent works note this phenomenon and propose possible mitigations. From an empirical perspective, Wortsman et al. (2023) observes that gradient norms in standard parameterization models decrease with model size and approach 1e-8 for 1.2B parameter models, suggesting this epsilon underflow is relevant in practice. As a mitigation, they propose using a smaller constant value for epsilon and show that decreasing epsilon from 1e-8 to 1e-15 improves the loss in a 4.6B parameter model. From a theoretical perspective, Yang & Littwin (2023) notes that epsilon should be treated as part of the parameterization and propose per-layer epsilon scaling. For each layer, epsilon is proportional to the parameterization and layer-specific gradient scale shown in Table 1. Similar to our per-layer learning rate experiments, we implement this as ϵl = base epsilon · ( n b ) −gl for each layer l, where glis the negative exponent of the gradient scale, the base model dim b is 1024, and the base epsilon is determined empirically. This approach has not been empirically validated and adds significant implementation complexity, but should ensure that epsilon and the gradients scale in tandem across width.  We investigate the practical impact of epsilon across parameterizations. As seen in the gradients column of Table 1, gradients decrease as model width increases with an exponent specific to both the parameterization and layer. As such, we expect that different parameterizations will encounter epsilon underflow at different model sizes: in particular, the steepest exponent is the mean-field parameterization hidden layer, which scales like 1/n1.5. This suggests that meanfield parameterization should encounter epsilon underflow  ![15_image_0.png](15_image_0.png)  0.00 0.01 0.02 0.03 0.04 0.05  
body
Epsilon Underflow In Adaptive Optimizers
2407.05872
 Both constant epsilon and per-layer epsilon require hyperparameter tuning, to select the constant or the base constant multiplier, respectively. As shown in Figure 6 for mean field parameterization and Figure F1 for all parameterizations, constant values that are too large lead to suboptimal performance; values that are too small lead to instability, presumably by failing to prevent the numerical instability epsilon was originally intended to prevent. Rather than providing numerical stability for small number division with an additive constant in the denominator that breaks scaleinvariance, we propose *Adam-atan2*: a variant of Adam that replaces the standard division operation in the update rule with the standard library function atan2. The function atan2(*x, y*) returns arctan(x/y) in the appropriate quadrant, which is approximately equal to x/y due to small-angle approximation when x/y is close to zero and asymptotically approaches ± π/2 as the argument goes to ± ∞. In particular, the atan2 function is defined even at (0, 0) exactly and is scale-invariant up to precision limits. The single-line code change in Appendix C.5 to use arctangent in the Adam update equation eliminates the epsilon hyperparameter entirely and restores the scale invariance to Adam. For all parameterizations, we compare the three mitigations using the best choice of constant in each setting against a baseline epsilon (1e-9): small constant epsilon (1e-15), per-layer epsilon scaling (base epsilon = 1e-12), and our proposal *Adam-atan2*. In Figure 7(a) and (b), we see that all three mitigations have modest improvements in the eval loss for NTK and and substantial improvements for meanfield parameterization. In addition, the gap in performance increases with model size. The other two parameterizations show no performance changes in our model sizes (see Figure F2) with any of the epsilon mitigations. The particular sensitivity of mean-field parameterization to epsilon is consistent with our theoretical motivation: due to its hidden layer gradients scaling like 1/n1.5, we expect that among the parameterizations, mean-field will encounter epsilon underflow at the smallest model sizes and benefit most from these mitigations.  We therefore recommend care when setting epsilon. For standard parameterization models with up to a billion parameters, the typical default value of 1e-8 is likely acceptable but slightly smaller values of 1e-12 or 1e-15 may be preferable. For larger models or other parameterizations, using epsilon requires smaller constants or per-layer epsilon scaling, with tradeoffs between implementation complexity and, at least theoretically, hyperparameter tuning costs. In principle, the theoretical prescription for per-layer epsilon encapsulates the epsilon scaling in the exponents, allowing hyperparameter transfer of the constant multiplier similar to other parameterized quantities like learning rates. In contrast, the optimal constant epsilon should be scale-dependent in theory, but we note that with our model sizes and constant search resolution we did not see scale dependence of the optimal epsilon constant. However, this epsilon hyperparameter can be eliminated entirely with *Adam-atan2* with a one-line code change and the same improved performance.  Moreover, epsilon illustrates that finite precision plays an important role in parameterization in practice. Recall that standard and NTK parameterizations, and similarly muP and mean-field parameterizations, are theoretically equivalent under the equivalence relations in Appendix B.2, if we overlook the contribution of epsilon. However, using default epsilons in Figure 7(c), we see significant performance gaps between equivalent parameterizations that are closed when epsilon underflow is mitigated in Figure 7(d): now the pairs standard + NTK and muP + mean-field show approximately equal performance. It is plausible that the more widespread usage of standard and muP parameterizations over their equivalents has been influenced by this phenomenon.  ![16_image_0.png](16_image_0.png)  Lastly, we can now compare the performance across equivalence classes rather than individual parameterizations: we see that the standard parameterization equivalence class outperforms the muP equivalence class. This narrows the possible explanations for the performance differences to the elements that we saw distinguish the equivalence classes in Section 3.2: namely, the shift in the readout layer in muP and mean-field parameterization that initializes the logits to scale like 1/ √n appears to harm the empirical performance.  Instead, to obtain the best performance for Adam, we recommend the parameterizations that initialize the logits to be constant scale, in particular our prescription for standard parameterization with per-layer learning rates where the embedding layer learning rate scales like O(1) and hidden and readout layer learning rates scale like O(1/n). 
body
At Smaller Model Sizes Than Other Parameterizations.
2407.05872
 In addition to the prior work in Yang & Hu (2021); Yang & Littwin (2023) discussed earlier in the paper, the literature on width-scaling parameterizations (Lee et al., 2017; Matthews et al., 2018; Jacot et al., 2018) includes the neural tangent kernel parameterization (Jacot et al., 2018), a modification to standard parameterization to enable a consistent infinite-width limit (Sohl-Dickstein et al., 2020), and a mean-field limit (Mei et al., 2018; Geiger et al., 2020; Rotskoff & Vanden-Eijnden, 2018; Chizat & Bach, 2018; Araujo et al. ´ , 2019) for single-hidden-layer MLPs. Bordelon & Pehlevan (2022) proposed the mean-field parameterization by extending this mean-field limit to deep neural networks using self-consistent dynamical field theory. Yaida (2022) proposed a one-parameter family of hyperparameter scaling strategies that interpolates between the neural tangent scaling and mean-field or muP scaling. This metaparameterized scaling strategy has been used to propose width-scaling initialization and training hyperparameters in Transformers (Dinan et al., 2023). Various "unit scaling" strategies proposed by Shazeer (2020); Kaplan (2019) are similar to NTK parameterization. In addition, Blake et al.  (2023) also modifies the per-layer gradients to keep the activations and parameters close to unit scale regardless of model size. Throughout this paper, as in Yang & Hu (2021) and Yang & Littwin (2023), we use the RMS norm to quantify the scale of activations and logits, but other choices of norm are possible including the spectral norm used in Yang et al. (2023a) to provide a spectral perspective on muP, and a modular norm proposed in Large et al. (2024). Ishikawa & Karakida (2023) extends muP to second-order optimizers (K-FAC, Shampoo).  Empirical evaluations of muP using Adam in Transformers were first presented in muTransfer (Yang et al., 2022) on models up to 6.7B parameters. Cerebras-GPT (Dey et al., 2023) found that muP with per-layer learning rates aided hyperparameter transfer and outperformed standard parameterization with global learning rates in compute-optimal models up to 2.7B parameters. In addition, Lingle (2024) showed that the optimal learning rate transfer across width for muP held on ablations of many architectural and optimal choices in models up to 1.2B parameters with learning rate granularity of 4×, but that the optimal learning rate drifted by a factor of 2× between 40M parameter and 10B parameter models.  Other important scaling dimensions include depth and batch size. For depth scaling, Yang et al. (2023b), Bordelon et al.  (2023), and Chizat & Netrapalli (2023) make similar proposals that add a 1/ √L scaling factor to the residual branch of a ResNet or Transformer where L is the depth. Batch size scaling is investigated in (Shallue et al., 2019; McCandlish et al., 2018; Zhang et al., 2019; Kaplan et al., 2020). It is common to increase the batch size with model size as in (Brown et al., 2020; Rae et al., 2021a; Hoffmann et al., 2022; Chowdhery et al., 2023; Scao et al., 2022; Dey et al., 2023; Bi et al., 2024) but there are exceptions that use fixed batch sizes (Thoppilan et al., 2022; Touvron et al., 2023a;b) across scale. 
body
Related Work
2407.05872
 This paper focuses specifically on width scaling, while in practice there are a number of dimensions that need to scale in tandem to achieve optimal performance in large models.  In particular, large models typically co-scale at least the width, depth, batch size, training horizon, weight decay, and learning rate. For specific pairs or subsets of these scaling dimensions, such as batch size and learning rate or the aspect ratio between width and depth, the literature contains some theoretical insights and empirical validation, but determining how to optimally co-scale this entire set of dimensions is a complex and open problem. As we discuss in Appendix I, co-scaling the width and training horizon as in the compute-optimal regime requires adjustments to the learning rate exponents, but developing theory with realistic implications would require alternate approaches.  In addition, our theory considers the input and output dimensionality, corresponding to the vocabulary size in a Transformer language model, to be an O(1) constant with respect to width. However, our experiments use a vocabulary size of 32,000 that is typical for our range of model sizes. We do hold the vocabulary size fixed across our model sizes, but the magnitude of 32,000 is comparable to model widths that range up to D = 16,384, and it might be more realistic to consider the vocabulary size to be O(n) rather than O(1). 
body
Limitations
2407.05872
 From our broad perspective across parameterizations and optimizers, we find that key assumptions about alignment in prior work require additional consideration. Using empirical measurements from our alignment ratio metric, we find that alignment is a dynamical quantity that depends significantly on the training step, parameterization and layer, and less heavily on the optimizer. The alignment measured during training gives intermediate values that indicate that alignment assumptions in prior work may be overly conservative, suggesting that a larger set of parameterizations is more interesting than previously thought.  By considering a more general space of parameterizations with respect to the alignment, we show that all parameterizations benefit from theoretically motivated learning rate exponent prescriptions. We also demonstrate that several hyperparameters should be chosen carefully. First, we show the necessity and practicality of tuning the constant factors in per-layer learning rate prescriptions. These constants transfer well across model sizes, showing that all parameterizations can perform hyperparameter transfer under the right theoretical prescription. Second, the epsilon hyperparameter in adaptive optimizers induces gradient underflow using typical defaults at realistic model sizes, in particular for meanfield parameterization. Theoretically, epsilon should be considered part of the parameterization and scaled per-layer, but practically, small constant values can perform just as well when selected carefully. To eliminate epsilon entirely, we propose making Adam scale-invariant with *Adam-atan2*.  Future work might consider alignment-aware learning rate schedules or alignment-aware optimizers. In addition, since the characterization of parameterizations into feature learning and kernel limits is specific to the alignment assumptions, this characterization could be extended to the general alignment setting. Beyond width scaling, future work should investigate the other scaling dimensions that are necessary for large model training, in particular depth and batch size. 
conclusion
Conclusions And Future Work
2407.05872
 The authors would like to thank Yasaman Bahri and Kevin Yin for detailed feedback on paper drafts, and Kevin Yin for the suggestion to scale the arctangent function as λ · atan2(*x, λy*). We also thank Ben Adlam, Kelvin Xu, Justin Gilmer, Gamaleldin Elsayed, Mark Kurzeja, Jiri Hron, Noah Fiedel, Zachary Nado, Justin Austin, Ben Poole, Clare Lyle and Tomas Lozano-P ´ erez for technical discussions and are ´ grateful to four anonymous ICML reviewers who provided especially high-quality reviews and discussion during the rebuttal period. 
body
Acknowledgements
2407.05872
 This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, in particular due to the large model sizes considered in this work, but we do not feel there are specific aspects of this work with broader impacts beyond the considerations relevant to all large machine learning models. 
body
Impact Statement
2407.05872
 Araujo, D., Oliveira, R. I., and Yukimura, D. A mean-field ´ limit for certain deep neural networks. *arXiv preprint* arXiv:1906.00193, 2019.  Ba, J. L., Kiros, J. R., and Hinton, G. E. Layer normalization.  arXiv preprint arXiv:1607.06450, 2016.  Babuschkin, I., Baumli, K., Bell, A., Bhupatiraju, S., Bruce, J., Buchlovsky, P., Budden, D., Cai, T., Clark, A., Danihelka, I., Dedieu, A., Fantacci, C., Godwin, J., Jones, C., Hemsley, R., Hennigan, T., Hessel, M., Hou, S., Kapturowski, S., Keck, T., Kemaev, I., King, M., Kunesch, M., Martens, L., Merzic, H., Mikulik, V., Norman, T., Papamakarios, G., Quan, J., Ring, R., Ruiz, F., Sanchez, A., Sartran, L., Schneider, R., Sezener, E., Spencer, S., Srinivasan, S., Stanojevic, M., Stokowiec, W., Wang, L., ´ Zhou, G., and Viola, F. The DeepMind JAX Ecosystem, 2020. URL http://github.com/deepmind.  Bi, X., Chen, D., Chen, G., Chen, S., Dai, D., Deng, C., Ding, H., Dong, K., Du, Q., Fu, Z., et al. Deepseek llm: Scaling open-source language models with longtermism.  arXiv preprint arXiv:2401.02954, 2024.  Blake, C., Orr, D., and Luschi, C. Unit scaling: Out-of-the-box low-precision training. *arXiv preprint* arXiv:2303.11257, 2023.  Bordelon, B. and Pehlevan, C. Self-consistent dynamical field theory of kernel evolution in wide neural networks.  Advances in Neural Information Processing Systems, 35: 32240–32256, 2022. Bordelon, B., Noci, L., Li, M. B., Hanin, B., and Pehlevan, C. Depthwise hyperparameter transfer in residual networks: Dynamics and scaling limit. arXiv preprint arXiv:2309.16620, 2023.  Bradbury, J., Frostig, R., Hawkins, P., Johnson, M. J., Leary, C., Maclaurin, D., Necula, G., Paszke, A., VanderPlas, J., Wanderman-Milne, S., and Zhang, Q. JAX: composable transformations of Python+NumPy programs, 2018. URL http://github.com/google/jax.  Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al. Language models are few-shot learners.  Advances in neural information processing systems, 33: 1877–1901, 2020.  Chizat, L. and Bach, F. On the global convergence of gradient descent for over-parameterized models using optimal transport. Advances in neural information processing systems, 31, 2018.  Chizat, L. and Netrapalli, P. Steering deep feature learning with backward aligned feature updates. arXiv preprint arXiv:2311.18718, 2023.  Choi, D., Shallue, C. J., Nado, Z., Lee, J., Maddison, C. J., and Dahl, G. E. On empirical comparisons of optimizers for deep learning. *arXiv preprint arXiv:1910.05446*, 2019.  Chowdhery, A., Narang, S., Devlin, J., Bosma, M., Mishra, G., Roberts, A., Barham, P., Chung, H. W., Sutton, C., Gehrmann, S., et al. Palm: Scaling language modeling with pathways. *Journal of Machine Learning Research*, 24(240) :1–113, 2023.  Dehghani, M., Djolonga, J., Mustafa, B., Padlewski, P., Heek, J., Gilmer, J., Steiner, A. P., Caron, M., Geirhos, R., Alabdulmohsin, I., et al. Scaling vision transformers to 22 billion parameters. In *International Conference on* Machine Learning, pp. 7480–7512. PMLR, 2023.  Dey, N., Gosal, G., Khachane, H., Marshall, W., Pathria, R., Tom, M., Hestness, J., et al. Cerebras-gpt: Open computeoptimal language models trained on the cerebras waferscale cluster. *arXiv preprint arXiv:2304.03208*, 2023.  Dinan, E., Yaida, S., and Zhang, S. Effective theory of transformers at initialization. *arXiv preprint* arXiv:2304.02034, 2023.  Du, N., Huang, Y., Dai, A. M., Tong, S., Lepikhin, D., Xu, Y., Krikun, M., Zhou, Y., Yu, A. W., Firat, O., et al. Glam: Efficient scaling of language models with mixture-ofexperts. In *International Conference on Machine Learning*, pp. 5547–5569. PMLR, 2022. Duchi, J., Hazan, E., and Singer, Y. Adaptive subgradient methods for online learning and stochastic optimization.  Journal of machine learning research, 12(7) , 2011.  Fedus, W., Zoph, B., and Shazeer, N. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. *Journal of Machine Learning Research*, 23(120) :1–39, 2022.  Geiger, M., Spigler, S., Jacot, A., and Wyart, M. Disentangling feature and lazy training in deep neural networks.  Journal of Statistical Mechanics: Theory and Experiment, 2020(11) :113301, 2020.  Glorot, X. and Bengio, Y. Understanding the difficulty of training deep feedforward neural networks. In *Proceedings of the thirteenth international conference on* artificial intelligence and statistics, pp. 249–256. JMLR Workshop and Conference Proceedings, 2010.  Golovin, D., Solnik, B., Moitra, S., Kochanski, G., Karro, J., and Sculley, D. Google vizier: A service for blackbox optimization. In *Proceedings of the 23rd ACM* SIGKDD International Conference on Knowledge Discovery and Data Mining, Halifax, NS, Canada, August 13 - 17, 2017, pp. 1487–1495. ACM, 2017. doi: 10.  1145/3097983.3098043. URL https://doi.org/ 10.1145/3097983.3098043.  He, K., Zhang, X., Ren, S., and Sun, J. Delving deep into rectifiers: Surpassing human-level performance on imagenet classification. In *Proceedings of the IEEE international conference on computer vision*, pp. 1026–1034, 2015.  Heek, J., Levskaya, A., Oliver, A., Ritter, M., Rondepierre, B., Steiner, A., and van Zee, M. Flax: A neural network library and ecosystem for JAX, 2023. URL http:// github.com/google/flax.  Hendrycks, D. and Gimpel, K. Gaussian error linear units (gelus) . *arXiv preprint arXiv:1606.08415*, 2016.  Hinton, G., Srivastava, N., and Swersky, K. Neural networks for machine learning lecture 6a overview of mini-batch gradient descent. *Cited on*, 14(8) :2, 2012.  Hoffmann, J., Borgeaud, S., Mensch, A., Buchatskaya, E., Cai, T., Rutherford, E., Casas, D. d. L., Hendricks, L. A., Welbl, J., Clark, A., et al. Training compute-optimal large language models. *arXiv preprint arXiv:2203.15556*, 2022.  Ishikawa, S. and Karakida, R. On the parameterization of second-order optimization effective towards the infinite width. *arXiv preprint arXiv:2312.12226*, 2023. Jacot, A., Gabriel, F., and Hongler, C. Neural tangent kernel: Convergence and generalization in neural networks.  Advances in neural information processing systems, 31, 2018.  Kaplan, J. Notes on contemporary machine learning for physicists, 2019.  Kaplan, J., McCandlish, S., Henighan, T., Brown, T. B., Chess, B., Child, R., Gray, S., Radford, A., Wu, J., and Amodei, D. Scaling laws for neural language models.  arXiv preprint arXiv:2001.08361, 2020.  Kingma, D. P. and Ba, J. Adam: A method for stochastic optimization. *arXiv preprint arXiv:1412.6980*, 2014.  Kudo, T. and Richardson, J. Sentencepiece: A simple and language independent subword tokenizer and detokenizer for neural text processing. *arXiv preprint* arXiv:1808.06226, 2018.  Large, T., Liu, Y., Huh, M., Bahng, H., Isola, P., and Bernstein, J. Scalable optimization in the modular norm. *arXiv* preprint arXiv:2405.14813, 2024.  Lee, J., Bahri, Y., Novak, R., Schoenholz, S. S., Pennington, J., and Sohl-Dickstein, J. Deep neural networks as gaussian processes. *arXiv preprint arXiv:1711.00165*, 2017.  Lingle, L. A large-scale exploration of µ-transfer. arXiv preprint arXiv:2404.05728, 2024.  Liu, P. J., Novak, R., Lee, J., Wortsman, M., Xiao, L., Everett, K., Alemi, A. A., Kurzeja, M., Marcenac, P., Gur, I., Kornblith, S., Xu, K., Elsayed, G., Fischer, I., Pennington, J., Adlam, B., and Dickstein, J.-S. Nanodo: A minimal transformer decoder-only language model implementation in JAX., 2024. URL http://github.  com/google-deepmind/nanodo.  Loshchilov, I. and Hutter, F. Decoupled weight decay regularization. In *International Conference on Learning* Representations, 2018.  Matthews, A. G. d. G., Rowland, M., Hron, J., Turner, R. E., and Ghahramani, Z. Gaussian process behaviour in wide deep neural networks. *arXiv preprint arXiv:1804.11271*, 2018.  McCandlish, S., Kaplan, J., Amodei, D., and Team, O. D.  An empirical model of large-batch training. *arXiv* preprint arXiv:1812.06162, 2018.  Mei, S., Montanari, A., and Nguyen, P.-M. A mean field view of the landscape of two-layer neural networks. *Proceedings of the National Academy of Sciences*, 115(33) : E7665–E7671, 2018.
references
References