🤖 machine learning

FOAM: Blocked State Folding for Memory-Efficient LLM Training

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

Publicado 2026-05-14

📖 5 min de lectura🧠 Análisis profundo

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

Artículo original bajo licencia CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Esta es una explicación generada por IA del artículo a continuación. No ha sido escrita ni avalada por los autores. Para mayor precisión técnica, consulte el artículo original. Leer descargo de responsabilidad completo

Imagina que estás intentando enseñar a un robot gigante y superinteligente (un Modelo de Lenguaje Grande) a escribir, chatear y razonar. Para lograrlo, le muestras una biblioteca masiva de libros (datos de entrenamiento). El robot aprende cometiendo errores, revisando su trabajo y ajustando su configuración interna de "cerebro" (parámetros) para mejorar.

La forma estándar de realizar este ajuste es como un contable muy cuidadoso llamado Adam. Adam lleva un libro de cuentas detallado para cada configuración individual del cerebro, rastreando no solo el error actual, sino también la historia de errores pasados para decidir cuánto cambiar las cosas.

El Problema: El Contable es Demasiado Pesado
El problema es que este "contable" (el optimizador) es increíblemente pesado. Para un robot con miles de millones de configuraciones cerebrales, el libro de cuentas del contable ocupa el doble de memoria que el propio cerebro del robot.

Analogía: Imagina intentar mover una casa. Los muebles (el modelo) son grandes, pero el camión de mudanza (los estados del optimizador) es aún más grande. Si solo tienes un camión pequeño (memoria de computadora limitada), no puedes mover la casa en absoluto, sin importar lo buena que sea la casa. Esta es la "cuello de botella de memoria" que impide a los investigadores entrenar robots más grandes e inteligentes.

Las Soluciones Antiguas: Recortar Esquinas
Los intentos anteriores de reducir el camión de mudanza tenían defectos:

Congelar partes de la casa: Dejas de mover algunos muebles y solo ajustas unas pocas cajas pequeñas. Esto ahorra espacio pero hace que la casa sea menos flexible (menor rendimiento).
Tomar fotos borrosas: En lugar de mover los muebles reales, tomas fotos de baja resolución de ellos para ahorrar espacio. Pero calcular estas fotos requiere mucho tiempo y energía (sobrecarga computacional).
Compartir libros de cuentas: Haces que diferentes habitaciones compartan la misma libreta. Esto ahorra espacio pero hace que los ajustes sean menos precisos.

La Nueva Solución: FOAM (El Método de Plegado Inteligente)
El artículo introduce FOAM (Optimizador Plegado con Momento Aproximado). Piensa en FOAM como un empacador maestro que utiliza una técnica de "plegado" inteligente para meter el camión de mudanza masivo en una camioneta pequeña, sin perder ningún detalle importante.

Así funciona FOAM, usando metáforas simples:

1. El "Plegado por Bloques" (Compresión)

En lugar de rastrear cada configuración individual del cerebro, FOAM las agrupa en pequeños bloques (como una cuadrícula de 8 o 16 configuraciones).

La Metáfora: Imagina que tienes una fila larga de 100 personas y necesitas recordar su altura promedio. En lugar de escribir 100 números separados, los agrupas en 10 equipos de 10. Escribes solo un número: la altura promedio de cada equipo.
El Resultado: Esto reduce la memoria necesaria para el "libro de cuentas" en una cantidad enorme (hasta un 90% menos).

2. La "Corrección Residual" (La Red de Seguridad)

Si solo usaras los promedios, perderías los detalles únicos de cada persona. Quizás una persona en un equipo es muy alta y otra es muy baja. El promedio oculta eso.

La Metáfora: FOAM es inteligente. Después de escribir el "promedio del equipo", calcula rápidamente la diferencia (el error) entre las personas reales y el promedio. Llama a esto "Residual".
El Truco: No mantiene esta diferencia para siempre. La calcula, la usa para corregir la actualización en ese momento específico y luego la desecha. Es como un chef que prueba una sopa, ajusta la sal y luego olvida el sabor exacto de esa cucharada, en lugar de guardar un registro de cada cucharada probada.
Por qué importa: Esto asegura que el robot aún aprenda los detalles únicos de cada configuración individual del cerebro, aunque la memoria esté comprimida.

3. El "Desplegado" (Restauración)

Cuando es hora de actualizar realmente el cerebro del robot, FOAM toma esos "promedios del equipo" comprimidos y los expande de nuevo al tamaño completo, agregando las correcciones "Residual" encima.

El Resultado: El robot recibe una actualización precisa, exactamente como lo haría con el contable pesado, pero la computadora solo tuvo que llevar la versión plegada y pequeña.

Lo que Afirma el Artículo (Los Resultados)
Los autores probaron FOAM entrenando varios modelos de robots (desde modelos pequeños de 60 millones de parámetros hasta modelos grandes de 7 mil millones de parámetros). Esto es lo que encontraron:

Ahorro Masivo de Memoria: FOAM reduce la memoria necesaria para el optimizador hasta en un 90%. En total, reduce la huella de memoria del entrenamiento en aproximadamente un 50%. Esto significa que puedes entrenar modelos más grandes en el mismo hardware, o entrenar los mismos modelos mucho más rápido.
Sin Pérdida de Rendimiento: A pesar de la compresión pesada, los modelos entrenados con FOAM funcionan igual de bien (o incluso mejor en ocasiones) que los modelos entrenados con el contable pesado estándar. Aprenden más rápido y alcanzan una mayor precisión.
Velocidad: Como no tiene que llevar una carga pesada, el proceso de entrenamiento es más rápido.
Versatilidad: Funciona en diferentes tipos de arquitecturas de robots (como LLaMA, Qwen y otros) e incluso puede combinarse con otros trucos de ahorro de memoria.

En Resumen
FOAM es como una maleta mágica. Te permite guardar una cantidad masiva de información (la memoria del optimizador) en un espacio pequeño plegándola ordenadamente y añadiendo una rápida "nota de corrección" para asegurar que nada se pierda. Esto permite a los investigadores construir modelos de IA más inteligentes y grandes sin necesidad de computadoras masivas y supercostosas.

Resumen Técnico: FOAM – Plegado de Estado Bloqueado para Entrenamiento Eficiente en Memoria de LLM

1. Declaración del Problema

Los Modelos de Lenguaje Grandes (LLM) han logrado un rendimiento notable mediante el aumento de la cantidad de parámetros y de los datos de entrenamiento, pero su escala introduce cuellos de botella significativos en el uso de memoria durante el entrenamiento. El optimizador de facto, Adam (y AdamW), requiere almacenar estimaciones de momentos de primer y segundo orden para cada parámetro, lo que resulta en una sobrecarga de memoria del estado del optimizador equivalente al doble del tamaño del modelo. Esto hace que el preentrenamiento y el ajuste fino de modelos grandes (por ejemplo, con más de 7 mil millones de parámetros) estén limitados por la memoria, requiriendo a menudo cientos de gigabytes de VRAM incluso con tamaños de lote pequeños.

Las soluciones existentes para la eficiencia en memoria enfrentan compensaciones:

Congelamiento de Parámetros (por ejemplo, LoRA): Restringe las actualizaciones a subespacios de rango bajo, limitando potencialmente la expresividad y el rendimiento del modelo.
Proyecciones de Rango Bajo (por ejemplo, GaLore, APOLLO): Comprimen los estados del optimizador utilizando Descomposición en Valores Singulares (SVD) o proyecciones aleatorias. Aunque son efectivas, la SVD introduce una sobrecarga computacional sustancial, y las proyecciones aleatorias pueden degradar el rendimiento bajo restricciones estrictas de memoria.
Compartición de Bloques (por ejemplo, Adam-Mini): Comparte las estimaciones del segundo momento entre bloques para reducir el almacenamiento, pero a menudo deja alta la sobrecarga del primer momento o limita la reducción de memoria a aproximadamente un 50%.

2. Metodología: FOAM

Los autores proponen FOAM (Optimizador Plegado con Momento Aproximado), un marco estructurado de plegado de estados diseñado para comprimir los estados del optimizador sin sacrificar la fidelidad de la optimización. FOAM opera bajo el principio de promedio bloqueado combinado con corrección de residuos.

Mecanismo Central

Dada una matriz de gradiente $G_t \in \mathbb{R}^{m \times n}$ en el paso de tiempo $t$ y un nivel de plegado $l$ :

Promedio Bloqueado (Plegado): El gradiente se divide en bloques de $2^l$ elementos consecutivos. Un operador de plegado $A^{(l)}$ calcula la media de cada bloque, comprimiendo la dimensión del gradiente de $n$ a $n/2^l$ . Los momentos de primer y segundo orden ( $\tilde{M}_t, \tilde{V}_t$ ) se actualizan basándose en estos gradientes comprimidos.
Desplegado: Para aplicar actualizaciones al espacio completo de parámetros, un operador de plegado $E^{(l)}$ replica las entradas comprimidas de vuelta a las dimensiones originales.
Corrección de Residuos: Replicar directamente las actualizaciones hace que todos los elementos dentro de un bloque compartan valores idénticos, dependiendo únicamente de la inicialización para la diferenciación. Para recuperar la información perdida, FOAM calcula un término residual $R_t = G_t - \tilde{G}_t E^{(l)}$ $R_{t} = G_{t} - \tilde{G}_{t} E^{(l)}$ , que representa el error entre el gradiente original y la versión plegada-desplegada.
- El momento de primer orden se reconstruye como: $M_t = \tilde{M}_t E^{(l)} + R_t$ .
- El momento de segundo orden se reconstruye como: $V_t = \tilde{V}_t E^{(l)} + R_t^2$ .
- Crucialmente, $R_t$ se calcula sobre la marcha y se libera inmediatamente; no se almacena en memoria.

Aplicación Híbrida

FOAM se aplica específicamente a los módulos MLP y de Atención, mientras que Adam estándar se utiliza para las capas de Embedding. Un coeficiente de escalado $\alpha$ ajusta la relación de la tasa de aprendizaje entre los módulos FOAM y Adam, una práctica común en optimizadores eficientes en memoria.

3. Contribuciones Clave

Diseño Innovador de Optimizador: FOAM logra una reducción de hasta un 90% en la sobrecarga de memoria del estado del optimizador y una reducción de aproximadamente un 50% en la huella total de memoria de extremo a extremo en comparación con Full-Adam. A diferencia de los métodos basados en proyección, no requiere matrices de proyección y soporta la optimización de todos los parámetros.
Garantías Teóricas: Bajo supuestos estándar de optimización no convexa (suavidad L, varianza acotada, gradientes acotados), los autores demuestran que FOAM alcanza la misma tasa de convergencia de $O(\frac{\log T}{\sqrt{T}})$ que Adam vanilla. La relación de error residual $\delta_l$ está acotada por 1 y no impacta negativamente la tasa de convergencia principal.
Evaluación Exhaustiva: El método se valida en tamaños de modelo (de 60M a 7B), arquitecturas (LLaMA, Qwen, GPT, DeBERTa) y tareas (preentrenamiento en C4, ajuste fino en GLUE y MMLU).

4. Resultados Experimentales

El artículo reporta los siguientes hallazgos empíricos:

Rendimiento de Preentrenamiento: FOAM supera consistentemente a las líneas base eficientes en memoria (GaLore, APOLLO, Adam-Mini, Muon) en términos de perplejidad de validación (PPL) mientras mantiene un uso de memoria significativamente menor. Por ejemplo, en LLaMA-1.3B, FOAM-3 reduce el consumo total de memoria en aproximadamente un 50% en comparación con Adam, logrando al mismo tiempo un PPL más bajo.
Velocidad de Convergencia: FOAM demuestra una convergencia más rápida, requiriendo a menudo menos iteraciones de entrenamiento para alcanzar un PPL objetivo en comparación con otros métodos eficientes en memoria.
Escalabilidad: FOAM mantiene un rendimiento superior en modelos más grandes (LLaMA-3B, 7B) y bajo regímenes de entrenamiento de secuencias largas (hasta 2048 tokens).
Ajuste Fino: En las pruebas de GLUE y MMLU, FOAM logra un rendimiento comparable o superior al de Full-Adam y otras líneas base, demostrando su utilidad más allá del preentrenamiento.
Robustez: Los estudios de ablación confirman que el término residual es crítico para el rendimiento; su eliminación conduce a una convergencia más lenta y un PPL final más alto. FOAM también permanece robusto cuando se combina con cuantización de 8 bits y cuando se integra con otros optimizadores como Adam-Mini y Muon.

5. Significado y Afirmaciones

El artículo afirma que FOAM ofrece una solución práctica, altamente eficiente y agnóstica al optimizador para los cuellos de botella de memoria en el entrenamiento de modelos de lenguaje a gran escala. Al eliminar la necesidad de congelamiento de pesos o cálculos de proyección costosos, FOAM permite la optimización de todos los parámetros con una sobrecarga de memoria drásticamente reducida. Los autores posicionan a FOAM como una alternativa complementaria a las técnicas de proyección de rango bajo, proporcionando un método que teóricamente iguala la convergencia de Adam mientras que empíricamente ofrece una eficiencia de memoria superior y, a menudo, un mejor rendimiento de entrenamiento.

El trabajo sugiere que FOAM puede permitir el entrenamiento de modelos más grandes en hardware existente o reducir los costos computacionales del entrenamiento de los modelos actuales de última generación, haciendo que el entrenamiento de LLM de alto rendimiento sea más accesible.