From Weights to Features: SAE-Guided Activation Regularization for LLM Continual Learning

Resumen Técnico: De los pesos a las características: Regularización de la activación guiada por SAE para el aprendizaje continuo de LLM

Planteamiento del problema
El aprendizaje continuo (CL) en los Modelos de Lenguaje de Gran Escala (LLMs) enfrenta el desafío del olvido catastrófico, donde el ajuste fino secuencial sobrescribe las representaciones necesarias para tareas previas. Aunque los métodos basados en la regularización como la Consolidación Elástica de Pesos (EWC) son un estándar para modelos más pequeños, su rendimiento es significativamente inferior en los LLMs. Los autores argumentan que este fallo es estructural, derivado de la naturaleza "polisemántica" de los LLMs. Debido a la hipótesis de la superposición, los neuronas individuales y sus pesos asociados codifican múltiples conceptos no relacionados simultáneamente. En consecuencia, las estimaciones de importancia en el espacio de pesos (por ejemplo, la información de Fisher) son demasiado gruesas; proteger un peso para preservar el Concepto A restringe inadvertidamente al Concepto B, creando una penalización no selectiva que sofoca la plasticidad. Además, los métodos existentes a menudo no logran imponer explícitamente la plasticidad requerida para aprender nuevas tareas, lo que conduce a la "intransigencia", donde el modelo retiene el conocimiento antiguo pero no puede adaptarse.

Metodología
El artículo propone la regularización de la activación guiada por SAE, un método que desplaza el objetivo de la regularización desde el espacio de parámetros entrelazados hacia un espacio de activación más selectivo utilizando Autoencoders Dispersos (SAEs) pre-entrenados.

1. Cambio de sistema de coordenadas: En lugar de penalizar la deriva de los pesos, el método utiliza un SAE pre-entrenado y congelado como un diccionario de características monosemánticas fijas. Los SAE descomponen las activaciones densas del modelo en direcciones de características dispersas que suelen ser más interpretables y menos entrelazadas que las neuronas individuales.
2. Formulación de optimización con restricciones: El objetivo de entrenamiento se deriva de un problema de optimización con restricciones en lugar de un diseño de penalización ad-hoc. El modelo debe minimizar la pérdida de la tarea actual sujeto a dos restricciones:
   • Estabilidad: Las características irrelevantes para la tarea (baja relevancia) no deben derivar más allá de un presupuesto específico ($\varepsilon_s$).
   • Plasticidad: Las características relevantes para la tarea (alta relevancia) deben derivar lo suficiente ($\ge \varepsilon_p$) para evitar soluciones degeneradas donde el modelo no aprende nada.
3. Función de pérdida: Utilizando la relajación de Lagrange y una penalización de bisel cuadrado (squared-hinge penalty), los autores derivan una pérdida de entrenamiento que comprende una pérdida de protección (que penaliza la deriva excesiva en las características protegidas) y una pérdida de guía (que penaliza la deriva insuficiente en las características relevantes). Se demuestra que EWC es un caso especial de este marco donde el espacio de representación son los pesos y la restricción de plasticidad está ausente.
4. Construcción de la máscara: Para cada nueva tarea, se construye una máscara de relevancia continua ($M_t$) utilizando datos de la tarea actual pasados a través del modelo base congelado y el SAE. Esta máscara identifica las características relevantes mediante la magnitud de la activación y propaga la relevancia a las características vecinas usando k-NN en los vectores del decodificador. Crucialmente, una vez construida la máscara, todos los datos de la tarea se descartan; solo se retiene la máscara compacta.
5. Entrenamiento: Durante el entrenamiento, el modelo actual (con LoRA) y el modelo base congelado son codificados por el SAE. La deriva de las características se pondera mediante la máscara para calcular las pérdidas de protección y de guía.

Contribuciones clave

• Regularización guiada por SAE: Un método práctico de CL que regulariza en el espacio de activación utilizando características de SAE. No requiere datos de tareas previas, no almacena activaciones, no requiere parámetros específicos de la tarea, ni enrutamiento en el tiempo de inferencia tras la construcción de la máscara.
• Derivación fundamentada: El compromiso entre estabilidad y plasticidad se formula como un problema de optimización con restricciones en el espacio de características de SAE, produciendo una pérdida de entrenamiento de bisel cuadrado mediante la relajación de Lagrange.
• Evidencia empírica de la polisemanticidad: El artículo proporciona evidencia mecánica de que la protección en el espacio de pesos es no selectiva. Los experimentos muestran que las representaciones relevantes para la tarea son linealmente separables en el espacio de características de SAE (AUC ~0.88) pero indistinguibles del azar en el espacio de pesos/neuronas (AUC ~0.50). Además, proteger los pesos para una tarea restringe las características para una tarea subsecuente a una tasa del 91–96% de la tarea original, mientras que la protección en el espacio de características reduce esta restricción colateral a un 43–61%.
• Escalabilidad: Debido a la menor dimensionalidad del espacio de características en comparación con el espacio de parámetros, el almacenamiento por tarea se reduce de gigabytes (anclas de pesos) a sub-megabytes (máscaras de características).

Resultados
El método fue evaluado en TRACE-5000 (8 tareas de dominios cruzados) y MedCL (10 tareas biomédicas) utilizando Gemma-2 9B-it con LoRA.

• Rendimiento: En TRACE, la regularización guiada por SAE logró el mejor desempeño entre los enfoques no arquitectónicos (OP = 0.545), superando a los regularizadores de espacio de pesos (EWC: 0.447, SI: 0.513), a los métodos de proyección de gradiente y a los baselines de replay (incluso con el 10% de los datos). En MedCL, lideró todos los métodos que no retienen ejemplos de tareas previas (OP = 0.510).
• Compromiso estabilidad-plasticidad: A diferencia de EWC, que logró una alta retención (ARA = 0.987) sacrificando severamente la plasticidad (0.453), el método propuesto mantuvo una alta plasticidad (0.567, comparable al ajuste fino sin protección) mientras retenía el 95% del rendimiento aprendido.
• Eficiencia: El método añade un overhead computacional mínimo (~1.88s/paso frente a 1.70s del baseline) y reduce drásticamente los requisitos de almacenamiento (412 KB por tarea frente a 6.5 GB para las anclas de EWC).

Significancia y afirmaciones
El artículo afirma que el bajo rendimiento de la regularización en el espacio de pesos en los LLMs no es un problema de ajuste de hiperparámetros, sino una consecuencia estructural de la polisemanticidad. Al mover la restricción de regularización a una base de características monosemánticas, el método logra una protección más selectiva del conocimiento. Los autores sostienen que la efectividad de la regularización depende críticamente del espacio de representación en el cual se imponen las restricciones. El método demuestra que es posible lograr un fuerte desempeño de aprendizaje continuo sin almacenar datos pasados o modificar la arquitectura del modelo en el tiempo de inferencia, siempre que la regularización opere en un sistema de coordenadas alineado con la estructura conceptual interna del modelo. El artículo concluye que la regularización basada en SAE ofrece una alternativa más selectiva y escalable a los métodos tradicionales de espacio de pesos para los LLMs.