Resumen Técnico: La Restauración Selectiva de Capas Recupera la Diversidad

Declaración del Problema

Las técnicas de post-entrenamiento, como el ajuste de instrucciones y el Aprendizaje por Refuerzo a partir de la Retroalimentación Humana (RLHF), mejoran significativamente las capacidades de seguimiento de instrucciones y la utilidad de los Grandes Modelos de Lenguaje (LLMs). Sin embargo, estos métodos suelen inducir un colapso de modo (mode collapse), un fenómeno donde el modelo asigna una masa de probabilidad desproporcionadamente alta a un rango estrecho de respuestas, suprimiendo alternativas diversas pero válidas. Esto conduce a salidas repetitivas en entornos de código abierto (por ejemplo, escritura creativa, QA de código abierto y razonamiento), limitando la utilidad del modelo en aplicaciones que valoran respuestas distintas.

Las estrategias de mitigación existentes se dividen en tres categorías, cada una con limitaciones:

1. Métodos basados en la decodificación (por ejemplo, escalado de temperatura, muestreo min-p) aumentan la estocasticidad pero no pueden recuperar la diversidad perdida debido a los cambios en los parámetros durante el post-entrenamiento.
2. Métodos basados en prompts (por ejemplo, solicitar múltiples respuestas) aumentan los costos de inferencia y son sensibles a la redacción del prompt.
3. Métodos basados en el entrenamiento (por ejemplo, rediseño de objetivos) requieren un reentrenamiento y una recolección de datos costosos.

Los autores buscan un método que recupere la diversidad de manera efectiva, sea computacionalmente eficiente durante la adaptación e introduzca ningún costo adicional durante la inferencia.

Metodología

El enfoque propuesto, Restauración Selectiva de Capas (SLR, por sus siglas en inglés), es una intervención en el espacio de pesos sin necesidad de entrenamiento basada en dos observaciones empíricas:

1. Las capas de los LLMs exhiben roles funcionales distintos (por ejemplo, las asociaciones de hechos suelen estar localizadas en capas de profundidad media).
2. El post-entrenamiento principalmente orienta qué comportamientos pre-entrenados se expresan, suprimiendo a menudo otros (olvido catastrófico), en lugar de adquirir conocimiento completamente nuevo.

Los autores hipotetizan que el colapso de modo está localizado en capas específicas y que restaurar un rango de capas cuidadosamente elegido en un modelo post-entrenado ($M_{post}$) a sus pesos pre-entrenados ($M_{pre}$) puede recuperar la diversidad manteniendo la calidad.

El Marco de Trabajo SLR

Dado un modelo con $N$ capas, SLR construye un modelo híbrido $M_{i:j}$ reemplazando un intervalo contiguo de capas $[i, j]$ en $M_{post}$ con las capas correspondientes de $M_{pre}$. El objetivo es encontrar el intervalo $(i, j)$ que maximice la diversidad ($D$) sujeto a un umbral mínimo de calidad ($Q \ge q_{min}$).

La Tarea Proxy CRC

La optimización directa del intervalo de restauración en tareas descendentes complejas es computacionalmente costosa y difícil de evaluar debido a la subjetividad de la "calidad". Para abordar esto, los autores introducen el Carácter Aleatorio Restringido (CRC, por sus siglas en inglés) como una tarea proxy.

• Mecanismo: CRC implica generar caracteres aleatorios (dígitos o letras) dentro de un conjunto válido estrictamente definido (por ejemplo, "Genera un número entero aleatorio en [0, 5]").
• Métricas:
  • Calidad ($Q_{CRC}$): La masa de probabilidad asignada a salidas válidas (mide la adherencia a la instrucción).
  • Diversidad ($D_{CRC}$): La entropía de la distribución sobre las salidas válidas (mide la variedad).
• Proceso de Selección: Los autores realizan una búsqueda de cuadrícula (grid search) sobre todos los posibles intervalos de capas contiguos utilizando CRC. Seleccionan el intervalo que maximiza $D_{CRC}$ mientras cumple con $Q_{CRC} \ge 0.9 \times Q_{CRC}(M_{post})$. Este intervalo seleccionado se aplica luego a las tareas descendentes sin ajustes adicionales.

Contribuciones Clave

1. Restauración Selectiva de Capas (SLR): Un método sin entrenamiento que restaura un intervalo contiguo específico de capas en un LLM post-entrenado a sus parámetros pre-entrenados. Preserva la arquitectura original y el conteo de parámetros, sin incurrir en costos de inferencia adicionales. Se demuestra que es complementario a los métodos de diversidad basados en decodificación y prompts.
2. Carácter Aleatorio Restringido (CRC): Una tarea proxy con un conjunto de validez explícito y un objetivo de diversidad natural. CRC permite una exploración eficiente y de bajo costo de los intervalos de restauración y proporciona evidencia empírica para la hipótesis de localización de capas.
3. Validación Empírica: Evaluación exhaustiva a través de tres familias de modelos (Llama, Qwen, Gemma) y tres tareas descendentes (escritura creativa, QA de código abierto, razonamiento de múltiples pasos).

Resultados Experimentales

Los autores evaluaron SLR frente al modelo base post-entrenado y un modelo base de "Model Soup" (promedio ponderado de parámetros de modelos pre-entrenados y post-entrenados).

• Escritura Creativa: SLR mejoró consistentemente la diversidad (aumento promedio del +42.5% en la disimilitud basada en embeddings) con una pérdida de calidad mínima (promedio de -2.1% en las puntuaciones del juez de LLM). En contraste, Model Soup mostró ganancias de diversidad menores (+27.3%) y una degradación de calidad más significativa (-6.0%).
• QA de Código Abierto: SLR aumentó sustancialmente la diversidad de las respuestas, mostrando un incremento del +112.7% en la entropía y un incremento del +169.5% en la cobertura (fracción de respuestas correctas únicas), manteniendo una alta precisión (caída promedio de -1.8%).
• Razonamiento (GSM8K): En problemas donde los modelos post-entrenados fallaban bajo decodificación codiciosa (greedy decoding), SLR mejoró el Pass@k en todos los presupuestos de muestreo ($k \in \{1, \dots, 64\}$), indicando un desplazamiento de la masa de probabilidad hacia caminos de razonamiento diversos pero correctos.
• Componibilidad: SLR siguió siendo efectivo cuando se combinó con decodificación de temperatura más alta ($T=1.5$) y muestreo verbalizado (métodos basados en prompts), lo que sugiere que SLR modifica la distribución subyacente del modelo mientras que otros métodos manipulan el proceso de muestreo o la entrada.
• Estudios de Ablación: La restauración de capas guiada por la tarea proxy CRC produjo resultados superiores en comparación con la restauración de intervalos arbitrarios (por ejemplo, las capas más tempranas o más tardías). La restauración de capas arbitrarias a menudo no logró mejorar la diversidad o degradó la calidad, confirmando que la selección del intervalo específico es crítica.

Significado y Reivindicaciones

El artículo afirma que SLR ofrece un camino simple pero efectivo para mitigar el colapso de modo sin la sobrecarga computacional del reentrenamiento o los costos de inferencia de los prompts complejos. Al demostrar que la diversidad puede recuperarse mediante intervenciones dirigidas en el espacio de pesos, el trabajo respalda una visión modular de los LLMs post-entrenados:

1. Las funcionalidades y comportamientos (incluyendo el colapso de modo) están localizados en capas específicas en lugar de estar distribuidos uniformemente.
2. El post-entrenamiento principalmente orienta la expresión de los comportamientos pre-entrenados en lugar de alterar fundamentalmente la base de conocimiento central del modelo.

Los autores concluyen que las intervenciones de pesos de grano fino como SLR proporcionan una vía prometedora para controlar la relación calidad-diversidad en los LLMs, permitiendo la recuperación de propiedades pre-entrenadas deseables (como la diversidad) mientras se retienen los beneficios del post-entrenamiento (como la adherencia a las instrucciones). El artículo se mantiene modesto respecto a la escalabilidad, señalando que los resultados se limitan actualmente a modelos de aproximadamente ~7–9B de parámetros y que se necesita trabajo futuro para determinar si estos hallazgos se mantienen en escalas mayores o con intervenciones más granulares (por ejemplo, restauración de subcapas).