Resumen Técnico: Consistencia Optimizador-Modelo

Enunciado del Problema

En el paradigma de entrenamiento de Modelos de Lenguaje Grandes (LLM), los modelos pasan por una etapa de preentrenamiento en corpus masivos, seguida de una etapa de Ajuste Fino Supervisado (SFT) en conjuntos de datos más pequeños y de alta calidad para adquirir habilidades específicas. Un desafío crítico en el SFT es la compensación entre aprendizaje y olvido: el modelo debe aprender nuevo conocimiento específico de la tarea mientras preserva las capacidades generales de lenguaje adquiridas durante el preentrenamiento.

Aunque se emplean diversos optimizadores (por ejemplo, AdamW, Muon) y métodos eficientes en parámetros (por ejemplo, LoRA), sigue sin estar claro qué combinación produce el equilibrio óptimo entre aprender nuevas tareas y minimizar el olvido catastrófico. La literatura existente ofrece intuiciones contradictorias; por ejemplo, algunos sugieren que LoRA olvida inherentemente menos debido a su naturaleza de bajo rango, mientras que otros enfatizan la importancia de la elección del optimizador. Este artículo investiga si la elección del optimizador en el SFT debe desacoplarse o alinearse con el optimizador utilizado en el preentrenamiento.

Metodología

Los autores emplean una combinación de experimentos empíricos y análisis teórico para investigar la relación entre la elección del optimizador y el rendimiento del modelo durante el SFT.

Experimentos Empíricos

1. Modelos y Conjuntos de Datos: Los experimentos se realizaron en GPT-2-small (preentrenado desde cero) y Llama-2-7B. El preentrenamiento utilizó el conjunto de datos OpenWebText. El SFT se realizó en tres tareas distintas: MetaMathQA (razonamiento matemático), Alpaca (seguimiento de instrucciones) y Magicoder (generación de código).
2. Optimizadores Comparados: El estudio compara AdamW, Muon (un optimizador de estructura matricial) y SGD. También evalúa LoRA (Adaptación de Bajo Rango) con varios rangos ($r=4, 8, 32, 64, 128, 256$).
3. Métrica de Evaluación: Los autores definen una frontera de Pareto para visualizar la compensación entre aprendizaje y olvido.
   • Olvido: Medido por la pérdida de validación en el conjunto de datos de preentrenamiento (OpenWebText) o la caída en las pruebas de capacidad general. Una pérdida menor indica menos olvido.
   • Aprendizaje: Medido por la pérdida de validación en la nueva tarea de SFT o la precisión específica de la tarea. Una pérdida menor/precisión mayor indica un mejor aprendizaje.
   • Protocolo: Para cada algoritmo de entrenamiento, se realizó una búsqueda en cuadrícula sobre las tasas de aprendizaje para identificar la frontera de Pareto, asegurando una comparación justa de los límites del algoritmo en lugar de una única configuración de hiperparámetros.
4. Análisis Sintético: Para comprender por qué ciertos optimizadores se comportan de manera diferente, los autores realizaron un experimento sintético utilizando un corpus de lenguaje corrupto. Mezclaron el orden de los tokens en secuencias limpias para destruir los patrones lingüísticos, forzando al modelo a depender de la memorización mecánica. Esto les permitió distinguir entre la capacidad de un optimizador para aprender patrones y su tendencia a memorizar datos.

Análisis Teórico

El artículo proporciona un marco teórico basado en normas inducidas por matrices (siguiendo a Pethick et al., 2025).

• Optimizador como Regularizador: Los autores postulan que los optimizadores no solo convergen a una solución, sino que moldean los pesos y las activaciones del modelo. Específicamente, un optimizador derivado de una norma $\|\cdot\|_{\alpha, \beta}$ regulariza implícitamente las activaciones de entrada mediante $\|\cdot\|_\alpha$ y las activaciones de salida mediante $\|\cdot\|_\beta$.
  • AdamW (derivado de $A_{1,\infty}$) regulariza las entradas mediante $\|\cdot\|_1$, lo que conduce a activaciones dispersas.
  • Muon (derivado de $A_{2,2}$) regulariza las entradas mediante $\|\cdot\|_2$, lo que conduce a activaciones más densas.
• Aproximación del Olvido: Utilizando una expansión de Taylor, el aumento en la pérdida de preentrenamiento (olvido) después del SFT se aproxima como proporcional a la norma cuadrática esperada de la actualización de pesos aplicada a las entradas de preentrenamiento ($\mathbb{E}[\|\Delta W x\|_2^2]$).
• Hipótesis de Consistencia: La teoría sugiere que, para minimizar el olvido, la actualización de SFT ($\Delta W$) debe alinearse con la curvatura del paisaje de pérdida establecido por el preentrenamiento. Esta alineación se logra cuando el optimizador de SFT pertenece a la misma familia (mismos $\alpha, \beta$) que el optimizador de preentrenamiento.

Resultados Clave

1. Consistencia Optimizador-Modelo

El hallazgo central es la Consistencia Optimizador-Modelo: el ajuste fino completo con el mismo optimizador (o familia de optimizadores) utilizado en el preentrenamiento logra consistentemente la mejor compensación entre aprendizaje y olvido.

• Preentrenamiento con AdamW: Cuando AdamW se utiliza para el preentrenamiento, el ajuste fino completo con AdamW supera al ajuste fino completo con Muon o SGD, y también supera a LoRA en todas las tareas (matemáticas, código, instrucciones).
• Preentrenamiento con Muon: De manera similar, cuando Muon se utiliza para el preentrenamiento, el ajuste fino completo con Muon produce la mejor compensación en comparación con el uso de AdamW o SGD para el SFT.
• Comparación con LoRA: Contrario a la intuición de que LoRA olvida menos, el análisis de la frontera de Pareto muestra que para cualquier punto de control de LoRA, existe un punto de control de ajuste fino completo (utilizando el mismo optimizador) que logra tanto un mejor aprendizaje (pérdida de tarea menor) como menos olvido (pérdida de preentrenamiento menor). Esto se mantiene cierto cuando se considera un barrido adecuado de la tasa de aprendizaje.

2. Rendimiento de Muon en Tareas de Razonamiento

Aunque Muon a menudo produce puntos de control preentrenados superiores (pérdida de preentrenamiento menor), su rendimiento se degrada en relación con AdamW cuando se ajusta fino en tareas de razonamiento (por ejemplo, MetaMathQA).

• Memorización vs. Aprendizaje de Patrones: El experimento sintético reveló que Muon tiende a memorizar datos corruptos (sin patrones) más rápido que AdamW o SGD.
• Implicación: En el SFT, donde los datos son escasos, la tendencia de Muon a priorizar la memorización mecánica sobre la adquisición de patrones puede obstaculizar el aprendizaje de habilidades de razonamiento complejas, lo que lleva a un rendimiento final peor en tareas que requieren generalización, a pesar de sus métricas sólidas de preentrenamiento.

3. El Papel de las Tasas de Aprendizaje

El artículo destaca que la compensación entre aprendizaje y olvido es altamente sensible a la tasa de aprendizaje. Estudios anteriores que sugerían que LoRA olvida menos a menudo fijaban la tasa de aprendizaje a la óptima para la tarea. Cuando se realiza un barrido de la tasa de aprendizaje, el ajuste fino completo puede lograr una compensación superior simplemente reduciendo la tasa de aprendizaje, una estrategia que LoRA no puede replicar tan efectivamente en este contexto específico.

Significado y Afirmaciones

El artículo afirma establecer la Consistencia Optimizador-Modelo como un principio fundamental para el entrenamiento de LLM:

1. Métrica Unificada: La compensación entre aprendizaje y olvido en toda la tubería de preentrenamiento-SFT es una métrica más completa para evaluar un optimizador que el rendimiento de preentrenamiento por sí solo.
2. Orientación Práctica: Para minimizar el olvido mientras se maximiza el aprendizaje de la tarea, los practicantes deben emplear la misma familia de optimizadores para el SFT que la utilizada para el preentrenamiento.
3. Reevaluación de Muon: Aunque Muon es un optimizador de preentrenamiento potente, su mecanismo de actualización (rango completo, regularización $L_2$) puede hacerlo propenso a la sobre-memorización en regímenes de SFT con pocos datos, particularmente para tareas de razonamiento. Esto sugiere que podría ser necesaria una estrategia de ajuste fino dedicada, potencialmente diferente, para modelos preentrenados con Muon para desbloquear todo su potencial.
4. Perspectiva Teórica: El trabajo cierra la brecha entre el diseño del optimizador (normas) y el comportamiento del modelo (dispersión de activaciones), ofreciendo una explicación teórica de por qué importa la consistencia en la elección del optimizador.

Los autores reconocen limitaciones, señalando que los experimentos se realizaron principalmente en modelos más pequeños (GPT-2, Llama-2-7B) debido a restricciones computacionales y que las pruebas teóricas dependen de suposiciones simplificadas (por ejemplo, Hessianos bloque-diagonales). Sugieren que el trabajo futuro debería escalar estos hallazgos a modelos más grandes y explorar optimizadores diseñados específicamente para mitigar las tendencias de memorización de los optimizadores de estructura matricial durante el SFT.