🤖 machine learning

FOAM: Blocked State Folding for Memory-Efficient LLM Training

The paper introduces FOAM, a memory-efficient optimizer that compresses optimizer states via block-wise gradient means with residual correction, achieving convergence rates equivalent to vanilla Adam while eliminating up to 90% of optimizer memory overhead and accelerating training without degrading performance.

Original authors: Ziqing Wen, Jiahuan Wang, Ping Luo, Dongsheng Li, Tao Sun

Published 2026-05-14

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Ziqing Wen, Jiahuan Wang, Ping Luo, Dongsheng Li, Tao Sun

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a giant, super-smart robot (a Large Language Model) how to write, chat, and reason. To do this, you show it a massive library of books (training data). The robot learns by making mistakes, checking its work, and adjusting its internal "brain settings" (parameters) to get better.

The standard way to do this adjustment is like a very careful accountant named Adam. Adam keeps a detailed ledger for every single brain setting, tracking not just the current mistake, but also the history of past mistakes to decide how much to change things.

The Problem: The Accountant is Too Heavy
The problem is that this "accountant" (the optimizer) is incredibly heavy. For a robot with billions of brain settings, the accountant's ledger takes up twice as much memory as the robot's brain itself.

Analogy: Imagine trying to move a house. The furniture (the model) is big, but the moving truck (the optimizer states) is even bigger. If you only have a small truck (limited computer memory), you can't move the house at all, no matter how good the house is. This is the "memory bottleneck" that stops researchers from training bigger, smarter robots.

The Old Solutions: Cutting Corners
Previous attempts to shrink the moving truck had flaws:

Freezing parts of the house: You stop moving some furniture and only adjust a few small boxes. This saves space but makes the house less flexible (lower performance).
Taking blurry photos: Instead of moving the actual furniture, you take low-resolution photos of it to save space. But calculating these photos takes a lot of time and energy (computational overhead).
Sharing ledgers: You make different rooms share the same notebook. This saves space but makes the adjustments less precise.

The New Solution: FOAM (The Smart Folding Method)
The paper introduces FOAM (Folded Optimizer with Approximate Moment). Think of FOAM as a master packer who uses a clever "folding" technique to fit the massive moving truck into a tiny van, without losing any important details.

Here is how FOAM works, using simple metaphors:

1. The "Block Folding" (Compression)

Instead of tracking every single brain setting individually, FOAM groups them into small blocks (like a grid of 8 or 16 settings).

The Metaphor: Imagine you have a long line of 100 people, and you need to remember their average height. Instead of writing down 100 separate numbers, you group them into 10 teams of 10. You write down just one number: the average height of each team.
The Result: This shrinks the memory needed for the "ledger" by a huge amount (up to 90% less).

2. The "Residual Correction" (The Safety Net)

If you only used the averages, you'd lose the unique details of each person. Maybe one person in a team is very tall, and another is very short. The average hides that.

The Metaphor: FOAM is smart. After it writes down the "team average," it quickly calculates the difference (the error) between the real people and the average. It calls this the "Residual."
The Trick: It doesn't keep this difference forever. It calculates it, uses it to fix the update for that specific moment, and then throws it away. It's like a chef tasting a soup, adjusting the salt, and then forgetting the exact taste of the spoonful, rather than keeping a record of every spoonful ever tasted.
Why it matters: This ensures the robot still learns the unique details of every single brain setting, even though the memory is compressed.

3. The "Unfolding" (Restoration)

When it's time to actually update the robot's brain, FOAM takes those compressed "team averages" and expands them back out to the full size, adding the "Residual" corrections on top.

The Result: The robot gets a precise update, just like it would with the heavy accountant, but the computer only had to carry the tiny folded version.

What the Paper Claims (The Results)
The authors tested FOAM on training various robot models (from small 60-million-parameter models to large 7-billion-parameter models). Here is what they found:

Massive Memory Savings: FOAM cuts the memory needed for the optimizer by up to 90%. In total, it reduces the memory footprint of training by about 50%. This means you can train bigger models on the same hardware, or train the same models much faster.
No Performance Loss: Despite the heavy compression, the models trained with FOAM perform just as well (or sometimes even better) than models trained with the standard heavy accountant. They learn faster and reach higher accuracy.
Speed: Because it doesn't have to carry a heavy load, the training process is faster.
Versatility: It works on different types of robot architectures (like LLaMA, Qwen, and others) and can even be combined with other memory-saving tricks.

In a Nutshell
FOAM is like a magic suitcase. It allows you to pack a massive amount of information (the optimizer's memory) into a tiny space by folding it up neatly and adding a quick "correction note" to ensure nothing is lost. This lets researchers build smarter, larger AI models without needing super-expensive, massive computers.

Technical Summary: FOAM – Blocked State Folding for Memory-Efficient LLM Training

1. Problem Statement

Large Language Models (LLMs) have achieved remarkable performance through increased parameter counts and training data, but their scale introduces significant memory bottlenecks during training. The de facto optimizer, Adam (and AdamW), requires storing first- and second-order moment estimates for every parameter, resulting in an optimizer state memory overhead of twice the model size. This makes pre-training and fine-tuning large models (e.g., 7B+ parameters) memory-bound, often requiring hundreds of gigabytes of VRAM even with small batch sizes.

Existing memory-efficient solutions face trade-offs:

Parameter Freezing (e.g., LoRA): Restricts updates to low-rank subspaces, potentially limiting model expressiveness and performance.
Low-Rank Projections (e.g., GaLore, APOLLO): Compress optimizer states using Singular Value Decomposition (SVD) or random projections. While effective, SVD introduces substantial computational overhead, and random projections may degrade performance under strict memory constraints.
Block-Sharing (e.g., Adam-Mini): Shares second-moment estimates across blocks to reduce storage but often leaves first-moment overhead high or caps memory reduction at roughly 50%.

2. Methodology: FOAM

The authors propose FOAM (Folded Optimizer with Approximate Moment), a structured state-folding framework designed to compress optimizer states without sacrificing optimization fidelity. FOAM operates on the principle of blocked averaging combined with residual correction.

Core Mechanism

Given a gradient matrix $G_t \in \mathbb{R}^{m \times n}$ at time step $t$ and a fold level $l$ :

Blocked Averaging (Fold): The gradient is partitioned into blocks of $2^l$ consecutive elements. A fold operator $A^{(l)}$ computes the mean of each block, compressing the gradient dimension from $n$ to $n/2^l$ . The first- and second-order moments ( $\tilde{M}_t, \tilde{V}_t$ ) are updated based on these compressed gradients.
Unfolding: To apply updates to the full parameter space, a fold operator $E^{(l)}$ replicates the compressed entries back to the original dimensions.
Residual Correction: Directly replicating updates causes all elements within a block to share identical values, relying solely on initialization for differentiation. To recover lost information, FOAM computes a residual term $R_t = G_t - \tilde{G}_t E^{(l)}$ $R_{t} = G_{t} - \tilde{G}_{t} E^{(l)}$ , representing the error between the original gradient and the folded-unfolded version.
- The first-order moment is reconstructed as: $M_t = \tilde{M}_t E^{(l)} + R_t$ .
- The second-order moment is reconstructed as: $V_t = \tilde{V}_t E^{(l)} + R_t^2$ .
- Crucially, $R_t$ is computed on-the-fly and released immediately; it is not stored in memory.

Hybrid Application

FOAM is applied specifically to the MLP and Attention modules, while standard Adam is used for Embedding layers. A scaling coefficient $\alpha$ adjusts the learning rate ratio between FOAM and Adam modules, a common practice in memory-efficient optimizers.

3. Key Contributions

Innovative Optimizer Design: FOAM achieves up to 90% reduction in optimizer state memory overhead and approximately 50% reduction in total end-to-end memory footprint compared to Full-Adam. Unlike projection-based methods, it requires no projection matrices and supports full-parameter optimization.
Theoretical Guarantees: Under standard non-convex optimization assumptions (L-smoothness, bounded variance, bounded gradients), the authors prove that FOAM achieves the same convergence rate of $O(\frac{\log T}{\sqrt{T}})$ as vanilla Adam. The residual error ratio $\delta_l$ is bounded by 1 and does not negatively impact the primary convergence rate.
Extensive Evaluation: The method is validated across model sizes (60M to 7B), architectures (LLaMA, Qwen, GPT, DeBERTa), and tasks (pre-training on C4, fine-tuning on GLUE and MMLU).

4. Experimental Results

The paper reports the following empirical findings:

Pre-training Performance: FOAM consistently outperforms memory-efficient baselines (GaLore, APOLLO, Adam-Mini, Muon) in terms of validation perplexity (PPL) while maintaining significantly lower memory usage. For instance, on LLaMA-1.3B, FOAM-3 reduces total memory consumption by ~50% compared to Adam while achieving lower PPL.
Convergence Speed: FOAM demonstrates faster convergence, often requiring fewer training iterations to reach a target PPL compared to other memory-efficient methods.
Scalability: FOAM maintains superior performance on larger models (LLaMA-3B, 7B) and under long-sequence training regimes (up to 2048 tokens).
Fine-tuning: On GLUE and MMLU benchmarks, FOAM achieves performance comparable to or better than Full-Adam and other baselines, demonstrating its utility beyond pre-training.
Robustness: Ablation studies confirm that the residual term is critical for performance; removing it leads to slower convergence and higher final PPL. FOAM also remains robust when combined with 8-bit quantization and when integrated with other optimizers like Adam-Mini and Muon.

5. Significance and Claims

The paper claims that FOAM offers a practical, highly efficient, and optimizer-agnostic solution to the memory bottlenecks in large-scale language model training. By eliminating the need for weight freezing or expensive projection computations, FOAM allows for full-parameter optimization with drastically reduced memory overhead. The authors position FOAM as a complementary alternative to low-rank projection techniques, providing a method that theoretically matches Adam's convergence while empirically delivering superior memory efficiency and often better training performance.

The work suggests that FOAM can enable the training of larger models on existing hardware or reduce the computational costs of training current state-of-the-art models, making high-performance LLM training more accessible.