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

Resumo Técnico: Dos Pesos para as Características: Regularização de Ativação Guiada por SAE para Aprendizado Contínuo de LLMs

Declaração do Problema
O aprendizado contínuo (CL - Continual Learning) em Grandes Modelos de Linguagem (LLMs) enfrenta o desafio do esquecimento catastrófico, onde o ajuste fino sequencial sobrescreve representações necessárias para tarefas anteriores. Embora métodos baseados em regularização como o Elastic Weight Consolidation (EWC) sejam padrão para modelos menores, eles apresentam um desempenho significativamente inferior em LLMs. Os autores argumentam que essa falha é estrutural, decorrente da natureza "polissemântica" dos LLMs. Devido à hipótese da superposição, neurônios individuais e seus pesos associados codificam múltiplos conceitos não relacionados simultaneamente. Consequentemente, estimativas de importância no espaço de pesos (ex: informação de Fisher) são muito grosseiras; proteger um peso para preservar o Conceito A restringe inadvertidamente o Conceito B, criando uma penalidade não seletiva que sufoca a plasticidade. Além disso, os métodos existentes frequentemente falham em impor explicitamente a plasticidade necessária para aprender novas tarefas, levando à "intransigência", onde o modelo retém o conhecimento antigo, mas não consegue se adaptar.

Metodologia
O artigo propõe a regularização de ativação guiada por SAE, um método que desloca o alvo da regularização do espaço de parâmetros emaranhado para um espaço de ativação mais seletivo usando Autoencoders Esparsos (SAEs) pré-treinados.

1. Mudança de Sistema de Coordenadas: Em vez de penalizar o desvio de pesos, o método utiliza um SAE pré-treinado e congelado como um dicionário de características monosemânticas fixo. Os SAEs decompõem ativações densas do modelo em direções de características esparsas que são frequentemente mais interpretáveis e menos emaranhadas do que neurônios individuais.
2. Formulação de Otimização Restrita: O objetivo de treinamento é derivado de um problema de otimização restrita, em vez de um design de penalidade ad hoc. O modelo deve minimizar a perda da tarefa atual sujeito a duas restrições:
   • Estabilidade: Características irrelevantes para a tarefa (baixa relevância) não devem desviar além de um orçamento específico ($\varepsilon_s$).
   • Plasticidade: Características relevantes para a tarefa (alta relevância) devem desviar suficientemente ($\ge \varepsilon_p$) para evitar soluções degeneradas onde o modelo não aprende nada.
3. Função de Perda: Usando relaxação Lagrangiana e uma penalidade de hinge quadrada, os autores derivam uma perda de treinamento que compreende uma perda de proteção (penalizando o desvio excessivo em características protegidas) e uma perda de guia (penalizando o desvio insuficiente em características relevantes). O EWC é mostrado como um caso especial deste framework, onde o espaço de representação são os pesos e a restrição de plasticidade está ausente.
4. Construção de Máscara: Para cada nova tarefa, uma máscara de relevância contínua ($M_t$) é construída usando dados da tarefa atual passados pelo modelo base congelado e pelo SAE. Esta máscara identifica características relevantes via magnitude de ativação e propaga a relevância para características vizinhas usando k-NN em vetores de decodificação. Crucialmente, uma vez construída a máscara, todos os dados da tarefa são descartados; apenas a máscara compacta é retida.
5. Treinamento: Durante o treinamento, o modelo atual (com LoRA) e o modelo base congelado são ambos codificados pelo SAE. O desvio de característica é ponderado pela máscara para calcular as perdas de proteção e de guia.

Principais Contribuições

• Regularização Guiada por SAE: Um método prático de CL que regulariza no espaço de ativação usando características de SAE. Ele não requer dados de tarefas anteriores, não armazena ativações, não possui parâmetros específicos de tarefa e não exige roteamento em tempo de inferência após a construção da máscara.
• Derivação Principiada: O compromisso estabilidade-plasticidade é formulado como um problema de otimização restrita no espaço de características SAE, resultando em uma perda de treinamento de hinge quadrada via relaxação Lagrangiana.
• Evidência Empírica de Polissemantismo: O artigo fornece evidência mecanística de que a proteção no espaço de pesos é não seletiva. Experimentos mostram que representações relevantes para a tarefa são linearmente separáveis no espaço de características SAE (AUC ~0,88), mas indistinguíveis do acaso no espaço de pesos/neurônios (AUC ~0,50). Além disso, proteger pesos para uma tarefa restringe as características para uma tarefa subsequente a uma taxa de 91–96% da tarefa original, enquanto a proteção no espaço de características reduz essa restrição colateral para 43–61%.
• Escalabilidade: Devido à menor dimensionalidade do espaço de características em comparação ao espaço de parâmetros, o armazenamento por tarefa é reduzido de gigabytes (âncoras de pesos) para sub-megabytes (máscaras de características).

Resultados
O método foi avaliado no TRACE-5000 (8 tarefas de domínios cruzados) e MedCL (10 tarefas biomédicas) usando Gemma-2 9B-it com LoRA.

• Desempenho: No TRACE, a regularização guiada por SAE alcançou o melhor desempenho entre abordagens não arquiteturais (OP = 0,545), superando regularizadores de espaço de pesos (EWC: 0,447, SI: 0,513), métodos de projeção de gradiente e baselines de replay (mesmo com 10% de dados). No MedCL, liderou todos os métodos que não retêm exemplos de tarefas anteriores (OP = 0,510).
• Compromisso Estabilidade-Plasticidade: Ao contrário do EWC, que alcançou alta retenção (ARA = 0,987) sacrificando severamente a plasticidade (0,453), o método proposto manteve alta plasticidade (0,567, comparável ao ajuste fino sem proteção) enquanto retinha 95% do desempenho aprendido.
• Eficiência: O método adiciona overhead computacional mínimo (~1,88s/passo vs 1,70s do baseline) e reduz drasticamente os requisitos de armazenamento (412 KB por tarefa vs 6,5 GB para âncoras EWC).

Significância e Alegações
O artigo afirma que o baixo desempenho da regularização de espaço de pesos em LLMs não é um problema de ajuste de hiperparâmetros, mas uma consequência estrutural da polissemantia. Ao mover a restrição de regularização para uma base de características monosemânticas, o método alcança uma proteção mais seletiva do conhecimento. Os autores afirmam que a eficácia da regularização depende criticamente do espaço de representação no qual as restrições são impostas. O método demonstra que é possível alcançar um forte desempenho de aprendizado contínuo sem armazenar dados passados ou modificar a arquitetura do modelo em tempo de inferência, desde que a regularização opere em um sistema de coordenadas alinhado com a estrutura de conceitos interna do modelo. O artigo conclui que a regularização baseada em SAE oferece uma alternativa mais seletiva e escalável aos métodos tradicionais de espaço de pesos para LLMs.