Adam-mini: Use Fewer Learning Rates To Gain More

🎯 Pitch

More than 99.9% of Adam’s coordinate-wise learning rates are redundant within dense Hessian sub-blocks; Adam-mini exploits this by partitioning parameters block-wise and assigning just one good learning rate per block, matching or beating AdamW on models up to 13B while halving memory use. This slashes distributed-training communication overhead, yielding 49.6% higher throughput on 7B-parameter Llama 2 pre-training and 33% less wall-clock time—without changing any hyperparameters.

1. Executive Summary

This paper proposes Adam-mini, an optimizer that reduces Adam's memory footprint by 50% by cutting down the number of learning rates—replacing coordinate-wise 1/√v entries with block-wise scalars. Through empirical analysis on language models sized from 39M to 13B—including GPT-2 and Llama series pre-training, supervised fine-tuning, and RLHF—the authors demonstrate that Adam-mini performs on par or better than AdamW while using the same hyperparameters. The core mechanism is a Hessian-structure-based parameter partitioning that assigns a single learning rate per dense Hessian sub-block (e.g., partitioning query and key by attention heads, value and attn.proj by output neurons), founded on the finding that ≥99.9% of Adam's coordinate-wise learning rates can be harmlessly removed within each sub-block. The reduced memory enables 49.6% higher throughput than AdamW when pre-training Llama 2-7B on 2× A800-80GB GPUs—translating to 33% less wall-clock time—establishing that memory-efficient optimization benefits substantially from architecture-aware partitioning rather than generic low-rank factorization like Adafactor, but that these gains are most pronounced on mainstream Transformer architectures where the Hessian's near-block-diagonal structure is readily exploitable.

2. Context and Motivation

The Core Problem: Adam's Memory Cost Is a Bottleneck for LLM Training

The fundamental problem this paper addresses is deceptively simple: Adam is memory-hungry, and that hunger is expensive. Specifically, Adam requires storing two optimizer states per parameter—the first-order momentum m and the second-order momentum v—which together consume at least 2× the memory of the model parameters themselves. For a 7B-parameter model, this translates to approximately 56 GB for m and v alone, with gradients pushing the total optimizer memory to roughly 86 GB. That is larger than a single cutting-edge GPU (e.g., A100-80GB), forcing practitioners to resort to CPU offloading and optimizer state sharding (Rajbhandari et al., 2020). These workarounds increase latency and slow down training (Rajbhandari et al., 2021).

This memory burden matters for several reasons the authors articulate (Section 1):

Democratizing LLM research: Reducing optimizer memory "lowers the threshold of training LLMs and encourages participation from more diverse researchers, especially those with limited GPU resources."
Cost and energy savings: Fewer GPUs are required to train a model of a given size, leading to "substantial savings in both cost and energy."
Throughput and latency: Easing the burden of CPU offloading and model sharding can "enhance the throughput and accelerate the training process."

The memory problem is not merely a hardware limitation—it is a structural inefficiency in the optimizer design itself. Adam assigns an individual learning rate η/√v_i to every single parameter, meaning a billion-parameter model requires billions of learning rates. Whether this level of granularity is necessary, or whether much of v is redundant, was an open question before this work.

Why Modifying Adam Is Difficult—and Where Prior Work Falls Short

Modifying Adam's optimizer states is challenging for a fundamental reason: we lack a clear understanding of the role of m and v. The paper highlights this uncertainty (Section 1):

"It remains uncertain which components in Adam are indispensable for superior performance, and which components could be re-designed or improved."

This uncertainty has not prevented prior attempts at memory-efficient optimization, but it has limited their effectiveness. The paper identifies three categories of prior work, each with significant shortcomings:

Adafactor (Shazeer & Stern, 2018) and its variants. The most prominent memory-efficient alternative to Adam, Adafactor, reduces memory by applying low-rank factorization to v—approximating it as the outer product of two vectors rather than storing the full matrix. This cuts memory usage similarly to what Adam-mini achieves. However, the paper notes two critical problems:

Performance degradation: Adafactor is "not easy to tune and often performs worse than Adam" (Section 1). The authors cite evidence from Luo et al. (2023) and provide their own extensive experiments (Section 3.4) showing that both the original Adafactor and a modified version from Zhai et al. (2022) consistently underperform AdamW across model scales, despite extensive hyperparameter tuning.
Generic approach, not architecture-aware: Low-rank factorization is a "generic approach that could be applied broadly, but it does not leverage much problem-specific structure" (Section 1). Because it ignores the specific structure of neural network Hessians, it "does not work well on specific neural-net tasks."

The paper goes further to show that Adafactor has a higher latency than Adam-mini (Figure 13c) because it requires summing across both rows and columns of the weight matrix, while Adam-mini only computes means along one dimension—and operates on a much smaller state vector.

CAME (Luo et al., 2023), SM3 (Anil et al., 2019), and other lightweight optimizers. These methods similarly attempt to reduce v's memory by various means—SM3 by using a minimal-value selection over predetermined covers of the gradient, CAME by building on Adafactor's factorization. All achieve memory savings, but all show performance degradation in the authors' experiments (Figures 8, 10). The paper attributes this to the same fundamental issue: they make broad modifications to v without understanding which aspects of v are functionally necessary.

LAMB (You et al., 2019). An important clarification the authors make (Appendix A) is that LAMB—despite using the phrase "layer-wise learning rates"—is fundamentally different from Adam-mini. LAMB retains the full coordinate-wise 1/√v from Adam and adds an additional layer-wise scaling on top. It saves no memory over Adam and targets large-batch training, not memory reduction.

The Central Mystery: Is Coordinate-Wise Granularity Necessary?

The paper identifies a critical conceptual gap in the existing literature: nobody has systematically asked whether each parameter truly needs its own learning rate. The key observation is that neural network Hessians exhibit block heterogeneity—different parameter blocks (e.g., the query projection matrix vs. the MLP weight matrix) have dramatically different eigenvalue distributions, as shown by Zhang et al. (2024) and restated in Appendix E.2. This means different blocks do need different learning rates.

However, Adam goes substantially further: it assigns a unique learning rate to every individual parameter within each block. The number of parameters is orders of magnitude larger than the number of blocks. This asymmetry raises a natural question that the paper frames as its driving inquiry:

"Is it necessary to use a customized learning rate for each parameter? If not, how much can we save?" (Section 2.1, Q1)

Prior work had not asked this question in a principled, structure-aware way. Adafactor answered a related but different question: "can we approximate v with less memory?"—and the answer, shown by its performance gaps, is "yes, but poorly." The paper's insight is that there may be a structural reason within the Hessian itself that makes most coordinate-wise learning rates redundant.

Connecting to Hessian Structure Theory

The paper grounds its investigation in a classical but under-appreciated finding: the Hessian of neural networks is near-block-diagonal. This was reported by Collobert (2004, Section 7) and the paper reproduces it empirically for both MLPs (Figure 3) and Transformers (Figure 7). Collobert's analysis shows that after even one training step, cross-entropy loss causes off-diagonal-block Hessian entries to shrink to zero, because the factor p(x)(1-p(x)) in the Hessian expression rapidly vanishes as the model's predictions become more confident (Appendix C, Equation 3).

This observation is crucial because it changes the framing of the memory problem. If the Hessian consists of dense sub-blocks (not fully diagonal, but organized into independent blocks), then the question becomes: within each dense sub-block, does Adam's coordinate-wise preconditioning provide value? Or is a single learning rate per sub-block sufficient?

The paper provides both theoretical and empirical evidence for the latter (Section 2.1). Through numerical experiments on random quadratic problems with block-diagonal Hessians (Figure 4), they show that:

Across different blocks: Adam (coordinate-wise) outperforms a single global learning rate, confirming that block-level differentiation is necessary.
Within a single dense block: A single optimal learning rate outperforms Adam's coordinate-wise approach, even though Adam uses tens or hundreds of parameters for that block.
Combining block-wise optimal rates: A block-wise GD method using one learning rate per block converges faster than Adam on the full problem.

The paper explains this counterintuitive result through a linear algebra lens: Adam is a diagonal preconditioner, but there is no guarantee that a diagonal preconditioner reduces the condition number of a dense matrix. For a dense Hessian sub-block H_b, the condition number κ(D_Adam * H_b) can actually be larger than κ(H_b)—meaning the preconditioner can make the problem harder to optimize. Figure 5 quantifies this numerically: as the "diagonal-over-off-diagonal ratio" τ decreases (matrix becomes less diagonal/more dense), Adam's preconditioner effectiveness r (ratio of preconditioned to raw condition number) increases, often substantially. In dense regions, Adam's coordinate-wise learning rates may be actively unhelpful.

How This Paper Positions Itself

The paper's positioning can be understood through four strategic choices:

1. Not a wholesale replacement of Adam, but a structurally motivated simplification. Unlike prior work that attempted to replace v with a fundamentally different construct (low-rank factorization, minimal-value selection), Adam-mini preserves the core update rule but reduces the granularity of v to match the Hessian's structure. This is a more conservative modification anchored in an empirical observation about neural networks rather than a generic compression technique.

2. Hessian structure as the organizing principle. The paper explicitly frames its contribution around Principle 1 (Section 2.3): parameters should be partitioned such that each block corresponds to the smallest dense sub-block in the Hessian. This is not an optimization trick—it is a principle for designing optimizers that are aware of the problem's structure. For Transformers, this principle specifies exactly how to partition: query and key by attention heads; value, attn.proj, and mlp by output neurons; embed and output by tokens.

3. A direct rebuttal to the necessity of full coordinate-wise adaptivity. The paper challenges a widespread assumption in the optimization community. As the authors note (Section 1):

"This is possible as most existing Adam variants that attempt to modify v to varying extents have been reported to perform worse than Adam (Orabona, 2020)."

The implicit narrative had been: "Adam's v is essential in its full precision; any simplification hurts performance." Adam-mini provides counter-evidence at scale—across 39M to 13B parameter models, pre-training, supervised fine-tuning, and RLHF—that the granularity can be reduced by ≥99.9% without harming performance, if the reduction respects Hessian structure.

4. Practical, not just theoretical. Beyond the empirical performance results, the paper demonstrates concrete throughput improvements: 49.6% higher throughput than AdamW and 33% less wall-clock time for Llama 2-7B pre-training (Table 2, Figure 1). This connects the theoretical insight (Hessian block structure) to a practical outcome (faster training) that directly addresses the motivation of lowering barriers to LLM training.

The Gap This Paper Fills

To summarize the landscape before Adam-mini:

AdamW: High memory cost (2× the model size for m and v), but reliable performance.
Adafactor and variants: Lower memory cost (comparable to Adam-mini), but inconsistent performance and difficult tuning (9 hyperparameters vs. Adam's standard 3-4).
Hessian analysis work: Had established the near-block-diagonal structure of neural network Hessians, but this knowledge had not been translated into optimizer design.
Block-level adaptivity research: Zhang et al. (2024) showed that Transformers need different learning rates for different parameter blocks, but did not address whether individual parameters within blocks needed distinct rates.

Adam-mini fills the gap at the intersection of these lines of work: it uses the Hessian's block structure (from the analysis literature) to determine parameter partitions, applies the insight that block-level differentiation is necessary but coordinate-wise differentiation within blocks may be redundant (building on Zhang et al., 2024), and produces an optimizer that matches Adam's performance while matching Adafactor's memory savings—all without introducing new hyperparameters. The paper is the first to demonstrate that this specific synthesis—structure-aware, Hessian-guided reduction of learning rate granularity—is both theoretically grounded and practically effective at scale.

3. Technical Approach

3.1 Reader Orientation

Adam-mini is a drop-in replacement for AdamW that stores the second-order momentum v as a small collection of scalars—one per parameter block—rather than as a full tensor matching the model's parameter count. It solves the problem of Adam's excessive memory consumption by reducing the number of stored learning rates by ≥99.9%, but critically achieves this through an architecture-aware partitioning strategy derived from the near-block-diagonal structure of neural network Hessians rather than through generic tensor compression techniques that have historically degraded performance.

3.2 Big-Picture Architecture (Diagram in Words)

The system has four conceptual components, though in practice they collapse to a single optimizer implementation:

Hessian Structure Analyzer (offline, conceptual) — determines how the model's parameters should be partitioned into sub-blocks, based on empirical observation of which Hessian sub-blocks are dense (and thus need only one learning rate per block). This component produces a fixed partitioning rule encoded in the optimizer constructor.
Parameter Partitioner — at initialization time, applies the partitioning rule to the model's named parameters, grouping each parameter's gradient into sub-vectors g_b according to the architecture-specific mapping (e.g., query by heads, value by output neurons).
Block-wise Second-Moment Estimator — replaces Adam's element-wise v = β_2·v + (1-β_2)·(g⊙g) with a block-wise version v_b = β_2·v_b + (1-β_2)·mean(g_b⊙g_b), producing one scalar v_b per block rather than one scalar v_i per parameter. This is the core memory-saving mechanism.
Standard Adam Update Engine — applies the Adam update rule identically to AdamW, but with the block-level v_b scalar broadcast back to all parameters in the block: each parameter in block b receives the learning rate η/√(v_b/(1-β_2^t)). The first-order momentum m is stored at full precision (unchanged from Adam).

Information flows as follows: model parameters are partitioned once at optimizer construction → at each training step, gradients are computed normally → second-order momentum is accumulated as block-wise scalars → the scalar v_b for each block is broadcast and used to scale the first-order momentum for the parameter update → weight decay is applied identically to AdamW.

3.3 Roadmap for the Deep Dive

This is primarily an empirical analysis and design paper whose core idea is that Adam's coordinate-wise learning rates are redundant within dense Hessian sub-blocks, and that a Hessian-aware partitioning strategy enables dramatic memory reduction without performance degradation. The deep dive follows four stages:

First, the formal observation that motivates the entire approach: v provides learning rates, and the number of learning rates can be dramatically reduced if we understand the Hessian's block structure. This includes the evidence from random quadratic problems (Figures 4–5) and Transformers (Figures 6–7) that Adam's coordinate-wise granularity is unnecessary within dense Hessian blocks.
Second, the parameter partitioning principle (Principle 1) and its concrete realization for Transformers. This is the critical design choice that distinguishes Adam-mini from naive approaches: partitioning must follow the Hessian's smallest dense sub-blocks, not arbitrary layer boundaries. I walk through the specific partitioning rules for query/key (by attention heads), value/attn.proj/mlp (by output neurons), and embed/output (by tokens), and explain why the PyTorch default partition fails.
Third, the Adam-mini algorithm itself (Algorithm 1), focusing on the two-step procedure—partition, then apply block-wise v averaging—and its concrete form for Transformers (Algorithms 2–3). This includes a detailed example showing the precise difference in u (the effective learning rate vector) between Adam and Adam-mini.
Fourth, the characteristics of Adam-mini as a practical tool: memory savings (50% of Adam's total optimizer memory), throughput improvements (49.6% for Llama 2-7B), the relationship to Adam's trajectory, and the design choice of using mean(v) over alternatives.

3.4 Detailed, Sentence-Based Technical Breakdown

This is primarily an empirical analysis and design paper whose core idea is that Adam assigns far more learning rates than necessary, and that a Hessian-structure-aware reduction to block-level learning rates preserves or improves performance while halving memory usage. The paper makes its case through two parallel tracks: (1) numerical experiments on abstract quadratic problems and small Transformers that establish the conceptual claim that coordinate-wise learning rates are redundant within dense Hessian blocks, and (2) the design and large-scale validation of Adam-mini, a concrete optimizer that exploits this redundancy through a principled parameter partitioning strategy.

The Conceptual Engine: Redundancy of Coordinate-Wise Learning Rates in Dense Hessian Sub-Blocks

The paper does not assume that all of Adam's v is redundant—that would contradict evidence from Zhang et al. (2024) showing that different parameter blocks need different learning rates due to Hessian eigenvalue heterogeneity. Instead, the paper makes a more nuanced claim: within a single dense Hessian sub-block, coordinate-wise learning rates are redundant; a single learning rate per block can match or exceed Adam's performance. This sub-section builds the evidence for this claim, which is the intellectual foundation for Adam-mini's design.

Evidence from random quadratic problems (Section 2.1, Figures 4–5). The authors construct a controlled experiment using random quadratic minimization: min_w (1/2) w^T H w, where H is a random positive definite matrix with a block-diagonal structure composed of three dense sub-blocks (Figure 4a). The eigenvalues within each sub-block are deliberately spread across different scales: block 1 has eigenvalues sampled from {1, 2, 3}, block 2 from {99, 100, 101}, and block 3 from {4998, 4999, 5000}—creating a condition number mismatch across blocks.

The experiment compares three methods:

Adam (coordinate-wise): Each parameter w_i receives its own learning rate η/√v_i based on its individual gradient history, using β_2 = 1 (no decay on the second moment, which is necessary for convergence on quadratics, following Da Silva & Gazeau, 2020 Proposition 12, Figure 1) and β_1 = 0 (to isolate the effect of the second-order term).
Optimal single-learning-rate method: Gradient descent with the globally optimal constant learning rate 2/(L + μ), where L and μ are the largest and smallest eigenvalues of the full Hessian H.
Optimal blockwise-learning-rate method: Gradient descent where each of the three blocks receives its own optimal constant learning rate, using the largest and smallest eigenvalues of only that block's sub-matrix.

The results (Figure 4b) are striking:

On the full problem: Adam outperforms the optimal single-learning-rate method, confirming that block-level differentiation is necessary—the eigenvalue spread across blocks (from 1–3 in block 1 to 4998–5000 in block 3) means no global learning rate can handle all blocks well.
On each individual dense sub-block (Figures 4c, 4d): A single optimal learning rate outperforms Adam. For the first sub-block (eigenvalues 1–3), the optimal single rate converges substantially faster than Adam, even though Adam uses 30 distinct learning rates (one per parameter) while the optimal block method uses 1.
Combining block-wise optimal rates: The blockwise GD method (green line in Figure 4b) converges faster than Adam on the full problem, using only 3 learning rates total.

The paper's explanation for this phenomenon is based on linear algebra (Section 2.1):

"Adam can be viewed as a diagonal preconditioned method, i.e., at the t-th step: w_{t+1} = w_t - η_t D_t m_t, where D_t = Diag(1/√v_t) is a diagonal matrix."

Where D_t is the preconditioner (a diagonal matrix whose i-th entry is 1/√{v_{t,i}}), m_t is the first-order momentum, w_t are the model parameters, η_t is the global learning rate, and t indexes the training step.

The effectiveness of any diagonal preconditioner D is measured by how much it reduces the condition number: κ(D·H) compared to κ(H), where κ(·) is the condition number (the ratio of largest to smallest eigenvalue—smaller values mean faster linear convergence). The paper makes a critical observation:

"Unfortunately, there is no guarantee of κ(DH) ≤ κ(H) and this inequality often requires strict assumptions on both D and H."

In fact, for a dense Hessian sub-matrix, a diagonal preconditioner can increase the condition number, making optimization slower rather than faster. The paper quantifies this through the numerical experiment in Figure 5, which examines the relationship between two metrics:

$\tau = \frac{\sum_i |H_{b,i,i}|}{\sum_{i,j} |H_{b,i,j}|},\quad r = \frac{\kappa(D_{\text{Adam}}H_b)}{\kappa(H_b)}$

where H_b ∈ R^{d×d} is a random dense positive definite matrix (a proxy for a neural network's dense Hessian sub-block), τ ∈ [0, 1] is the "diagonal-over-off-diagonal ratio"—measuring how diagonally dominant the matrix is (τ = 1 means pure diagonal, τ ≈ 0 means highly dense), and r ≥ 0 measures the effectiveness of Adam's preconditioner (r < 1 means the preconditioner helps; r > 1 means it hurts). D_Adam is diag(1/√v), where v = g ⊙ g, g = H_b·x, and each entry x_i ∼ N(0, 1/√d) (Xavier initialization).

What this computes: τ sums the absolute values on the diagonal of H_b and divides by the sum of absolute values of all entries—a measure of how concentrated the matrix's mass is on the diagonal. r computes the condition number after applying Adam's preconditioner and divides by the raw condition number. If r < 1, the preconditioner reduces the condition number (helpful); if r > 1, the preconditioner makes the problem more ill-conditioned (harmful).

Why this form: the ratio-based formulation isolates the effect of off-diagonal entries. The experiment varies τ by rotating the eigenvectors of H_b without changing its eigenvalues, so the eigenvalue distribution is held constant—any change in r is purely due to how diagonal the matrix is, not due to changes in the eigenvalue spread.

The key findings from Figure 5:

For most dimensions d (10–100) and condition numbers κ(H_b) (200–1000), r increases (worsens) as τ decreases (matrices become less diagonal). When τ is small (dense matrices), r can be 3–15×, meaning Adam's preconditioner makes the condition number 3 to 15 times worse than no preconditioning at all.
Adam's preconditioner is only clearly beneficial when τ ≈ 1—i.e., when the Hessian is already nearly diagonal.

This explains the results in Figure 4: within each dense sub-block, Adam's coordinate-wise learning rates are not only unnecessary but counterproductive—a single well-chosen learning rate per block does better because it avoids the false precision of a diagonal preconditioner on a dense matrix.

Connecting to neural networks. The paper does not claim that the random dense matrix experiment directly represents neural network Hessians. Rather, it uses it to establish a mathematical possibility: if the Hessian has dense sub-blocks (which neural networks do, as shown in Figures 3 and 7), then within those sub-blocks, coordinate-wise learning rates may not help, and a block-wise approach could be superior. This reframes the research question from "can we compress v?" to "can we identify the natural blocks within which v's granularity is wasted?"

Evidence from Transformers (Section 2.1, Figure 6). The paper extends the block-wise insight to Transformers through a "leave-x-out" experiment. Using a 4-layer Transformer with PyTorch default parameter blocks, the authors randomly select one parameter block (e.g., query in layer 2, mlp.fc_1 in layer 3), replace Adam's coordinate-wise learning rates for that block with a single learning rate (grid-searched), keep Adam's standard treatment for all other blocks, and measure training loss.

The results (Figures 6a–6c) show that:

Leaving one block out (Figure 6a): For all five randomly selected blocks, the single-learning-rate version matches or beats Adam.
Leaving two blocks out (Figure 6b): Still matches or beats Adam for all five trials.
Leaving three blocks out (Figure 6c): Again matches or beats Adam for all five trials, though with more variance.

Figure 6d extends this to all possible blocks in the 4-layer Transformer, plotting the performance gap loss(Adam(leave-one-out)) - loss(Adam). A negative value means the leave-one-out method performs better. The paper observes:

"We find Adam (leave-one-out) always performs on par with Adam, and for most blocks, Adam (leave-one-out) performs better."

The authors do not extend this experiment beyond three blocks because "the cost of grid search grows exponentially"—searching k blocks' learning rates each from L candidates requires L^k trials. This motivates the need for a cheap, principled way to assign block-wise learning rates without grid search, which leads directly to Adam-mini.

The Hessian Structure of Transformers and the Partitioning Principle (Section 2.3)

The paper's central design principle is:

Principle 1: "We should partition parameters into blocks, such that each parameter block is associated with the smallest dense sub-block in Hessian."

This principle follows from the analysis in Section 2.1: if a Hessian sub-block is dense, coordinate-wise learning rates within that block are at best unnecessary and at worst harmful. But if the partition is too coarse—merging parameters that belong to different dense sub-blocks—then crucial inter-block learning rate distinctions are lost, causing training instability.

Why does coarseness fail? Appendix E.2 restates a finding from Zhang et al. (2024): different parameter blocks in Transformers have dramatically different Hessian eigenvalue distributions. The query projection matrix might have eigenvalues spanning one range, while the mlp weight matrix spans a completely different range. If these two blocks were merged into a single partition and given a single learning rate, the optimizer would be forced to use a compromise rate that is simultaneously too large for the well-conditioned block and too small for the ill-conditioned block—or vice versa. This is why the PyTorch default partition, which groups parameters by layer (layer_i.query, layer_i.key, etc.), fails on larger models (Figure 7i, Llama 2-1B shows training loss spikes).

Transformer Hessian exploration (Section 2.3, Figure 7). To determine the correct partition for Transformers, the authors compute the Hessian of a small 1-layer Transformer (n_emb=16, n_head=4, mlp.fc_1 width=32) after 1% of training steps and visualize the block structure (Figures 7a–7h). The Hessian is computed exactly using the Pearlmutter (1994) method—two passes of backpropagation to compute the Hessian-vector product, applied repeatedly to construct the full matrix.

The key findings, organized by parameter class:

Class 1: query and key (Figures 7a, 7b). The Hessians exhibit a clear near-block-diagonal structure where the number of diagonal blocks equals the number of attention heads (4 in this case). Each attention head corresponds to a separate dense Hessian sub-block. This makes structural sense: each head's query projection learns to attend to different positional relationships, and the cross-entropy loss factor p(x)(1-p(x)) naturally decouples the heads' second-order interactions (via the analysis in Appendix C). The correct partition is to split query and key by attention heads, treating each head's parameters as an independent block.

Class 2: attn.proj and MLPs (Figures 7d, 7e, 7f). The Hessians of attn.proj, mlp.fc_1, and mlp.proj exhibit block-diagonal structures where the number of blocks equals the number of output neurons (16 for attn.proj and mlp.proj, 32 for mlp.fc_1). This follows the same logic as the MLP analysis in Collobert (2004): each output neuron corresponds to an independent dense Hessian sub-block. The correct partition is to split these by output neurons—for a weight matrix W ∈ R^{d_out × d_in}, each row (associated with one output neuron) becomes its own block.

Class 3: value (Figure 7c). The Hessian of value is structurally ambiguous:

"For value, the structure of Hessian seems less clear. It seems to have the hint of 16 diagonal blocks (16 is the number of output neurons), but the pattern is less obvious."

The paper identifies several distinguishing properties of value: (1) its Hessian entries are about 10^6 larger than those of query and key, (2) the block structure is less sharply defined because value is positioned outside the softmax operator (unlike query and key, which pass through the softmax's nonlinearity), and (3) this structural difference is significant enough to warrant different treatment. The paper proposes two partition strategies for value—partitioning by output neurons (Strategy I, which "works well when the number of total training steps is large") and treating value as a whole (Strategy II, which "works better when the number of total training steps is small"). By default, the paper uses Strategy I, but provides optimizer.wv_names = {} as a one-line override for switching to Strategy II.

Class 4: embed and output (Figures 7g, 7h). The Hessians of the embedding layer and output projection layer have a near-block-diagonal structure where the number of blocks equals the number of tokens (8 in this case). This is distinct from the neuron-based or head-based structures elsewhere. The correct partition is to split embed and output by tokens—each token's embedding row becomes its own block.

Why the default PyTorch partition fails. PyTorch's default parameter grouping treats each named parameter tensor (e.g., transformer.layers.0.attention.query.weight) as one block. For a Transformer, this means all attention heads in the query matrix are grouped together, all output neurons in attn.proj are grouped together, etc. This is coarser than the smallest dense Hessian sub-block (since, e.g., query contains n_head dense blocks merged into one), which violates Principle 1. The consequence is demonstrated in Figure 7i: Adam-mini with PyTorch default partition suffers clear loss spikes when training Llama 2-1B, while the Hessian-aware partition stabilizes training and matches AdamW.

The Collobert (2004) mechanism and why it matters. Appendix C restates the theoretical mechanism behind the near-block-diagonal Hessian. For a 1-hidden-layer neural network with cross-entropy loss, the off-diagonal-block Hessian entries between weights of different output neurons contain the factor:

$\frac{\partial^2 \ell(f(\theta, x), y)}{\partial w_i \partial w_j} = p(x)(1 - p(x)) v_i v_j \phi'(w_i^T x) \phi'(w_j^T x) x x^T \quad \text{for } i \neq j$

where p(x) = 1/(1 + \exp(-y f(\theta, x))) is the model's predicted probability for the correct class, \phi'(\cdot) is the derivative of the activation function, v_i are the output weights, w_i are the input weights, x is the input, and y is the label.

What this equation means in practice: As training progresses and p(x) approaches 1 (the model becomes confident in its predictions), the factor p(x)(1-p(x)) rapidly shrinks toward zero. This factor multiplies the entire off-diagonal Hessian entry ∂²ℓ/(∂w_i ∂w_j) when i ≠ j. As a result, the cross-entropy Hessian naturally decouples into blocks corresponding to different output neurons, just from the mechanics of confidence-driven optimization. The authors note this can happen "after 1 training step" (Collobert, 2004) and reproduce it numerically for both MLPs (Figure 3b–d) and Transformers (Figure 7).

Why this mechanism matters for partitioning: The natural decoupling at output neuron boundaries means that a block-wise learning rate assignment that respects these boundaries is not an approximation—it is aligning the optimizer's granularity with the problem's actual second-order structure. Conversely, merging across neuron boundaries (as PyTorch default does) forces the optimizer to use a single learning rate for what are actually independent Hessian sub-blocks, causing tension and potential instability.

The Adam-mini Algorithm (Section 2.2, Algorithms 1–3)

Adam-mini is defined by a two-step procedure applied at each optimizer step: parameter partitioning (determined once at construction time based on Principle 1) and block-wise second-moment estimation (applied at every training step).

The general form (Algorithm 1).

Algorithm 1 presents the conceptual structure:

Input: weight-decay coefficient λ and current step t.
Partition: group parameters into param_blocks by Principle 1 (using Algorithm 3 for Transformers or Algorithm 3's simpler version for non-Transformers).
For each parameter block:
- Retrieve the gradient g = param.grad.
- Apply decoupled weight decay: param = param - η_t · λ · param (this exactly matches AdamW's decoupled weight decay, where η_t is the scheduled learning rate at step t and λ is the weight decay coefficient).
- Update first-order momentum: m = (1 - β_1) · g + β_1 · m (standard Adam exponential moving average).
- Compute bias-corrected momentum: \hat{m} = m / (1 - β_1^t).
- Compute block-wise second moment (the key difference from Adam): v = (1 - β_2) · mean(g ⊙ g) + β_2 · v, where mean(g ⊙ g) is the scalar average of the squared gradients across all parameters in the block. This produces one scalar v per block, not one per parameter. β_2 is the second-moment decay rate (standard Adam uses 0.95 or 0.999 depending on the task; the paper uses β_2 = 0.95 for all LLM experiments).
- Compute bias-corrected second moment: \hat{v} = v / (1 - β_2^t).
- Apply the parameter update: param = param - η_t · \hat{m} / (\sqrt{\hat{v}} + ε), where ε is the numerical stability constant (typically 1e-8).

What changes from Adam to Adam-mini: Only the second-moment accumulation (step 8 in Algorithm 1). Adam computes v = (1-β_2) · (g ⊙ g) + β_2 · v, which produces a tensor the same shape as g with element-wise squared gradient accumulation. Adam-mini replaces g ⊙ g (element-wise square) with mean(g ⊙ g) (scalar mean across all parameters in the block), producing a single scalar v per block. This single scalar is then broadcast back to all parameters in the block during the update step.

The concrete example from the paper (Section 2.2).

To make this concrete, the paper provides a 5-parameter example. For a model with 5 parameters w ∈ R^5, the effective learning rate vectors u (the vector that multiplies the momentum m in the update) differ as follows:

Adam: u_Adam = [η/√v_1, η/√v_2, η/√v_3, η/√v_4, η/√v_5] — five distinct learning rates, one per parameter.
Adam-mini with partition (1, 2, 3) and (4, 5): u_mini = [η/√((v_1+v_2+v_3)/3), η/√((v_1+v_2+v_3)/3), η/√((v_1+v_2+v_3)/3), η/√((v_4+v_5)/2), η/√((v_4+v_5)/2)] — only two distinct learning rates.

Where v_i = (1-β_2) · g_i^2 + β_2 · v_i in Adam (before averaging), and the averaging pools the v values within each block.

Why this averaging specifically (Appendix C). The paper provides three reasons for choosing the mean of v within each block rather than other candidates:

Grid-search would be too expensive. Optimal blockwise learning rates (as used in the leave-x-out experiments) require exponential search cost. The authors explicitly state that "such a searching procedure is not scalable."
The mean is the natural quantity to borrow from Adam. Among candidates tested (1-norm(v), 2-norm(v), max(v), min(v)), mean(v) performed best in ablation studies (Figure 15). The 1-norm and 2-norm diverged; max and min underperformed.
The mean keeps Adam-mini's trajectory close to Adam's. The authors argue from the backpropagation rule: for a weight matrix W ∈ R^{d×d}, the gradient decomposes as G = ∂ℓ/∂W = e·z^T, where e is the backpropagation error vector and z is the input feature. For the i-th row of G, all entries share the same error term e_i. Therefore, gradients within a row (which associates with one output neuron) are typically similar, and their mean is a good representative. The consequence is that "Adam-mini's trajectory closely resembles that of Adam" (Figure 9b, Figure 10), which makes it easier to reuse Adam's hyperparameters and maintain Adam's scaling behavior.

The partitioning algorithms (Algorithms 3 and 3).

The paper distinguishes two cases:

Algorithm 3 (Partition for non-Transformers): Simply uses the PyTorch default parameter names as blocks—each named parameter tensor (param_blocks[name] = param) becomes its own block. No special partitioning is applied. The authors note this is tested on CNNs, diffusion models, and graph neural networks, but caution that "in the future, it is possible that we will have more complicated non-Transformer architectures on which Algorithm 3 fails"—in which case Principle 1 would need to be applied manually.

Algorithm 3 (Partition for Transformers): Applies the Hessian-aware partitioning rules derived from Figure 7:

embed and output: Partition by tokens. For each token index i = 0...tokens-1, extract param[i] as a separate block. This treats each embedding row independently.
query and key: Partition by attention heads. For each head index i = 0...heads-1, extract param[i] (the corresponding slice of the weight tensor) as a separate block. This gives each attention head its own learning rate.
value, attn.proj, and mlp: Partition by output neurons. For each output neuron index i = 0...output_neurons-1, extract param[i] as a separate block.
Everything else: Keep as is (whole-tensor blocks for any parameter that does not match the above patterns, ensuring no parameters are left unassigned).

Why value gets partitioned by output neurons despite its ambiguous Hessian. Appendix D.6 provides guidance based on training length: partitioning by output neurons (Strategy I) works better when total training steps are large—it provides finer granularity for long runs where the value matrix has time to develop clear per-neuron structure. Treating value as a whole (Strategy II) works better for short runs, likely because the Hessian's per-neuron structure has not yet fully crystallized and a single learning rate avoids premature specialization. The paper's default recommendation is Strategy I.

Design Choices and Their Justifications

Choice 1: Keep m at full precision, only reduce v. The paper makes a deliberate asymmetry: the first-order momentum m is stored at full parameter resolution (unchanged from Adam), while only the second-order momentum v is compressed to block-wise scalars. This is justified by the theoretical framing: coordinate-wise v provides a diagonal preconditioner that is unnecessary in dense Hessian blocks, while m provides a per-parameter running average of gradient direction that remains valuable even within blocks—the gradient direction can vary across parameters in the same block even if the optimal learning rate is the same. No experiment tests reducing m's precision because the theoretical motivation does not suggest it.

Choice 2: Decoupled weight decay (AdamW-style), not Adam-style L2 regularization. Algorithm 1 applies weight decay as param = param - η_t · λ · param before the momentum-based update, exactly matching AdamW (Loshchilov & Hutter, 2017) rather than the original Adam's L2 regularization that couples weight decay with the adaptive learning rate. This is consistent with "Adam-mini performs well using the same hyperparameters as AdamW", since all LLM training baselines use AdamW.

Choice 3: Identical hyperparameters to AdamW. The paper emphasizes repeatedly that Adam-mini uses the same β_1, β_2, ε, learning rate schedule, and weight decay coefficient as the corresponding AdamW baseline. No new hyperparameters are introduced. This is a design choice grounded in the trajectory preservation property: since Adam-mini's mean(v) stays close to Adam's per-parameter v on average, the same hyperparameters that work for AdamW generally transfer. The sensitivity analysis in Figure 12c confirms this—Adam-mini's performance is "not overly sensitive to hyperparameters" on GPT-2-125M, with a range of learning rates from 2e-4 to 8e-4 producing similar validation losses.

Choice 4: β_2 = 0.95 for LLM training. While Adam's default β_2 = 0.999 is standard for many tasks, the paper uses β_2 = 0.95 for all LLM experiments (GPT-2, Llama). This is not specific to Adam-mini—it mirrors the standard LLM training practice where shorter second-moment memory is beneficial for the rapidly changing loss landscape. The authors follow the same β_2 as their AdamW baselines.

Choice 5: The partition is static, not adaptive. The parameter blocks are determined once at optimizer construction time and never change during training. This is a practical choice: dynamic re-partitioning would require ongoing Hessian computation during training (prohibitively expensive) and would introduce additional noise into the learning rate schedule. The static partition works because the near-block-diagonal Hessian structure is a persistent feature of neural network training (Figures 3b–d show it is present at initialization and maintained through training), not a transient one.

Choice 6: mean(g ⊙ g) rather than a per-parameter g ⊙ g. In Algorithm 1 step 8, the squared gradient is averaged before the exponential moving average is applied: v = (1-β_2) · mean(g ⊙ g) + β_2 · v. An alternative would be to compute v element-wise first (as in Adam) and then average the resulting per-parameter v values. The paper chooses to average the raw squared gradient per step because it maintains a consistent interpretation: v is an EMA of the block's average squared gradient, which directly corresponds to the variance of the gradient within that block. Computing element-wise v then averaging would (due to the EMA's different history per parameter) produce a less interpretable quantity.

Memory, Throughput, and Trajectory Characteristics (Section 2.4)

Memory savings: 50% of Adam's total optimizer memory. Adam stores two tensors per parameter (m and v), each the size of the model. Adam-mini stores one full-sized tensor (m) and one scalar per block (v_b). Since the number of blocks is tiny compared to the number of parameters (for a Transformer with n_layers, n_heads, and d_model, the number of blocks scales as n_layers · n_heads for attention projections plus n_layers · d_model for MLP/attention-output projections, compared to n_layers · d_model² total parameters), the memory for v is reduced by ≥99.9%. Table 1 quantifies this:

GPT-2-1.5B: AdamW requires 12.48 GB; Adam-mini requires 6.24 GB (50% reduction).
Llama 2-1B: AdamW requires 8.80 GB; Adam-mini requires 4.40 GB (50% reduction).
Llama 2-7B: AdamW requires 53.92 GB; Adam-mini requires 26.96 GB (50% reduction).
Llama 3-8B: AdamW requires 64.24 GB; Adam-mini requires 32.12 GB (50% reduction).
Llama 2-13B: AdamW requires 104.16 GB; Adam-mini requires 52.08 GB (50% reduction).

All calculations assume float32 storage for optimizer states, which is standard practice.

Why exactly 50%? Adam stores m and v, each equal to the model size. Adam-mini stores m at full size and v at negligible size. The total optimizer memory goes from 2 × model_size to ~1 × model_size, hence the 50% reduction. The gradients add one more model's worth of memory, making the total memory savings roughly 33% (from 3 × model_size to 2 × model_size), but the paper's claim of "50% of Adam's memory" refers to the optimizer states specifically, excluding gradients.

Throughput gains: up to 49.6% higher, translating to 33% less wall-clock time. Table 2 reports results for Llama 2-7B pre-training on 2× A800-80GB GPUs:

Adam-mini with batch size 4 per GPU, total batch size 256: 5572.19 tokens/second.
AdamW with batch size 2 per GPU: out of memory (cannot run).
AdamW with batch size 1 per GPU, total batch size 256: 3725.59 tokens/second.

Adam-mini's throughput is 49.6% higher than the maximum achievable AdamW configuration on the same hardware. This translates to processing 1B tokens in 49.85 GPU-hours for Adam-mini vs. 74.56 GPU-hours for AdamW—a 33.1% reduction in wall-clock time.

Why the throughput improves. The paper identifies two factors (Section 2.4):

No extra computation in per-step updates. The mean(g ⊙ g) operation is negligible compared to the element-wise g ⊙ g in Adam, and it "significantly reduces the number of vector-square-root and vector-division operations" since sqrt and division are performed on scalars rather than on tensors with millions of elements.
Reduced communication overhead. The memory cut-down allows larger batch sizes per GPU (Adam-mini can fit batch size 4 where AdamW can only fit batch size 1 on the same hardware). This means fewer gradient accumulation steps and correspondingly fewer all-reduce operations between GPUs, which is "known to be a major overhead" (Rajbhandari et al., 2021).

Trajectory similarity to AdamW. The paper emphasizes that Adam-mini produces training trajectories that closely resemble AdamW's. Figure 9b quantifies this: for a small Transformer, the Euclidean distance between checkpoints trained with Adam-mini and AdamW is much smaller than the distance between AdamW and other memory-efficient optimizers (Adafactor, CAME). This is attributed to the small modification: Adam-mini replaces only the learning rate assignment within blocks, preserving the overall momentum dynamics that govern trajectory evolution. The practical benefit is that "Adam-mini can maintain the scaling laws of LLMs trained by Adam" (Appendix C)—as evidenced by Figure 11 where Adam-mini's scaling law curves closely parallel AdamW's across model sizes from 39M to 1B.

Summary of the Key Technical Relationships

To make the dependencies explicit:

Collobert (2004) provides the structural insight → the Hessian of neural networks is near-block-diagonal with dense sub-blocks corresponding to output neurons or attention heads. This is a persistent property, not a transient one.
The dense sub-block property explains why Adam's v is redundant → within each dense block, a diagonal preconditioner cannot reliably reduce the condition number and may worsen it (Figure 5). Coordinate-wise learning rates offer no benefit over a single well-chosen rate per block (Figures 4, 6).
This redundancy is exploitable, but only if the blocks are correctly identified → too-coarse partitions (like PyTorch default) merge distinct Hessian sub-blocks, destroying the inter-block learning rate differentiation that is necessary for Transformers (Zhang et al., 2024). The training instability in Figure 7i demonstrates this directly.
Adam-mini's design is the minimal modification to Adam that exploits this redundancy → it preserves the momentum update, weight decay, and bias correction exactly; it changes only the granularity of v to match Hessian structure; and it chooses mean(v) as the per-block learning rate because it is the cheapest quantity that keeps the trajectory close to Adam's while avoiding grid search.

The result is an optimizer that is simultaneously: (a) memory-efficient (50% savings), (b) performance-preserving (matches AdamW across scales and tasks), (c) throughput-improving (49.6% faster on Llama 2-7B), and (d) hyperparameter-compatible (same settings as AdamW). The key intellectual contribution is that this combination is possible because the reduction is structure-aware, not generic—previous methods achieved (a) at the cost of (b); the Hessian-based partitioning is what breaks this trade-off.

4. Key Insights and Innovations

Innovation 1: A Structural Explanation for Why Adam's Memory Is Wasteful — and When It Isn't

The paper's most distinctive conceptual contribution is not the optimizer itself but the diagnostic framework that identifies where in a neural network Adam's coordinate-wise learning rates are redundant. Prior work on memory-efficient optimization either treated Adam's v as a monolithic structure to be compressed (Adafactor's low-rank factorization, SM3's minimal-value selection) or attempted to remove it entirely (sign-free methods). These approaches share a common assumption: that memory savings must come at the cost of some approximation error, and that the optimizer designer's job is to find the best compression trade-off.

This paper reframes the problem in a fundamentally different way. Rather than asking "how can we compress v with minimal information loss?", it asks "does v actually carry the information we think it does?" The answer, grounded in Hessian structure, is no — at least not within the dense sub-blocks that make up the vast majority of neural network parameters. The key finding from Figure 5 is that within dense Hessian regions, Adam's diagonal preconditioner can make the condition number worse than no preconditioning at all (r > 1 for small τ). This means coordinate-wise learning rates are not just unnecessary in these regions — they can be actively counterproductive.

What distinguishes this from prior diagnostic work on Adam (Zhang et al., 2020; Kunstner et al., 2023; Zhang et al., 2024) is that those papers focused on explaining why Adam works — the importance of sign descent, the role of block heterogeneity, the convergence properties. This paper identifies a regime where Adam's core mechanism becomes dysfunctional, and shows that this regime (dense Hessian sub-blocks) is precisely the regime that neural networks naturally create through training. The Collobert (2004) mechanism — p(x)(1-p(x)) shrinking off-diagonal-block entries toward zero — produces Hessians that are simultaneously dense within blocks and sparse between them. Adam enters a worst-of-both-worlds situation: it wastes memory on fine-grained learning rates within blocks where they don't help, while the true inter-block variation (which does require distinct learning rates) is captured regardless.

This is a fundamental rather than incremental shift in how the optimization community should think about Adam's v. Before this paper, the default assumption was "v is essential in its full precision; any simplification hurts." After this paper, the correct framing is "v's precision is wasted on intra-block variation; what matters is inter-block variation, which can be captured with orders of magnitude fewer learning rates if the block boundaries are correctly identified." This reframing explains why prior compression methods underperformed (they either compressed v uniformly without respecting block structure, or they used partitions that were too coarse and merged distinct Hessian blocks) and why Adam-mini succeeds (it compresses precisely where the Hessian permits and retains differentiation precisely where it's needed).

Innovation 2: Architecture-Aware Optimizer Design as a Principle, Not a Heuristic

Prior to this paper, parameter grouping in optimizers was either practical (Adafactor groups by tensor shape for factorization convenience; SM3 groups by predetermined covers) or purpose-agnostic (PyTorch's default grouping by named parameter tensors). The paper introduces Principle 1 — partition parameters by the smallest dense Hessian sub-blocks — as an explicit design principle that connects optimization theory to neural architecture structure.

This is a conceptual advance, not just a useful rule of thumb. Principle 1 unifies three previously separate observations into a single framework:

Collobert (2004): That the cross-entropy Hessian is near-block-diagonal with blocks corresponding to output neurons (a structural fact about neural networks).
Zhang et al. (2024): That different parameter blocks in Transformers need different learning rates due to eigenvalue heterogeneity (a requirement for effective optimization).
The paper's own numerical finding (Figures 4-6): That within each dense block, a single learning rate suffices and coordinate-wise granularity provides no benefit (a limit on how fine-grained the differentiation needs to be).

The principle says: use exactly one learning rate per dense Hessian sub-block — no more, no fewer. More would be wasteful (the random quadratic and leave-x-out experiments); fewer would cause instability (Figure 7i, where PyTorch default partition merges distinct Hessian blocks and produces loss spikes on Llama 2-1B).

What makes this a genuine principle rather than an ad-hoc rule is its explanatory and predictive power. It explains why Adafactor underperforms (low-rank factorization doesn't respect Hessian block boundaries), why naive layer-wise partitioning fails (it's too coarse for Transformers), and why the specific Transformer partitioning in Algorithm 3 works (it aligns with the Hessian's natural block decomposition). It also makes a falsifiable prediction: if a new architecture were to exhibit a different Hessian block structure — say, blocks at the level of attention head sub-spaces rather than whole heads — then an optimizer following Principle 1 would need to adapt its partitioning accordingly. The paper's cautious labeling of Algorithm 3 as "Partition for Transformers" (not "Partition for Neural Networks") reflects this: the principle is general, but its realization is architecture-specific.

This is a fundamental shift from how the field has approached optimizer design. The dominant paradigm has been to develop generic algorithms (Adam, Adafactor, LAMB) and then tune them per-task. Adam-mini inverts this: start with the structure of the problem, and design the optimizer's granularity to match. This aligns optimization with architecture in a way that prior work — even architecture-aware methods like LAMB, which adds layer-wise scaling but retains full coordinate-wise v — did not.

Innovation 3: Empirical Resolution of a Contradiction in the Optimization Landscape

The paper resolves a latent tension in the adaptive optimization literature through a result that is straightforward to state but had not been demonstrated at scale: a memory-efficient optimizer can match AdamW's performance without introducing new hyperparameters or requiring per-task tuning. This sounds like a performance claim, but its significance is conceptual: it demonstrates that the memory cost of Adam is not an inherent price of its optimization quality.

The field had long operated under a de facto trade-off assumption. AdamW delivers reliable, state-of-the-art performance across LLM training tasks, but at high memory cost (2× the model size for optimizer states). Adafactor — the most prominent memory-efficient alternative — saves comparable memory to Adam-mini (~48% according to Appendix A) but has a documented history of being difficult to tune and often underperforming. The paper's exhaustive hyperparameter sweep for Adafactor (Figures 13 and 19, spanning learning rates from 1e-5 to 1e-2, three β_2 values, four ε settings, and six warm-up ratios) shows that no configuration reliably closes the gap to AdamW on Llama 2-20M or Llama 2-1B, and the "Adafactor-Zhai-version" that performed better at 20M becomes unstable at 1B.

This pattern had created a de facto consensus in parts of the LLM training community: if you need reliable performance, use AdamW and pay the memory cost; if you're memory-constrained, use Adafactor and accept degraded performance or invest in hyperparameter tuning. The scaling law experiments (Figure 11) demonstrate that this trade-off is not inherent — Adam-mini's scaling curves essentially overlay AdamW's across 39M to 1B parameters, using the same hyperparameters and 50% less memory. The fitted lines in Figure 11b further suggest this holds for larger models (subject to the scaling law's extrapolation).

This is a fundamental rather than incremental empirical finding because it breaks what appeared to be a hard trade-off. The reason the trade-off was breakable — rather than an inescapable consequence of information theory — traces back to the Hessian structure insight (Innovation 1): prior methods were losing performance not because memory savings inherently degrade optimization, but because they were applying the wrong type of compression to the wrong structure. Adam-mini demonstrates that Hessian-aware compression achieves both goals simultaneously. The implication for practitioners is a concrete regime change: "use the same hyperparameters as AdamW" means adoption does not require a separate tuning budget, removing the primary barrier that kept many practitioners from using memory-efficient alternatives.

Innovation 4: Verifying That Hessian Structure Persists and Is Exploitable Across Scales and Tasks

The paper makes a distinctive empirical contribution that goes beyond proposing a new optimizer: it provides the first systematic evidence that the near-block-diagonal Hessian structure of neural networks is stable enough across architectures, scales, and training phases to serve as a reliable basis for optimizer design. This is not a given — the Collobert (2004) analysis was on small MLPs; the paper's own Hessian visualizations (Figure 7) are on a 1-layer Transformer with tiny dimensions (n_emb=16, n_head=4). Whether this structure persists at scale (Llama 2-7B, 13B; GPT-2-1.5B) and across training phases (pre-training, SFT, RLHF) — and whether partitioning strategies derived from small-scale Hessian analysis transfer — was an open empirical question.

The evidence comes from multiple angles. The scaling law experiments (Figure 11) show that the Hessian-aware partition yields stable training across 39M to 1B parameters, with no loss spikes or degradation relative to AdamW — indirect evidence that the structure is robust enough to generalize. The trajectory comparison (Figure 9b) shows that Adam-mini stays close to AdamW's optimization path, suggesting the learning rate assignment captures the essential second-order information even as the loss landscape evolves. The across-task validation — pre-training (GPT-2, Llama), SFT, RLHF, diffusion models, vision Transformers, graph neural networks — demonstrates that the principle works beyond the specific setting where the Hessian was analyzed.

But what makes this more than a routine ablation study is the negative result embedded within it: the PyTorch default partition fails specifically on larger models (Figure 7i, Llama 2-1B), and the failure mode (loss spikes) is consistent with violating Principle 1 — merging distinct Hessian blocks forces the optimizer to use a single learning rate where multiple are needed. This failure is informative because it demonstrates that the Hessian structure is not just an interesting pattern that happens to correlate with good partitions — it is causally important. Violating it causes degradation; respecting it prevents degradation. This transforms the Hessian structure from a descriptive observation into a prescriptive design tool, and the paper's extensive across-scale validation establishes the reliability of that tool.

The significance for future work is substantial: any new architecture (mixture-of-experts, state-space models, hybrid architectures) can in principle be analyzed for its Hessian block structure using the same methodology (small-scale Hessian visualization + leave-x-out experiments), and a custom Adam-mini partition can be derived following Principle 1 — without requiring a full hyperparameter sweep at the target scale. This makes the paper's contribution not just a specific optimizer, but a methodology for designing architecture-aware optimizers that extends beyond Transformers.

5. Experimental Analysis

Evaluation Methodology

Dataset. All main experiments use the C4 dataset (Raffel et al., 2020) for Llama series pre-training and OpenWebText (Gokaslan et al., 2019) for GPT-2 series pre-training. Supervised fine-tuning and RLHF experiments use the UltraFeedback dataset (Cui et al., 2023), with 40% of the data used for SFT and 60% for reward model training / RL optimization. Non-LLM experiments use ImageNet (Deng et al., 2009) for vision tasks, CelebA for diffusion models, and OGBN-arxiv for graph neural networks.
Base model(s). The paper spans two model families: GPT-2 (Radford et al., 2019) at 125M, 330M, and 1.5B parameters, pre-trained on OpenWebText; and Llama 2 (Touvron et al., 2023) at scales from 20M to 13B parameters, pre-trained on C4. Additionally, Llama 3-8B is used in one pre-training comparison (Figure 10b). The scaling law experiments (Section 3.2) use Llama 2 architectures at sizes 39M, 67M, 102M, 162M, 271M, and 1B, all trained on approximately 20 × n_param tokens following Chinchilla's law (Hoffmann et al., 2022). For SFT and RLHF, the pre-trained Llama 2-7B model from Meta is used. The authors argue PaLM 2-S* is "representative of the capabilities of many contemporary LLMs" (Section 4, though the actual experiments use GPT-2 and Llama rather than PaLM).
Metrics. For pre-training, the primary metric is validation loss (cross-entropy) or validation perplexity (exp(val_loss)), reported at the end of training or as curves over tokens/iterations. For SFT, evaluation perplexity is used. For RLHF, evaluation reward from the trained reward model is reported, supplemented by MT-Bench scores (Zheng et al., 2024) where GPT-4 assigns a 0–10 score for multi-turn chatting quality. For vision tasks, validation accuracy or training loss is used; for diffusion models, FID (Fréchet Inception Distance, lower is better) and Inception Score (higher is better) evaluate image quality. Throughput is measured in tokens per second; wall-clock time is reported in GPU-hours.
Baselines. The primary baseline is AdamW (Loshchilov & Hutter, 2017). Memory-efficient alternatives include: Adafactor (Shazeer & Stern, 2018), both the original version and a modified version from Zhai et al. (2022) called "Adafactor-Zhai-version"; CAME (Luo et al., 2023); SM3 (Anil et al., 2019); and LAMB (You et al., 2019). All memory-efficient baselines incorporate momentum with β_1 = 0.9 to ensure fair comparison. For the Lion comparison (Appendix D.8), Lion (Chen et al., 2024b) is tuned following strategies from a concurrent work (authors, 2024). Adam-mini itself is compared in two partition variants: the Hessian-aware partition (Algorithm 3) and the PyTorch default partition (each named parameter tensor as one block), with the latter shown to fail on larger models.
Generation budget / compute accounting. Compute is measured primarily by wall-clock time and GPU-hours for throughput comparisons, and by total FLOPs for the scaling law experiments (Figure 11a). For fair optimizer comparison, all methods use the same learning rate tuning budget — the paper states it "tune[s] the learning rate for all methods, using the same tuning budget for each, and report[s] the best performance" (Section 3.1). Throughput tests are conducted on 2× A800-80GB GPUs without CPU offloading (Table 2). Pre-training token counts follow Chinchilla's law: models are trained on approximately 20× their parameter count in tokens (Table 9 details exact allocations: 39M model → 1.02B tokens, up to 1B model → 26.21B tokens).
Cross-validation / statistical protocol. No formal cross-validation is employed for the main pre-training experiments — each configuration is run once with a fixed seed, and validation loss curves are reported. For SFT and RLHF, learning rates are selected by grid search over validation loss or validation reward respectively. The scaling law experiments use systematic sweeps across model sizes but do not report confidence intervals or multiple seeds. The trajectory comparison (Figure 9b) uses the same random seed across optimizers to enable direct weight-space distance comparison.

Main Quantitative Results

Pre-Training: GPT-2 Series (125M–1.5B) on OpenWebText

Headline result: Adam-mini matches AdamW's validation loss curves across all three GPT-2 model sizes (125M, 330M, 1.5B) using 50% less memory, while Adafactor, CAME, SM3, and LAMB all underperform. Adam-mini with PyTorch default partition diverges and is stopped early due to instability.

GPT-2-125M (Figure 8a). At the end of 20B tokens of training:

Adam-mini achieves a validation loss indistinguishable from AdamW, with the two curves essentially overlaid throughout training.
Adafactor, SM3, and LAMB all converge to higher validation losses — the gap is visible from early in training and persists.
Adam-mini (PyTorch default partition) shows "clear unstable behavior" and the trial is stopped — the curve is not shown because it diverges.

GPT-2-330M and 1.5B (Figure 8b). The paper plots both model sizes on the same figure with dashed (AdamW) and solid (Adam-mini) lines of matching colors:

For GPT-2-330M: Adam-mini's validation loss closely tracks AdamW, with both converging to approximately 2.9–3.0 by 20B tokens.
For GPT-2-1.5B: Adam-mini again tracks AdamW, with both reaching approximately 2.55–2.65 validation loss by 20B tokens.
The paper notes an unexpected phenomenon on GPT-2-330M (Appendix D.9, Figure 21): at the recommended learning rate 3e-4, AdamW encounters loss spikes, while Adam-mini does not. The AdamW spike can be mitigated by changing ε from 1e-8 to 1e-6. The authors state: "we have not fully understood the cause of the loss spike and how Adam-mini prevents it in this experiment."

Training loss trajectories (Figure 9a). The training loss curves of Adam-mini "closely resemble those of AdamW" across all three GPT-2 sizes, with the curves for each model size essentially overlapping. This is consistent across the full 20B-token training run.

Trajectory comparison (Figure 9b). On a small 8-layer Transformer (11M parameters) trained on OpenWebText, the Euclidean distance between checkpoints trained with Adam-mini and AdamW (under the same random seed and learning rate 1e-5) is measured every 250 iterations. The distance for Adam-mini vs. AdamW stays under approximately 20 (in parameter-space L2 norm), while Adafactor vs. AdamW and CAME vs. AdamW diverge to approximately 250 and 100, respectively — more than an order of magnitude larger.

Pre-Training: Llama Series (20M–13B) on C4

Headline result: Adam-mini performs on par with AdamW across Llama 2-1B, Llama 3-8B, and Llama 2-13B, with training loss curves that closely match AdamW. Adafactor and CAME consistently underperform. On Llama 2-7B (Figure 1), Adam-mini achieves 49.6% higher throughput and 33% less wall-clock time.

Llama 2-1B (Figure 10a). Training for 10,000 iterations:

Adam-mini achieves a training loss curve nearly identical to AdamW, both converging from approximately 8 to roughly 3.2–3.5 by iteration 10,000.
CAME underperforms both, converging to a higher training loss of approximately 4.5.
Adafactor (standard) converges to roughly 6.0 — substantially worse.
The "Adafactor-Zhai-version" performs better than standard Adafactor but still underperforms Adam-mini, reaching approximately 4.0–4.5.

Llama 3-8B and Llama 2-13B (Figure 10b). Despite the smaller batch sizes forced by hardware constraints (seq_len = 2048, batch_size = 8):

Adam-mini and AdamW produce nearly overlapping training loss curves for both model sizes.
For Llama 3-8B: both converge from approximately 7 to roughly 3.5–4.0.
For Llama 2-13B: both converge from approximately 7 to roughly 3.2–3.5.

Llama 2-7B (Figure 1). On 2× A800-80GB GPUs:

Memory (Figure 1a): Adam-mini uses approximately 0.75× the memory of AdamW (consistent with the 50% optimizer-state reduction, though the total memory including activations/gradients makes the ratio slightly less than 0.5× overall).
Throughput (Figure 1a, Table 2): Adam-mini reaches 5572.19 tokens/second vs. AdamW's 3725.59 tokens/second — a 49.6% increase. AdamW with batch size 2 per GPU runs out of memory, so the comparison is Adam-mini (batch size 4/GPU) vs. AdamW (batch size 1/GPU), both at total batch size 256.
Validation loss (Figure 1b): Over 2.25B tokens, Adam-mini tracks AdamW's validation loss, both decreasing from roughly 6.0 to roughly 3.0.
Wall-clock time (Figure 1c, Table 2): Adam-mini processes the same number of tokens in 33.1% less time. At 1B tokens: 49.85 GPU-hours (Adam-mini) vs. 74.56 GPU-hours (AdamW). Projected to Chinchilla-optimal 70B tokens: 3,489.55 GPU-hours vs. 5,219.16 GPU-hours. At 140B tokens: 6,979.10 vs. 10,438.32 GPU-hours.

Llama 2-20M (complete run, Figure 16). Trained on 26B tokens following Chinchilla's law:

The complete validation loss curves for Adam-mini and AdamW are "closely resemble[d]" throughout all 200,000 iterations, with both converging to roughly 2.5–2.6 by the end.

Scaling Law Experiments: Llama 2 (39M–1B) by Chinchilla's Law

Headline result: Adam-mini's validation loss curves consistently parallel AdamW's across model sizes from 39M to 1B, with Adam-mini achieving slightly lower final validation perplexity for all model sizes tested. The fitted power-law lines suggest Adam-mini's scaling behavior extends to larger models.

Loss vs. FLOPs (Figure 11a). Both AdamW (solid lines) and Adam-mini (dashed lines) produce nearly identical scaling trajectories:

The curves for each model size (39M through 1B) are essentially overlaid between the two optimizers.
At the 1B scale (the rightmost point, 10^20 FLOPs), both converge to approximately 3 × 10^0 validation loss.
At the 39M scale (leftmost point, ~10^17 FLOPs), both converge to approximately 4 × 10^0 validation loss.

Loss vs. parameters (Figure 11b). The final validation perplexity after Chinchilla-optimal training:

Adam-mini achieves slightly lower perplexity than AdamW for every model size tested (Table 4):
- 39M: AdamW 40.795 vs. Adam-mini 40.407
- 67M: 29.319 vs. 29.014
- 102M: 24.670 vs. 24.192
- 162M: 20.360 vs. 20.172
- 271M: 17.178 vs. 17.035
- 1B: 12.452 vs. 12.372

The fitted power-law lines for both optimizers are nearly parallel, with Adam-mini's line shifted slightly lower (better). The paper states: "the fitted lines in Figure 11 (b) suggest that Adam-mini can be scaled up to larger models (if the scaling law holds)."

Supervised Fine-Tuning and RLHF (Llama 2-7B)

Headline result: Adam-mini achieves better SFT evaluation perplexity (lower) and higher RLHF evaluation reward than AdamW, while using 50% less memory. On MT-Bench with GPT-4 evaluation, Adam-mini scores higher than AdamW across SFT (both full and LoRA) and RLHF.

SFT — full parameter tuning (Figure 12a). Over 60,000 iterations:

Adam-mini consistently achieves lower evaluation perplexity than AdamW, with the gap appearing early and persisting.
At the end of training: Adam-mini reaches approximately 2.36 evaluation perplexity vs. AdamW's approximately 2.42.

SFT — LoRA (Figure 22, Appendix D.10). When LoRA (Hu et al., 2021) is used with rank 128:

Replacing the Adam steps in LoRA with Adam-mini yields lower evaluation perplexity: approximately 2.47 vs. AdamW's approximately 2.53 by iteration 60,000.
Both methods use learning rate 2e-5 found by grid search from {1e-6, 2e-6, 3e-6, 4e-6, 5e-6, 1e-5, 2e-5}.

RLHF with ReMax (Figure 12b). Over 600 iterations of reward optimization:

Adam-mini achieves higher and more stable evaluation reward than AdamW.
By iteration 600, Adam-mini reaches approximately 3.3–3.5 reward vs. AdamW's approximately 2.8–3.0.
Adam-mini's peak learning rate is 5e-7; AdamW's is 1e-6 (both selected from {5e-7, 1e-6, 2e-6} based on validation reward).

MT-Bench GPT-4 evaluation (Table 5, Appendix D.4). Averaged scores (0–10 scale, higher is better):

SFT (LoRA): AdamW 4.23 vs. Adam-mini 4.41
SFT (full): AdamW 5.37 vs. Adam-mini 5.40
RLHF: AdamW 5.54 vs. Adam-mini 5.68

Adam-mini outperforms AdamW in all three alignment settings, with the largest gap in the RLHF setting (0.14 points, or roughly 2.5% relative improvement). A qualitative example of model responses is shown in Figure 23 (Appendix D.11), where Adam-mini's model produces responses comparable to AdamW's on multi-turn instruction-following tasks.

Non-LLM Tasks: Vision, Diffusion, and Graph Neural Networks

Headline result: Adam-mini performs on par or slightly better than AdamW across ResNet-18, Swin-Transformer, DiT-XL-2, DC-AE-Diffusion, DDPM, GCN, and GAT, with comparable training curves and, in some cases, better final metrics.

Detailed metrics (Tables 6 and 7). Selected comparisons at 100% training steps:

DDPM on CelebA (training loss): Adam-mini 0.0388 vs. AdamW 0.0394 (lower is better).
ResNet-18 on ImageNet (validation accuracy): Adam-mini 0.6667 vs. AdamW 0.6669 (essentially tied).
Swin-Transformer on ImageNet (validation accuracy): Adam-mini 0.7300 vs. AdamW 0.7310 (within margin).
DiT-XL-2 on ImageNet (training loss): Adam-mini 0.1430 vs. AdamW 0.1431; (FID, lower is better): Adam-mini 88.20 vs. AdamW 91.83; (Inception Score, higher is better): Adam-mini 13.90 vs. AdamW 12.38.
DC-AE-Diffusion on ImageNet (FID): Adam-mini 33.15 vs. AdamW 34.72; (Inception Score): Adam-mini 44.38 vs. AdamW 41.79.
GAT on OGBN-arxiv (validation accuracy): Adam-mini 0.7429 vs. AdamW 0.7421.
GCN on OGBN-arxiv (validation accuracy): Adam-mini 0.7423 vs. AdamW 0.7374.

Training curves (Figures 17, 18). Across Swin-Transformer, DiT-XL-2, DC-AE-Diffusion, ResNet-18, and DDPM, Adam-mini's loss/accuracy curves closely track AdamW's throughout training. In several cases (DiT-XL-2, DDPM), the curves are essentially indistinguishable. No training instability is observed for non-Transformer architectures, consistent with the simpler partition strategy (Algorithm 3, using PyTorch default parameter blocks) being sufficient.

Ablation Studies and Robustness Checks

Partition strategy: PyTorch default vs. Hessian-aware (Figures 7i, 8a): The PyTorch default partition causes training instability on both Llama 2-1B and GPT-2-125M, while the Hessian-aware partition (Algorithm 3) stabilizes training and matches AdamW. On Llama 2-1B (Figure 7i), the default partition produces a clear loss spike around iteration 250–500 that is absent with Hessian-aware partitioning. On GPT-2-125M (Figure 8a), the default partition trial is stopped early due to "clear unstable behavior." This demonstrates that the partitioning choice is not cosmetic — it is necessary for stability at scale.

Partitioning value by output neurons vs. treating as a whole (Appendix D.6): The paper identifies that value's Hessian structure is ambiguous (Figure 7c) and tests both strategies. Partitioning by output neurons (Strategy I) works better when total training steps are large; treating value as a whole (Strategy II) works better for short runs. This finding emerges from comparing Llama experiments with 10,000 steps (where Strategy II works better, used in Figure 10) vs. GPT-2 and scaling law experiments with ≥50,000 steps (where Strategy I works better, used in Figure 8 and Figure 11). The paper provides optimizer.wv_names = {} as a one-line override for Strategy II.

Choice of mean(v) vs. alternatives for block-wise learning rate (Figure 15, Appendix D.2): Among candidates for the block-wise learning rate — 1-norm(v), 2-norm(v), max(v), min(v), and mean(v) — only mean(v) achieves stable, competitive performance on Llama 2-20M pre-training. The 1-norm and 2-norm variants diverge (the blue curve is "out of range" and not shown), while max and min produce higher validation losses than mean (e.g., max(v) reaches approximately 4.5 validation loss vs. mean(v)'s ~4.0). This justifies the paper's claim that "average of v is the most natural quantity to 'borrow'" from Adam.

Sensitivity to hyperparameters (Figure 12c): On GPT-2-125M pre-training, Adam-mini is tested across peak learning rates from 2e-4 to 8e-4 (with β_1 = 0.9, β_2 = 0.95 fixed). The validation loss after 2.5B tokens (Chinchilla-optimal) varies only modestly — from approximately 3.05 at 2e-4 to roughly 3.10 at 8e-4, with a minimum around 6e-4. The paper concludes Adam-mini "seems not overly sensitive to hyperparameters," consistent with its trajectory similarity to AdamW.

Adafactor hyperparameter sweep (Figures 13, 19, Appendix D.7): An extensive search over Adafactor and Adafactor-Zhai-version is conducted on Llama 2-20M and Llama 2-1B:

On Llama 2-20M (Figure 13a): Sweeping learning rates from 1e-5 to 1e-2, both Adafactor versions consistently underperform Adam-mini. The best Adafactor-Zhai-version configuration reaches approximately 4.75 validation loss vs. Adam-mini's ~4.0.
Additional sweeps on Llama 2-20M (Figure 19): Changing β_2 from 0.999 to 0.95 (Figure 19a), varying warm-up steps from 1% to 10% of total (Figure 19b), and varying ε from 10^{-30} to 10^{-6} (Figure 19c) all fail to close the gap to Adam-mini.
On Llama 2-1B (Figure 13b): The Adafactor-Zhai-version now suffers from training instability while the original version performs better, but both still underperform Adam-mini substantially (validation loss ~5.5–8 for Adafactor variants vs. ~3.2 for Adam-mini).

Throughput comparison: Adam-mini vs. Adafactor (Figure 13c): On Llama 2-1B with 2× A800-80GB GPUs, Adam-mini achieves 40% higher throughput than Adafactor (exact numbers not reported in tokens/second, but shown visually in Figure 13c as a bar chart). The paper attributes this to Adam-mini's simpler computation (row-wise mean vs. row-and-column summation) and smaller v dimension.

Lion hyperparameter sweep (Figure 20, Appendix D.8): Following the optimal tuning strategies from a concurrent work (authors, 2024) — including learning rates 10× smaller than AdamW, the "magical number" lr = 3.16e-4, and (β_1, β_2) = (0.95, 0.98) — Lion underperforms Adam-mini on both Llama 2-20M and GPT-2-125M. On GPT-2-125M (Figure 20b), Lion encounters loss spikes for all learning rate candidates tested (5e-5 through 6e-4). The paper notes: "we have not managed to make Lion work, and we haven't been able to reproduce (authors, 2024) on Llama 2-20M and GPT-2-125M."

Combination with LoRA (Figure 22, Appendix D.10): Replacing the Adam optimizer in LoRA with Adam-mini (keeping all other LoRA settings identical: rank 128, learning rate 2e-5) yields lower evaluation perplexity on Llama 2-7B SFT — approximately 2.47 vs. 2.53 for AdamW. This demonstrates that Adam-mini's benefits are orthogonal to and combinable with parameter-efficient fine-tuning methods.

Trajectory similarity across optimizers (Figure 9b): The Euclidean distance between model checkpoints trained with different optimizers vs. AdamW is quantified every 250 iterations on an 11M-parameter Transformer. Adam-mini's distance stays below ~20 throughout; CAME diverges to ~100; Adafactor diverges to ~250. This provides quantitative evidence for the paper's claim that "Adam-mini generates similar trajectories to that of AdamW, while other methods cannot."

Critical Assessment

Do the Experiments Genuinely Support the Paper's Central Claims?

Claim: Adam-mini "performs on par or better than AdamW with 50% less memory footprint" (abstract, Section 1). The evidence for this claim is extensive and consistent across model scales (39M to 13B), model families (GPT-2, Llama 2, Llama 3), and training paradigms (pre-training, SFT, RLHF). The pre-training loss curves (Figures 8, 10, 11) are nearly indistinguishable between Adam-mini and AdamW, and the scaling law experiments (Figure 11b) show Adam-mini achieving slightly lower perplexity at every scale. The SFT and RLHF results (Figures 12a, 12b) actually show Adam-mini modestly outperforming AdamW. Non-LLM tasks (Tables 6, 7) show comparable performance within measurement noise.

However, this claim needs qualification that the paper itself provides: the 50% memory reduction refers to optimizer state memory only (from 2 × model_size to ~1 × model_size for m and v combined). Total GPU memory including activations, gradients, and model parameters sees a smaller reduction — roughly 25–33% depending on configuration. The paper's memory numbers in Table 1 are for optimizer states only and explicitly exclude gradients.

Claim: The memory reduction "also alleviates communication overheads among GPUs, thereby increasing throughput" (abstract). The 49.6% throughput improvement (Table 2) is measured in a specific configuration: Llama 2-7B on 2× A800-80GB GPUs, comparing Adam-mini with batch size 4/GPU against AdamW with batch size 1/GPU (AdamW with batch size 2/GPU runs out of memory). This is a real and practically meaningful gain, but it conflates two effects: (1) reduced communication because larger per-GPU batch sizes mean fewer gradient accumulation steps (pure memory benefit), and (2) reduced per-step computation because mean(g⊙g) and scalar sqrt are cheaper than element-wise operations. The paper does not disentangle these, and the throughput improvement would likely be smaller in configurations where AdamW already fits comfortably in memory (e.g., 8× GPUs where per-GPU batch sizes are already high). The 33% wall-clock time reduction is a derived number from the throughput improvement and is accurate under the stated configuration, but should not be interpreted as a universal speedup factor.

Claim: "≥99.9% of these learning rates in v could be harmlessly removed if we carefully partition the parameters into blocks following our proposed principle on Hessian structure" (abstract). This is the strongest form of the paper's core theoretical claim. The evidence supporting it is:

Strong evidence: The random quadratic experiments (Figures 4, 5) demonstrate the mathematical possibility — within dense sub-blocks, coordinate-wise learning rates can be counterproductive.
Moderate evidence: The leave-x-out experiments on a 4-layer Transformer (Figure 6) show that individual blocks can be simplified without performance loss, but only up to 3 blocks simultaneously due to the combinatorial cost of grid search. The "≥99.9%" number is derived from counting parameters in the Hessian-aware blocks vs. total parameters, not from an experiment that actually replaced 99.9% of v entries with grid-searched optimal per-block rates.
Practical evidence: The scaling law and pre-training experiments show that Adam-mini's mean(v) approach works at scale, but this is a specific choice of block-wise learning rate, not necessarily the optimal one. The paper acknowledges this gap: "there is great room to improve the design of Adam-mini: currently Adam-mini uses a simple and cost-effective way to design a learning rate for each dense Hessian sub-block, but it might not be an optimal way" (Section 4).

So "harmlessly removed" is demonstrated in the sense that performance is maintained, but "≥99.9%" is more of a structural calculation than an empirically validated compression ratio — we do not actually know whether every single one of those removed learning rates was individually harmless, only that the aggregate effect of removing them (via block-wise averaging) is benign.

Claim: "Adam-mini performs well using the same hyperparameters as AdamW" (summary of Section 3). This is the claim that distinguishes Adam-mini most sharply from Adafactor and other memory-efficient alternatives. The evidence supports it well: across all experiments, Adam-mini uses the same β_1, β_2, ε, weight decay, warm-up schedule, and learning rate as the corresponding AdamW baseline, with no new hyperparameters introduced. The sensitivity analysis (Figure 12c) confirms that the learning rate does not need special tuning. The one exception is the RLHF experiment, where Adam-mini selects 5e-7 vs. AdamW's 1e-6 from the same search grid — a minor difference.

The paper's negative results on Adafactor (Figures 13, 19) and Lion (Figure 20) provide the contrast that makes this claim significant: those methods require extensive tuning, and even with careful sweeps, the paper could not match AdamW's performance. The claim should be understood as "Adam-mini inherits AdamW's hyperparameter robustness" rather than "Adam-mini works with any hyperparameters AdamW would work with" — the latter cannot be tested without exhaustive sweeps, but the paper's consistent reuse of AdamW's settings across diverse tasks makes a strong case.

Genuine Weaknesses in Experimental Design

1. The Hessian analysis is at a vastly smaller scale than the main experiments. Figure 7's Hessian visualizations are on a 1-layer Transformer with n_emb=16, n_head=4, and mlp.fc_1 width 32 — fewer than 10,000 parameters. The main experiments go up to 13B parameters. The paper does not verify (nor could it feasibly verify, given the O(n²) cost of Hessian computation) that the block structure observed at tiny scale persists at large scale. The indirect evidence — that the partition derived from the small-scale analysis works at large scale — is suggestive but not definitive; the large-scale training stability could be due to factors other than Hessian structure (e.g., the partition coincidentally aligning with gradient variance patterns). Without even a moderate-scale Hessian analysis (e.g., at the 1M-parameter level), the causal link between Hessian structure and partitioning effectiveness remains somewhat speculative.

2. Single-seed, single-run reporting throughout. Nearly all experiments appear to be run once with a fixed seed. The paper does not report confidence intervals, error bars, or multiple-seed averages for any of its main results. Given that LLM training is known to be sensitive to random seed (particularly for smaller models and shorter runs), the conclusion that Adam-mini "matches or outperforms AdamW" could be partially confounded with seed variance. The trajectory comparison (Figure 9b) uses a single seed; it is unclear whether the close Adam-mini-to-AdamW distance would replicate. This is a significant omission, particularly for the scaling law experiments where the differences between Adam-mini and AdamW (Table 4) are small enough (e.g., 12.452 vs. 12.372 at 1B scale) that seed variance could plausibly account for some fraction of the gap. Running 3–5 seeds at the 125M–350M scale would substantially strengthen the reliability claims.

3. The difficulty estimation cost comparison is missing. The paper's core practical claim — 33% wall-clock time reduction — is measured for a complete training run starting from initialization. But the Hessian-based partitioning requires knowing the architecture's Hessian structure in advance, which the paper derives from a separate small-scale analysis. If a practitioner encounters a new architecture (e.g., a custom Transformer variant with different attention patterns), they would need to either (a) run their own small-scale Hessian analysis (cost not quantified) or (b) guess at the partition. The paper does not discuss how much compute the Hessian analysis in Figure 7 cost, nor whether there is a cheaper way to derive the partition for new architectures. This is a practical gap analogous to the difficulty estimation cost problem in the reference example's Section 6.

4. The value matrix ambiguity is not fully resolved. The paper acknowledges that value's Hessian structure is "less clear" (Figure 7c) and that the optimal partition depends on training length (Appendix D.6). This means the partitioning principle is not fully automatic — it requires a judgment call (Strategy I vs. Strategy II) that the paper does not provide a rigorous criterion for. The recommendation ("Strategy I when training steps are large") is based on empirical observation rather than a structural criterion. For a new architecture, an analogous ambiguity could arise for other parameter groups, and the paper provides no systematic way to resolve it beyond trial and error.

5. Limited non-Transformer architecture testing. The paper's claim of general applicability (Algorithm 3 for non-Transformers) is supported by experiments on ResNet-18, Swin-Transformer, DiT-XL-2, DC-AE-Diffusion, DDPM, GCN, and GAT — but notably, all of these use the PyTorch default partition (Algorithm 3). Principle 1 is not actually applied to derive architecture-specific partitions for CNNs or GNNs; instead, the paper falls back to "one block per named parameter." This works in practice (Table 6), but it means the Hessian-structure-driven design principle is only really exercised for Transformers. For CNNs, it is possible that the default partition happens to align with the Hessian structure; for other architectures, it might not. The paper acknowledges this limitation in Appendix B: "In the future, it is possible that we will have more complicated non-Transformer architectures on which Algorithm 3 fails."

6. The throughput gains conflate memory capacity with optimizer efficiency. The 49.6% throughput improvement (Table 2) comes from Adam-mini's ability to use batch size 4/GPU where AdamW can only use batch size 1/GPU. This is a genuine practical advantage, but it is not an optimizer speed advantage per se — it is a capacity advantage. If both optimizers were run at the same batch size on hardware where neither was memory-constrained, the throughput difference would reflect only the per-step computational savings (cheaper mean and scalar sqrt operations), which the paper does not measure separately. In a well-provisioned multi-GPU setup where AdamW already runs at the optimal batch size, the throughput gain from Adam-mini would be much smaller.

7. The RLHF results are on a single algorithm (ReMax) and dataset (UltraFeedback). The paper's RLHF experiments (Figure 12b) show Adam-mini outperforming AdamW, but only in the specific configuration of ReMax on UltraFeedback with Llama 2-7B. RLHF is known to be sensitive to the choice of RL algorithm (PPO, ReMax, DPO), reward model quality, and prompt distribution. A single positive result does not establish that Adam-mini is generally superior for RLHF. The SFT results are similarly limited to a single dataset, though SFT is generally more stable across configurations.

Missing Experiments That Would Strengthen the Paper

Multi-seed runs at the 125M–350M scale to quantify variance. Even 3 seeds per configuration would allow error bars on the scaling law plots (Figure 11) and clarify whether the small perplexity differences in Table 4 are statistically reliable.
A moderate-scale Hessian analysis (at least one layer of a 125M-parameter model) to bridge the massive gap between Figure 7's toy model and the main experiments. This is computationally expensive but would substantially strengthen the causal claim that Hessian structure drives partitioning effectiveness.
Ablation on the number of blocks: what happens if the partition is coarser (e.g., entire layers) or finer (e.g., individual parameters within each head, reverting to Adam)? This would empirically bound the sensitivity to block granularity and test whether Principle 1's "smallest dense sub-block" specification is precisely optimal or just a safe default.
Throughput comparison isolating per-step computation from batch-size effects. Running both optimizers at identical batch sizes on memory-unconstrained hardware would measure the pure computational speedup from mean(g⊙g) and scalar sqrt, separate from the capacity benefit.
Experiments on a non-Transformer architecture that actually requires a custom Hessian-based partition. Testing Adam-mini on, e.g., a state-space model or a custom hybrid architecture where the PyTorch default partition is known to be suboptimal and Principle 1 must be applied from scratch would test the generality of the methodology.
Adam-mini with β_2 = 0.999 (the standard Adam default) vs. β_2 = 0.95. All LLM experiments use β_2 = 0.95, which is standard for LLM training but different from Adam's original default. An ablation showing that mean(v) works with both settings would clarify whether Adam-mini's behavior is tied to the shorter second-moment memory.

6. Limitations and Trade-offs

The Hessian Analysis That Drives Partitioning Is at a Scale 3–6 Orders of Magnitude Smaller Than the Target Models

The assumption or constraint. The entire partitioning principle (Principle 1) is derived from Hessian visualizations on a 1-layer Transformer with n_emb=16, n_head=4, and MLP width 32 — fewer than 10,000 parameters (Figure 7). The main experiments deploy this partitioning on models up to 13B parameters (Llama 2-13B) — a scale difference of roughly 6 orders of magnitude. The paper does not attempt even a moderate-scale Hessian analysis (e.g., a single layer of a 125M model) to verify that the near-block-diagonal structure observed in the toy model persists at scale. The computational cost of exact Hessian computation (O(n²) memory for the full matrix) makes large-scale verification infeasible with the Pearlmutter (1994) method used in the paper.

The consequence. The causal claim — that the Hessian-aware partition works because it respects Hessian sub-block boundaries — remains a hypothesis supported by indirect evidence rather than a verified mechanism. The large-scale training stability (Figures 8, 10, 11) could be explained by factors other than Hessian structure: the partition might coincidentally align with gradient variance patterns, or the block-wise averaging might simply provide beneficial regularization independent of Hessian geometry. If the Hessian structure changes qualitatively at scale — for instance, if the dense sub-blocks merge, split, or become less sharply defined as model depth increases — then Principle 1 might prescribe a suboptimal partition for larger models that the small-scale analysis cannot anticipate. The paper provides no diagnostic for detecting such regime shifts.

What evidence exists in the paper. Only indirect evidence. The leave-x-out experiments (Figure 6) are on a 4-layer Transformer, still far smaller than the target models. The failure of PyTorch default partitioning on Llama 2-1B (Figure 7i) demonstrates that the partition matters, but does not prove that Hessian structure is the reason the specific partition in Algorithm 3 works — it could be that any fine-grained partition (not necessarily aligned with dense sub-blocks) would stabilize training, and the Hessian-based partition simply happens to be fine-grained enough. The paper does not test alternative fine-grained partitions (e.g., random groupings of similar granularity) to rule this out.

Mitigation status. The paper does not explicitly acknowledge this scale gap as a limitation. It treats the small-scale Hessian analysis as sufficient evidence for the partitioning principle, and relies on the empirical success at scale as validation. The closest the paper comes to addressing this is the statement that "the near-block-diagonal structure maintains throughout training" (Section 2.3, referring to Figure 3b–d for MLPs), but this claim about persistence across training steps does not address persistence across model scales. A moderate-scale Hessian analysis (e.g., computing the Hessian for one layer of a 125M model using iterative methods like Lanczos, which can estimate eigenvalue/block structure without forming the full matrix) would substantially strengthen the causal link but is left to future work.

The `value` Matrix Partitioning Is Ambiguous and Requires Task-Specific Judgment

The assumption or constraint. Algorithm 3 prescribes partitioning value by output neurons, following the same rule as attn.proj and mlp. However, the Hessian visualization for value (Figure 7c) shows an ambiguous structure that the paper itself describes as having "the hint of 16 diagonal blocks (16 is the number of output neurons), but the pattern is less obvious" (Section 2.3). The Hessian entries for value are also about 10⁶ larger than those of query and key. The paper acknowledges this ambiguity and provides a second strategy — treating value as a whole (Strategy II) — recommending it for short training runs (Appendix D.6). The choice between strategies is based on empirical observation of which works better for a given training duration, not on a structural criterion derived from Principle 1.

The consequence. The partitioning strategy for value is not fully principled — it requires a judgment call that the paper does not reduce to a deterministic rule. This has two practical implications. First, a practitioner deploying Adam-mini on a new architecture or training regime would need to either (a) conduct their own ablation comparing Strategy I vs. Strategy II, which adds a tuning cost the paper otherwise claims to eliminate, or (b) default to Strategy I and accept potentially suboptimal performance on short training runs. Second, the value ambiguity suggests that Principle 1 may be incomplete as stated: it does not provide guidance for cases where the Hessian block structure is present but weak or in transition. For architectures with attention variants that position certain projections differently relative to nonlinearities (e.g., linearized attention, gated attention), analogous ambiguities could arise for other parameter groups, and the paper provides no systematic framework for resolving them short of trial and error.

What evidence exists in the paper. Appendix D.6 reports that Strategy I (partition by output neurons) "works well when the number of total training steps is large" and Strategy II (treat as a whole) "works better when the number of total training steps is small." The Llama experiments in Figure 10 (10,000 steps) use Strategy II; the GPT-2 experiments in Figure 8 and scaling law experiments in Figure 11 (50,000+ steps) use Strategy I. This is presented as a finding rather than a failure mode, but it reveals that the partitioning is contingent on a hyperparameter (training duration) that is not part of the partitioning principle itself.

Mitigation status. The paper partially mitigates this by providing both strategies and a one-line code override (optimizer.wv_names = {}) to switch between them. However, the criterion for choosing is imprecise (what counts as "large" vs. "small" number of steps? The paper does not define a threshold), and the paper does not test whether the performance difference between the two strategies is large enough to matter in contexts outside the specific training budgets they studied. No experiment quantifies the sensitivity of final performance to this choice across a range of training durations.

Throughput Gains Are Conflated with Memory Capacity Benefits and Depend on Hardware Configuration

The assumption or constraint. The paper's headline throughput improvement — 49.6% higher than AdamW when pre-training Llama 2-7B on 2× A800-80GB GPUs (Table 2) — is measured in a regime where Adam-mini can use batch size 4 per GPU while AdamW is limited to batch size 1 per GPU (batch size 2 per GPU for AdamW runs out of memory). The throughput comparison therefore reflects a combined effect of (a) reduced per-step computation from cheaper mean(g⊙g) and scalar sqrt operations, and (b) reduced communication overhead because larger per-GPU batch sizes mean fewer gradient accumulation steps and fewer all-reduce operations. The paper does not disentangle these factors or report the throughput gain in a setting where both optimizers run at the same per-GPU batch size.

The consequence. The 49.6% figure should not be interpreted as a universal speedup that applies regardless of hardware provisioning. In a setting with more GPUs where AdamW already fits the optimal batch size without memory pressure — e.g., 8× GPUs instead of 2× — the capacity-driven portion of the throughput gain would disappear entirely. The remaining gain would come only from per-step computational savings (cheaper second-moment operations), which the paper does not measure in isolation. The 33.1% wall-clock time reduction for Chinchilla-optimal training (Table 2: 3,489.55 vs. 5,219.16 GPU-hours for 70B tokens) inherits this conflation: it assumes a 2× GPU configuration where AdamW is memory-constrained, and would be smaller on better-provisioned hardware. Practitioners with different GPU counts, memory capacities, or model parallelism strategies may see substantially smaller throughput benefits.

What evidence exists in the paper. Table 2 reports throughput numbers only for the 2× A800-80GB configuration. No throughput measurements are reported for other GPU counts, for configurations where both optimizers run at the same per-GPU batch size, or for models smaller than 7B (where AdamW would be less memory-constrained). Figure 13c shows a throughput comparison with Adafactor on Llama 2-1B (Adam-mini achieves 40% higher throughput), but this is also at a fixed hardware configuration and the breakdown of capacity vs. computation is not analyzed.

Mitigation status. The paper does not explicitly separate computational savings from memory capacity gains in its throughput analysis. The statement that throughput improvements come from two factors — "Adam-mini does not introduce extra computation in per-step updates" and "the memory cut-down allows larger batch sizes per GPU... eases the burden of communication among GPUs" (Section 2.4) — correctly identifies both mechanisms but does not quantify their individual contributions. This makes it difficult for practitioners to estimate throughput benefits for their specific hardware setups without running their own benchmarks.

All Large-Scale Validation Is on a Single Family of Architectures (GPT-2 and Llama Transformers) and a Single Type of Task (Language Modeling)

The assumption or constraint. The paper's primary claims — that Adam-mini "performs on par or better than AdamW" and maintains its scaling laws — are validated exclusively on autoregressive Transformer language models (GPT-2 series, Llama series) trained on standard English text corpora (OpenWebText, C4). The non-LLM experiments (ResNet-18, Swin-Transformer, DiT, DDPM, GCN, GAT; Table 6 and Figures 17-18) use the simpler Algorithm 3 partition (PyTorch default parameter blocks), not the Hessian-aware Transformer partition in Algorithm 3. The Hessian analysis and partitioning principle (Principle 1) are therefore only actively exercised and validated for one architecture family under one training paradigm.

The consequence. The paper does not establish that the Hessian-aware partitioning methodology generalizes to architectures that are not standard autoregressive Transformers. For encoder-decoder Transformers (e.g., T5), encoder-only Transformers (e.g., BERT), mixture-of-experts models, state-space models (e.g., Mamba), or retrieval-augmented architectures, the Hessian block structure may differ — and Principle 1 would prescribe different partitions that the paper provides no guidance for deriving. The claim that "for different architectures, the principle will be realized in different forms" (Section 2.2) is aspirational rather than demonstrated. A practitioner applying Adam-mini to a new architecture would need to either (a) conduct their own small-scale Hessian analysis (following the methodology in Figure 7), which the paper does not provide a reproducible pipeline for, or (b) fall back to Algorithm 3 (PyTorch default) and hope it works, which the paper shows fails for Transformers at scale (Figure 7i).

Additionally, all pre-training experiments use standard English text. The behavior of Adam-mini on multilingual corpora, code, or domain-specific data (biomedical, legal, scientific) is not tested. If the Hessian's near-block-diagonal structure depends on the cross-entropy loss coupling properties that Collobert (2004) analyzed — specifically, the p(x)(1-p(x)) factor that shrinks off-diagonal blocks — then data distributions that produce different model confidence dynamics could potentially alter the Hessian structure and invalidate the partition.

What evidence exists in the paper. The non-LLM experiments (Table 6) show that the simpler PyTorch-default partition works for CNNs, vision Transformers, diffusion models, and graph networks — but these are using Algorithm 3, which does not implement Principle 1. They demonstrate that Adam-mini with coarse blocks can work for some architectures, but do not test whether Hessian-aware fine-grained partitioning (analogous to Algorithm 3) would be better. The paper does not report any experiment where Principle 1 was explicitly applied to derive a custom partition for a non-Transformer architecture and then validated.

Mitigation status. The paper partially acknowledges this in Appendix B: "In the future, it is possible that we will have more complicated non-Transformer architectures on which Algorithm 3 fails. In those cases, we need to investigate the Hessian structure of these new architectures (like what we did for Transformers) and then develop the concrete partition algorithms following our Principle 1." This is a clear statement of the limitation, but it delegates the solution entirely to future work without providing tools or a methodology for practitioners to conduct such investigations themselves. The paper does not release code for the Hessian visualization or the leave-x-out analysis that would enable others to replicate the partitioning discovery process on new architectures.

No Multi-Seed Evaluation or Statistical Quantification of Variance

The assumption or constraint. Nearly all experiments in the paper appear to be single-run: one seed, one training trajectory, one final metric per configuration. The main pre-training curves (Figures 8, 10, 11), the SFT and RLHF results (Figure 12), and the scaling law experiments (Table 4, Figure 11) report point estimates with no error bars, confidence intervals, or multi-seed averages. The paper does not state the number of seeds used, does not report min/max/standard deviation across runs, and does not discuss run-to-run variance.

The consequence. The paper's central empirical claim — that Adam-mini performs "on par or better than AdamW" — cannot be evaluated for statistical reliability from the reported results. LLM training is known to exhibit non-trivial seed sensitivity, particularly for smaller models (the 39M–271M range in the scaling law experiments) and shorter training runs. The final perplexity differences reported in Table 4 (e.g., AdamW 12.452 vs. Adam-mini 12.372 at the 1B scale, a difference of 0.080) could be within the range of seed-to-seed variance for these training configurations. The conclusion that "Adam-mini reaches a slightly lower perplexity than AdamW for all models" (Section 3.2) would be substantially weakened if any of these differences reversed under a different seed — and without error bars, the reader cannot assess how likely that is.

The leave-x-out experiments (Figure 6) — which provide the critical bridge between the abstract quadratic analysis and the full Adam-mini design — show performance gaps that vary across randomly selected blocks. Some leave-one-out trials perform better than Adam; some perform worse. The paper does not report the variance or statistical significance of these differences, making it unclear whether the finding that "Adam (leave-one-out) always performs on par with Adam" (Section 2.1) is robust or whether some blocks consistently degrade when simplified.

What evidence exists in the paper. The trajectory comparison (Figure 9b) uses a single seed (stated in Appendix F.2: "launch AdamW, Adam-mini, and other memory-efficient optimizers under the same random seed"). No other experiment explicitly states the number of seeds. The learning rate sensitivity analysis (Figure 12c) shows performance across a range of learning rates, providing some evidence of stability under hyperparameter variation, but does not address seed-to-seed variance at a fixed hyperparameter setting.

Mitigation status. The paper does not acknowledge the single-seed limitation or discuss variance. This is a significant methodological weakness for a paper that makes comparative performance claims across a grid of model sizes (the scaling law experiments) and across multiple competing optimizers (the pre-training comparisons with Adafactor, CAME, SM3, LAMB). Standard practice in the optimization literature for LLMs varies — some papers report multi-seed results (e.g., mean ± std over 3–5 seeds at moderate scales) while others rely on single-run trends at scale — but the absence of any variance quantification makes it difficult to assess which of the reported performance gaps are real versus within the noise floor. Running even 3 seeds at the 125M or 350M scale would substantially improve the reliability of the paper's empirical claims without prohibitive computational cost.

The Memory Savings Apply Only to `v`; First-Order Momentum and Gradients Are Unchanged

The assumption or constraint. Adam-mini reduces only the second-order momentum v to block-wise scalars. The first-order momentum m is stored at full parameter resolution, identical to AdamW. The total optimizer state memory goes from 2 × model_size (for m and v) to approximately 1 × model_size (for m only, with v reduced to negligible size). However, this means the overall GPU memory reduction — including model parameters, activations, and gradients — is smaller than 50%. With gradients adding another 1 × model_size of memory, the total memory for model + optimizer + gradients drops from roughly 3 × model_size to 2 × model_size — approximately a 33% reduction, not 50%.

The consequence. The paper's headline claim of "50% less memory footprint" (abstract) refers specifically to optimizer state memory and excludes activations, which can dominate GPU memory for large-batch or long-sequence training. A practitioner reading the abstract might expect their total GPU memory usage to halve; in practice, it drops by roughly one-third. This can matter for hardware provisioning: if a model requires 60 GB total GPU memory with AdamW (20 GB model + 40 GB optimizer + gradients), Adam-mini would reduce this to roughly 40 GB — a meaningful reduction but not enough to fit the model on a substantially smaller GPU class (e.g., going from A100-80GB to RTX 4090-24GB still requires an additional 2× compression factor via quantization or sharding). The throughput gains from larger per-GPU batch sizes are correspondingly moderated by the activation memory, which Adam-mini does not touch.

What evidence exists in the paper. Table 1 reports memory numbers for optimizer states only ("Calculation is based on float32, which is a standard choice for optimizer states"). The text in Section 1 states "Adam requires the memory for its optimizer states... These in total take at least 2× the memory of the model size." Figure 1a shows a bar chart labeled "Memory" where Adam-mini's bar is approximately 0.75× the height of AdamW's — consistent with a ~25% reduction in total GPU memory (including activations), not 50%. The paper does not explicitly reconcile the 50% optimizer-state claim with the ~25% total-memory reduction shown in the figure, which could confuse readers who do not distinguish between optimizer memory and total memory.

Mitigation status. The paper is precise in the abstract ("50% less memory footprint" is followed immediately by "Adam-mini reduces memory by cutting down the learning rate resources in Adam") but does not always maintain this precision in the main text. In Section 2.4, "which saves 50% of Adam's memory" could be read as referring to total memory. The distinction matters because the practical benefit — fitting larger batch sizes or larger models — depends on total memory, not optimizer memory alone. The paper partially addresses this by reporting actual per-GPU batch sizes in the throughput experiments (Table 2: Adam-mini can use batch size 4 while AdamW is limited to 1–2), which gives practitioners a concrete sense of the practical capacity improvement even if the percentage reduction is smaller than 50% of total memory.

7. Implications and Future Directions

How This Work Changes the Landscape

This paper introduces a new design paradigm for adaptive optimizers: rather than treating the optimizer as a generic algorithm applied uniformly to all parameters, optimization granularity should be matched to the problem's second-order structure. This is not a new optimizer trick — it is a reframing of what memory-efficient optimization should optimize for. Before Adam-mini, the dominant approach was to compress Adam's v using generic tensor decomposition (Adafactor's low-rank factorization, SM3's minimal-value selection) and accept degraded performance as the cost of memory savings. The implicit operating assumption was that Adam's coordinate-wise learning rates are jointly necessary, that any simplification loses information, and that the best we can do is find a compression scheme that minimizes the information loss per bit of memory.

Adam-mini falsifies this assumption. The paper shows that coordinate-wise learning rates are not jointly necessary — within dense Hessian sub-blocks, they can be actively counterproductive (Figure 5: the preconditioned condition number r can exceed 1, meaning 1/√v makes the problem harder). This means previous memory-efficient optimizers were solving the wrong problem: they were trying to approximate a function (coordinate-wise 1/√v) that is itself suboptimal in the regions where most parameters live. The correct problem — which Adam-mini solves — is to identify the natural granularity at which learning rates genuinely matter and discard the rest.

This shift has several concrete consequences for the field:

It reconciles the contradictory evidence on memory-efficient optimizers. Before this paper, the literature presented a confusing picture: Adafactor saves memory but underperforms (Luo et al., 2023; this paper's Figures 13, 19); layer-wise learning rate methods like LAMB add complexity without saving memory; sign-based methods sometimes work (Lion on some tasks; Chen et al., 2024b) but fail unpredictably (this paper's Figure 20, where Lion encounters loss spikes on GPT-2-125M). Adam-mini provides a unified explanation: these methods fail not because memory savings inherently degrade performance, but because they apply structure-agnostic compression. Low-rank factorization, minimal-value selection, and sign-based updates all ignore the Hessian's near-block-diagonal structure, so they inevitably either (a) merge distinct Hessian blocks (causing instability, as in Figure 7i with PyTorch default partition) or (b) preserve coordinate-wise granularity where it's unnecessary. Adam-mini succeeds because it respects the boundary between "different blocks need different rates" (which is true and important) and "different parameters within a block need different rates" (which is false in dense Hessian regions).

It redirects research attention from compression techniques to structure discovery. The paper's most influential contribution may ultimately be not Adam-mini itself, but the methodology it demonstrates: (1) analyze the Hessian structure of the target architecture at small scale, (2) identify the smallest dense sub-blocks, (3) assign one learning rate per sub-block, (4) validate at scale. This shifts the optimization community's focus away from developing increasingly clever generic compression schemes and toward understanding the second-order geometry of specific architectures. For a new architecture — a mixture-of-experts layer, a state-space model, a retrieval-augmented Transformer — the first question an optimizer designer should ask is no longer "how can we factor v?", but "what are the smallest dense Hessian sub-blocks, and how do we partition to match them?" This is a more principled starting point, and it makes optimizer design architecture-aware by default rather than as an afterthought.

It elevates the Hessian from a diagnostic tool to a design tool. Previous work used Hessian analysis for diagnosis — explaining why Adam works (Zhang et al., 2024), characterizing loss landscape geometry (Sagun et al., 2017; Ghorbani et al., 2019), or identifying pathologies (Dauphin et al., 2024). This paper shows that Hessian structure can directly prescribe optimizer hyperparameters (the parameter partition). This is a conceptual upgrade: the Hessian is not just something we measure to understand our optimizer, but something we use to build our optimizer. The Collobert (2004) near-block-diagonal result, which was known for two decades but had no practical optimizer-design consequence, now becomes directly actionable. This creates a template for future work: any time a structural property of the loss landscape is identified, the natural follow-up question is "can we use this to reduce the optimizer's degrees of freedom without losing performance?"

It makes the case that optimizer design and architecture design are coupled. The paper's partitioning principle (Principle 1) is architecture-specific: the Transformer partition in Algorithm 3 works because Transformers have attention heads and output neurons that naturally create dense Hessian sub-blocks. If the architecture were different — say, with depthwise-separable convolutions or gated linear units with different connectivity patterns — the Hessian blocks would likely be different, and the correct partition would change. This implies that the optimizer and the architecture are not independent design choices; the optimizer's parameter grouping should be informed by the architecture's Hessian structure. There is a future research program here: for each new architecture, derive its Hessian-aware partition, and compare the resulting "architecture-native" optimizer against generic Adam. If architecture-native optimizers consistently outperform generic ones, it would motivate co-designing architectures and optimizers from the start — much as hardware-software co-design has become standard in systems engineering.

The magnitude of this shift should not be overstated. This is not a paradigm shift on the scale of Adam itself replacing SGD for Transformers (Kunstner et al., 2023; Zhang et al., 2024). Adam-mini modifies Adam's v — it does not replace the adaptive update rule, the momentum mechanism, or the decoupled weight decay. It is better understood as a refinement of Adam's granularity that reveals a previously hidden degree of freedom: the coordinate-wise learning rate budget is largely wasted within dense Hessian blocks, and reclaiming that budget via structure-aware blocking yields memory savings with no performance cost. The paper's strongest conceptual contribution is not the optimizer itself but the diagnostic framework that identifies where in the parameter space Adam is over-provisioned and why reducing granularity there is safe.

Follow-Up Research This Work Enables

Cheap, scalable Hessian block discovery for arbitrary architectures. The paper's Hessian visualizations (Figure 7) require forming the full n×n Hessian matrix — feasible only for toy models with n < 1000. A critical practical gap is developing methods to identify the smallest dense Hessian sub-blocks at moderate scale (e.g., one layer of a 125M-parameter model) without O(n²) memory. Potential approaches include: (1) estimating block structure from the gradient outer product E[g·g^T] (which approximates the Fisher information and shares block structure with the Hessian under cross-entropy loss), computable via power iteration on the gradient covariance; (2) using randomized numerical linear algebra (Hutchinson's estimator, randomized SVD) to probe block-diagonal structure without forming the full matrix; (3) computing Hessian-vector products along random directions and clustering parameters by response similarity — parameters in the same dense block should have correlated second-order responses. A successful method would enable applying Principle 1 to new architectures (e.g., Mamba, mixture-of-experts, multimodal Transformers) without the paper's manual small-scale analysis step, making Adam-mini truly "plug-and-play" for any architecture. The evaluation would measure whether partitions derived from the scalable method produce the same blocks as the exact Hessian analysis on the small models where the exact analysis is feasible, and whether they yield stable training at scale.

Optimal block-wise learning rates beyond mean(v). The paper chooses mean(v) for practical reasons — it is cheap, it keeps the trajectory close to Adam, and it avoids grid search — but explicitly acknowledges it "might not be an optimal way" (Section 4). The block-wise GD experiments in Figure 4 show that block-wise optimal learning rates (derived from eigenvalues of each dense sub-block) converge faster than Adam and much faster than Adam-mini, but require computing or estimating the largest and smallest eigenvalues per block. A natural follow-up is to develop a lightweight eigenvalue estimator that runs online during training — e.g., maintaining a power iteration estimate of the top eigenvalue of each block's Hessian (using Hessian-vector products computed from the Pearlmutter trick) and setting the block learning rate to 2/(L_max + L_min) or a similar optimal rate for quadratics. The cost would be one extra backward pass per training step (for the Hessian-vector product), making it roughly 2× more expensive per step than Adam-mini, but potentially converging in significantly fewer steps — a compute-optimal trade-off analogous to the paper's own throughput-vs-memory analysis. A strong follow-up would measure whether the per-step overhead is offset by faster convergence on Llama 2-1B pre-training, comparing wall-clock time to reach a target validation loss against both AdamW and vanilla Adam-mini.

Combining Adam-mini with momentum compression. Adam-mini reduces only v; m is stored at full parameter resolution. Recent work on low-precision optimizers (8-bit Adam; Dettmers et al., 2021) and 4-bit optimizer states (Li et al., 2024) shows that m can be quantized with minimal performance loss. A natural combination is "Adam-mini-8bit": apply block-wise v (reducing v memory by ≥99.9%) and 8-bit quantization to m (reducing m memory by ~75% from float32 to int8), yielding total optimizer state memory of approximately 0.25 × model_size for m + negligible for v — a ~87.5% reduction from AdamW's 2 × model_size. This would be particularly impactful for on-device training or edge deployment of LLMs, where every GB of memory matters. The experiment would pre-train GPT-2-125M with Adam-mini-8bit vs. Adam-mini (float32 m) and AdamW-8bit, measuring whether the combination of both compression axes remains within the noise floor of float32 Adam-mini performance.

Stress-testing the Hessian structure assumption: when does it break? The paper's partitioning principle relies on the near-block-diagonal Hessian being a stable property of neural network training. But this property has only been verified empirically for MLPs and standard Transformers under cross-entropy loss on discrete prediction tasks. Important open questions: Does the block structure persist under contrastive losses (common in representation learning)? Under regression losses (MSE)? Under reinforcement learning objectives (where the loss surface is non-stationary due to the changing policy)? At extreme model widths or depths (where the Collobert (2004) analysis may not apply due to the p(x)(1-p(x)) factor behaving differently in overparameterized regimes)? A systematic stress-test would train small models under each of these conditions, compute Hessians using the Pearlmutter method, and quantify the "diagonal-over-off-diagonal ratio" τ (from Figure 5) for candidate blocks across training. If τ drops below some threshold — indicating the block is becoming less block-diagonal — Adam-mini's partition might need to be coarsened or refined dynamically. Negative results (tasks where the block structure degrades and Adam-mini underperforms AdamW) would be as informative as positive ones, because they would establish the boundary conditions for Principle 1's applicability.

Adam-mini for continual pre-training and domain adaptation. The paper tests Adam-mini on pre-training from scratch, supervised fine-tuning, and RLHF — but not on continual pre-training, where a pre-trained model is further trained on a new domain (e.g., continued pre-training of Llama 2 on code or scientific text). This setting is practically important (Ibrahim et al., 2024) and differs from pre-training in that the Hessian structure may have already crystallized during the initial pre-training phase. If the Hessian blocks from pre-training persist into the new domain, Adam-mini should work identically. If the new domain's data distribution alters which parameter groups are densely coupled, the pre-training partition might become suboptimal, and Adam-mini's performance could degrade relative to full AdamW. A concrete experiment: take the pre-trained Llama 2-7B, continue pre-training on the Stack (code) or PubMed (biomedical text) using both Adam-mini (with the same partition as the original pre-training) and AdamW, measuring validation loss on the new domain. If Adam-mini tracks AdamW, it confirms robustness to domain shift; if not, it suggests that Adam-mini may need domain-specific partitions, which would be a significant practical limitation for continual pre-training workflows.

Scaling laws for Adam-mini to 70B+ parameters. The paper's scaling law experiments (Figure 11) go up to 1B parameters and extrapolate via fitted power laws ("the fitted lines in Figure 11 (b) suggest that Adam-mini can be scaled up to larger models"). But extrapolation of optimizer behavior across an order of magnitude (1B → 70B) is unvalidated, and there are plausible failure modes: at extreme scale, the Hessian structure might change qualitatively (e.g., the dense sub-blocks within attention heads might start to couple across heads due to the increased representational capacity), or the mean(v) approximation might drift from Adam's per-parameter v as gradient distributions become heavier-tailed. Training a single Llama 2-70B model with both AdamW and Adam-mini to 1–2B tokens (enough to see whether the training trajectories diverge) would either validate the scaling law extrapolation or reveal a breakdown that would bound Adam-mini's applicability. The paper's own hardware constraints (4× A800-80GB GPUs) make 70B training infeasible, but a collaboration or a cloud compute grant could resolve this — and the result would determine whether Adam-mini can be adopted for frontier-scale training or is limited to smaller-scale and fine-tuning workloads.

Practical Applications and Downstream Use Cases

Reducing GPU count for LLM pre-training at the 7B–13B scale. The paper's throughput results (Table 2) directly translate to hardware savings: pre-training Llama 2-7B on 2× A800-80GB GPUs with Adam-mini processes the Chinchilla-optimal 140B tokens in approximately 6,979 GPU-hours vs. 10,438 GPU-hours for AdamW — a 33% reduction. For a startup or academic lab renting cloud GPUs at ~ $2/GPU-hour, this saves approximately$ 7,000 for a single 7B pre-training run. More significantly, it enables pre-training on fewer GPUs: a configuration that would require 4× GPUs with AdamW (due to memory constraints) might fit on 2× GPUs with Adam-mini, halving the hardware requirement. This directly addresses the paper's stated motivation of "lowering the threshold of training LLMs and encouraging participation from more diverse researchers." The key caveat is that the 33% figure assumes the specific 2× GPU configuration where AdamW is memory-constrained; on better-provisioned hardware, the savings would be smaller but still non-zero due to the reduced per-step computation from scalar sqrt and division operations.

Memory-efficient supervised fine-tuning and alignment of large models. The SFT and RLHF results (Figures 12a, 12b, Table 5) demonstrate that Adam-mini works not just for pre-training but also for the fine-tuning stages that practitioners most commonly interact with. For fine-tuning Llama 2-7B, Adam-mini reduces optimizer memory from ~56 GB to ~28 GB — a 28 GB saving that can be reallocated to larger batch sizes, longer sequence lengths, or fitting the full model on a single GPU where AdamW would require model parallelism. This is particularly impactful for RLHF, which typically requires maintaining multiple model copies (policy, reference, reward model, value model) simultaneously in memory. Adam-mini applied to all copies could reduce the total optimizer memory from 4 × 2 × model_size = 8 × model_size to 4 × 1 × model_size = 4 × model_size, freeing enough memory to use larger batch sizes or fit the pipeline on fewer GPUs. The MT-Bench results (Table 5) further show that Adam-mini's alignment quality is slightly better than AdamW's, meaning the memory savings come with no degradation — and potentially a small improvement — in the final model's chat capabilities.

Accelerating scaling law experiments for architecture search. The paper's scaling law results (Figure 11) demonstrate that Adam-mini can substitute for AdamW in the proxy model experiments used to predict optimal large-scale configurations. Scaling law experiments are typically run on models at 10M–1B scale to fit power laws and extrapolate to 70B+ — but even these proxy experiments are expensive (the paper reports 300 GPU-hours for its scaling law runs). Adam-mini's 33% wall-clock time reduction directly accelerates these experiments by the same proportion, or equivalently, allows testing 50% more configurations in the same compute budget. For an organization running regular scaling law experiments to guide architecture decisions, this is a direct cost saving with no methodological change required — since Adam-mini uses the same hyperparameters and produces the same scaling trends as AdamW, the fitted laws can be used identically.

Combining with parameter-efficient fine-tuning to push LoRA to larger models or higher ranks. Appendix D.10 shows that replacing Adam with Adam-mini inside LoRA fine-tuning improves evaluation perplexity (Figure 22). More practically, the memory savings from Adam-mini could enable higher LoRA ranks (e.g., rank 256 or 512 instead of the standard 128) on a given GPU budget — or allow LoRA fine-tuning of larger base models (e.g., Llama 2-13B with LoRA on a single consumer GPU where AdamW's optimizer states would exceed memory). Since LoRA applies Adam only to the low-rank adapters (which are small), the absolute memory savings from Adam-mini are modest in this context — but for high-rank LoRA on very large models where the adapters themselves contain billions of parameters, the 50% optimizer memory reduction on the adapter parameters becomes meaningful. A practical configuration: LoRA rank 512 on Llama 2-70B, where the adapters total ~280M parameters (requiring ~2.2 GB of optimizer states with AdamW), would require only ~1.1 GB with Adam-mini — a small absolute saving but one that might determine whether the configuration fits on a single 80 GB GPU alongside the frozen base model and activations.