🤖 machine learning

FOAM: Blocked State Folding for Memory-Efficient LLM Training

Autores originais: Ziqing Wen, Jiahuan Wang, Ping Luo, Dongsheng Li, Tao Sun

Publicado 2026-05-14

📖 5 min de leitura🧠 Leitura aprofundada

Autores originais: Ziqing Wen, Jiahuan Wang, Ping Luo, Dongsheng Li, Tao Sun

Artigo original sob licença CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Esta é uma explicação gerada por IA do artigo abaixo. Não foi escrita nem endossada pelos autores. Para precisão técnica, consulte o artigo original. Ler aviso legal completo

Imagine que você está tentando ensinar um robô gigante e superinteligente (um Modelo de Linguagem de Grande Escala) a escrever, conversar e raciocinar. Para fazer isso, você mostra a ele uma biblioteca massiva de livros (dados de treinamento). O robô aprende cometendo erros, verificando seu trabalho e ajustando suas "configurações cerebrais" internas (parâmetros) para melhorar.

A maneira padrão de fazer esse ajuste é como a de um contador muito cuidadoso chamado Adam. Adam mantém um livro-razão detalhado para cada configuração cerebral individual, rastreando não apenas o erro atual, mas também o histórico de erros passados para decidir quanto mudar as coisas.

O Problema: O Contador é Pesado Demais
O problema é que esse "contador" (o otimizador) é incrivelmente pesado. Para um robô com bilhões de configurações cerebrais, o livro-razão do contador ocupa o dobro da memória do próprio cérebro do robô.

Analogia: Imagine tentar mover uma casa. Os móveis (o modelo) são grandes, mas o caminhão de mudança (os estados do otimizador) é ainda maior. Se você tiver apenas um caminhão pequeno (memória de computador limitada), não conseguirá mover a casa de forma alguma, não importa o quão boa seja a casa. Este é o "gargalo de memória" que impede os pesquisadores de treinar robôs maiores e mais inteligentes.

As Soluções Antigas: Cortar Cantos
Tentativas anteriores de reduzir o caminhão de mudança apresentavam falhas:

Congelando partes da casa: Você para de mover alguns móveis e ajusta apenas algumas caixas pequenas. Isso economiza espaço, mas torna a casa menos flexível (menor desempenho).
Tirando fotos borradas: Em vez de mover os móveis reais, você tira fotos de baixa resolução deles para economizar espaço. Mas calcular essas fotos consome muito tempo e energia (sobrecarga computacional).
Compartilhando livros-razão: Você faz com que diferentes cômodos compartilhem o mesmo caderno. Isso economiza espaço, mas torna os ajustes menos precisos.

A Nova Solução: FOAM (O Método Inteligente de Dobrar)
O artigo apresenta o FOAM (Otimizador Dobrado com Momento Aproximado). Pense no FOAM como um empacotador mestre que usa uma técnica inteligente de "dobragem" para caber o caminhão de mudança massivo em uma van minúscula, sem perder nenhum detalhe importante.

Veja como o FOAM funciona, usando metáforas simples:

1. A "Dobragem em Blocos" (Compressão)

Em vez de rastrear cada configuração cerebral individualmente, o FOAM as agrupa em pequenos blocos (como uma grade de 8 ou 16 configurações).

A Metáfora: Imagine que você tem uma fila longa de 100 pessoas e precisa lembrar a altura média delas. Em vez de anotar 100 números separados, você as agrupa em 10 equipes de 10. Você anota apenas um número: a altura média de cada equipe.
O Resultado: Isso reduz a memória necessária para o "livro-razão" em uma quantidade enorme (até 90% a menos).

2. A "Correção Residual" (A Rede de Segurança)

Se você usasse apenas as médias, perderia os detalhes únicos de cada pessoa. Talvez uma pessoa em uma equipe seja muito alta e outra muito baixa. A média esconde isso.

A Metáfora: O FOAM é inteligente. Depois de anotar a "média da equipe", ele calcula rapidamente a diferença (o erro) entre as pessoas reais e a média. Ele chama isso de "Resíduo".
O Truque: Ele não mantém essa diferença para sempre. Ele a calcula, usa para corrigir a atualização naquele momento específico e depois a descarta. É como um chef provando uma sopa, ajustando o sal e depois esquecendo o gosto exato daquela colherada, em vez de manter um registro de cada colherada já provada.
Por que importa: Isso garante que o robô ainda aprenda os detalhes únicos de cada configuração cerebral individual, mesmo que a memória esteja comprimida.

3. O "Desdobramento" (Restauração)

Quando é hora de atualizar realmente o cérebro do robô, o FOAM pega essas "médias de equipe" comprimidas e as expande de volta para o tamanho completo, adicionando as correções "Residuais" por cima.

O Resultado: O robô recebe uma atualização precisa, exatamente como faria com o contador pesado, mas o computador só precisou carregar a versão dobrada minúscula.

O que o Artigo Afirma (Os Resultados)
Os autores testaram o FOAM no treinamento de vários modelos de robô (de modelos pequenos de 60 milhões de parâmetros a modelos grandes de 7 bilhões de parâmetros). Eis o que descobriram:

Economia Massiva de Memória: O FOAM reduz a memória necessária para o otimizador em até 90%. No total, reduz a pegada de memória do treinamento em cerca de 50%. Isso significa que você pode treinar modelos maiores no mesmo hardware ou treinar os mesmos modelos muito mais rápido.
Sem Perda de Desempenho: Apesar da compressão pesada, os modelos treinados com FOAM performam tão bem (ou às vezes até melhor) quanto os modelos treinados com o contador pesado padrão. Eles aprendem mais rápido e alcançam maior precisão.
Velocidade: Como não precisa carregar uma carga pesada, o processo de treinamento é mais rápido.
Versatilidade: Funciona em diferentes tipos de arquiteturas de robô (como LLaMA, Qwen e outros) e pode até ser combinado com outras truques de economia de memória.

Em Resumo
O FOAM é como uma mala mágica. Ele permite que você empacote uma quantidade massiva de informações (a memória do otimizador) em um espaço minúsculo, dobrando-a cuidadosamente e adicionando uma rápida "nota de correção" para garantir que nada seja perdido. Isso permite que os pesquisadores construam modelos de IA maiores e mais inteligentes sem precisar de computadores supercaros e massivos.

Resumo Técnico: FOAM – Dobramento de Estado Bloqueado para Treinamento de LLMs Eficiente em Memória

1. Declaração do Problema

Modelos de Linguagem de Grande Escala (LLMs) alcançaram desempenho notável através do aumento de contagem de parâmetros e dados de treinamento, mas sua escala introduz gargalos significativos de memória durante o treinamento. O otimizador de facto, Adam (e AdamW), exige o armazenamento de estimativas de momentos de primeira e segunda ordem para cada parâmetro, resultando em uma sobrecarga de memória de estado do otimizador duas vezes o tamanho do modelo. Isso torna o pré-treinamento e o ajuste fino de modelos grandes (por exemplo, 7B+ parâmetros) limitados pela memória, frequentemente exigindo centenas de gigabytes de VRAM mesmo com tamanhos de lote pequenos.

Soluções existentes eficientes em memória enfrentam compensações:

Congelamento de Parâmetros (por exemplo, LoRA): Restringe atualizações a subespaços de baixo posto, potencialmente limitando a expressividade e o desempenho do modelo.
Projeções de Baixo Posto (por exemplo, GaLore, APOLLO): Comprimem estados do otimizador usando Decomposição em Valores Singulares (SVD) ou projeções aleatórias. Embora eficazes, a SVD introduz uma sobrecarga computacional substancial, e projeções aleatórias podem degradar o desempenho sob restrições estritas de memória.
Compartilhamento de Blocos (por exemplo, Adam-Mini): Compartilha estimativas de segundo momento entre blocos para reduzir o armazenamento, mas frequentemente deixa a sobrecarga de primeiro momento alta ou limita a redução de memória a aproximadamente 50%.

2. Metodologia: FOAM

Os autores propõem o FOAM (Otimizador Dobrado com Momento Aproximado), um framework estruturado de dobramento de estado projetado para comprimir estados do otimizador sem sacrificar a fidelidade de otimização. O FOAM opera no princípio da média bloqueada combinada com correção de resíduo.

Mecanismo Central

Dada uma matriz de gradiente $G_t \in \mathbb{R}^{m \times n}$ no passo de tempo $t$ e um nível de dobra $l$ :

Média Bloqueada (Dobra): O gradiente é partitionado em blocos de $2^l$ elementos consecutivos. Um operador de dobra $A^{(l)}$ calcula a média de cada bloco, comprimindo a dimensão do gradiente de $n$ para $n/2^l$ . Os momentos de primeira e segunda ordem ( $\tilde{M}_t, \tilde{V}_t$ ) são atualizados com base nesses gradientes comprimidos.
Desdobramento: Para aplicar atualizações ao espaço completo de parâmetros, um operador de dobra $E^{(l)}$ replica as entradas comprimidas de volta às dimensões originais.
Correção de Resíduo: Replicar diretamente atualizações faz com que todos os elementos dentro de um bloco compartilhem valores idênticos, dependendo exclusivamente da inicialização para diferenciação. Para recuperar informações perdidas, o FOAM calcula um termo residual $R_t = G_t - \tilde{G}_t E^{(l)}$ $R_{t} = G_{t} - \tilde{G}_{t} E^{(l)}$ , representando o erro entre o gradiente original e a versão dobrada-desdobrada.
- O momento de primeira ordem é reconstruído como: $M_t = \tilde{M}_t E^{(l)} + R_t$ .
- O momento de segunda ordem é reconstruído como: $V_t = \tilde{V}_t E^{(l)} + R_t^2$ .
- Crucialmente, $R_t$ é calculado sob demanda e liberado imediatamente; não é armazenado na memória.

Aplicação Híbrida

O FOAM é aplicado especificamente aos módulos MLP e Atenção, enquanto o Adam padrão é usado para camadas de Embedding. Um coeficiente de escala $\alpha$ ajusta a razão da taxa de aprendizado entre os módulos FOAM e Adam, uma prática comum em otimizadores eficientes em memória.

3. Contribuições Principais

Design Inovador de Otimizador: O FOAM alcança até 90% de redução na sobrecarga de memória de estado do otimizador e aproximadamente 50% de redução na pegada total de memória fim a fim em comparação com o Full-Adam. Diferentemente de métodos baseados em projeção, não requer matrizes de projeção e suporta otimização de parâmetros completos.
Garantias Teóricas: Sob suposições padrão de otimização não convexa (suavidade L, variância limitada, gradientes limitados), os autores provam que o FOAM alcança a mesma taxa de convergência de $O(\frac{\log T}{\sqrt{T}})$ que o Adam vanilla. A razão de erro residual $\delta_l$ é limitada por 1 e não impacta negativamente a taxa de convergência primária.
Avaliação Extensa: O método é validado em tamanhos de modelo (60M a 7B), arquiteturas (LLaMA, Qwen, GPT, DeBERTa) e tarefas (pré-treinamento em C4, ajuste fino em GLUE e MMLU).

4. Resultados Experimentais

O artigo relata as seguintes descobertas empíricas:

Desempenho de Pré-treinamento: O FOAM supera consistentemente as linhas de base eficientes em memória (GaLore, APOLLO, Adam-Mini, Muon) em termos de perplexidade de validação (PPL), mantendo uso de memória significativamente menor. Por exemplo, no LLaMA-1.3B, o FOAM-3 reduz o consumo total de memória em ~50% em comparação com o Adam, enquanto alcança PPL menor.
Velocidade de Convergência: O FOAM demonstra convergência mais rápida, frequentemente exigindo menos iterações de treinamento para atingir um PPL alvo em comparação com outros métodos eficientes em memória.
Escalabilidade: O FOAM mantém desempenho superior em modelos maiores (LLaMA-3B, 7B) e sob regimes de treinamento de sequência longa (até 2048 tokens).
Ajuste Fino: Nos benchmarks GLUE e MMLU, o FOAM alcança desempenho comparável ou superior ao Full-Adam e outras linhas de base, demonstrando sua utilidade além do pré-treinamento.
Robustez: Estudos de ablação confirmam que o termo residual é crítico para o desempenho; removê-lo leva a convergência mais lenta e PPL final mais alto. O FOAM também permanece robusto quando combinado com quantização de 8 bits e quando integrado com outros otimizadores como Adam-Mini e Muon.

5. Significado e Alegações

O artigo alega que o FOAM oferece uma solução prática, altamente eficiente e agnóstica ao otimizador para os gargalos de memória no treinamento de modelos de linguagem em grande escala. Ao eliminar a necessidade de congelamento de pesos ou computações de projeção caras, o FOAM permite otimização de parâmetros completos com sobrecarga de memória drasticamente reduzida. Os autores posicionam o FOAM como uma alternativa complementar às técnicas de projeção de baixo posto, fornecendo um método que teoricamente iguala a convergência do Adam enquanto empiricamente entrega eficiência superior de memória e frequentemente melhor desempenho de treinamento.

O trabalho sugere que o FOAM pode permitir o treinamento de modelos maiores em hardware existente ou reduzir os custos computacionais do treinamento dos modelos atuais de última geração, tornando o treinamento de LLMs de alto desempenho mais acessível.