Resumo Técnico: SAE-FD para Aprendizado Contínuo de Modelos de Linguagem de Grande Escala

Declaração do Problema

Modelos de Linguagem de Grande Escala (LLMs) enfrentam o desafio do esquecimento catastrófico ao se adaptarem sequencialmente a múltiplas tarefas. Os métodos existentes de aprendizado contínuo (CL) geralmente operam em espaços de representação densa (espaços de pesos, gradientes ou saída). Uma limitação fundamental dessas abordagens é a superposição de características, onde múltiplos conceitos são codificados em dimensões sobrepostas dentro do espaço de ativação densa. Esse emaranhamento torna difícil proteger seletivamente o conhecimento previamente aprendido sem restringir inadvertidamente a aprendizagem de novas tarefas, levando a um compromisso subótimo entre plasticidade e estabilidade.

Metodologia: SAE-FD

Os autores propõem SAE-FD (Distilação de Características com Autoencoder Esparso), um método que ancora as representações do modelo no espaço de características esparsas de um Autoencoder Esparso (SAE) pré-treinado. Ao decompor ativações densas em uma base esparsa e sobrecompleta, o SAE-FD visa desentrelaçar representações superpostas, permitindo uma regularização mais direcionada.

O método opera como um pipeline de três etapas:

1. Treinamento do SAE (Etapa 1)

Antes do início da sequência de aprendizado contínuo, um SAE com Portão é treinado em ativações de texto diversas do modelo base.

• Entrada: Ativações da última camada MLP ($h \in \mathbb{R}^d$) de fontes diversas (classificação, QA, benchmarks de conhecimento).
• Arquitetura: O SAE com Portão separa a codificação em um caminho de portão ($g$) e um caminho de magnitude ($m$), produzindo características pré-ReLU ($f_{pre}$) e características pós-ReLU ($f$).
• Objetivo: Minimizar a perda de reconstrução ($\|\hat{h} - h\|_2^2$) e uma penalidade de esparsidade L1 ($\lambda_{\ell1}\|f\|_1$).
• Resultado: Um SAE congelado que decompõe ativações densas em características esparsas ($D \gg d$), efetivamente desentrelaçando conceitos.

2. Treinamento de Tarefa e Captura de Âncoras (Etapa 2)

Para cada tarefa $t$ na sequência:

• Ajuste Fino: O modelo é ajustado finamente usando LoRA com o objetivo padrão de modelagem de linguagem.
• Captura de Âncoras: Após o treinamento em uma tarefa, o modelo realiza uma passagem direta em um pequeno conjunto de amostras de âncora ($N_a$). As saídas por token da MLP ($H_i$) são extraídas e armazenadas.
• Buffer: Essas âncoras acumulam-se ao longo das tarefas. Durante o treinamento subsequente, um lote é amostrado uniformemente do buffer para calcular a perda de distilação, amortizando o custo em todas as tarefas anteriores.

3. Distilação de Características Consciente de Esparsidade (Etapa 3)

Durante o treinamento de uma nova tarefa ($t+1$), o modelo é regularizado comparando suas representações atuais com as âncoras armazenadas, ambas codificadas através do SAE congelado. A perda de distilação ($L_{FD}$) consiste em dois termos complementares:

• Perda de Direção Cosseno ($L_{cos}$): Penaliza o desvio angular entre o vetor de características pré-ReLU atual e o vetor de características da âncora. Isso preserva a direção global da representação.
• Perda de Magnitude Ativa ($L_{mag}$): Aplica o Erro Quadrático Médio (MSE) apenas nas características que estavam ativas (pós-ReLU $> 0$) na âncora. Isso restringe a magnitude das características que contribuíram para tarefas anteriores, deixando as características inativas livres para serem recrutadas para novas tarefas.

Mecanismo Adaptativo de $\lambda$

Para equilibrar a aprendizagem da tarefa e a regularização, os autores introduzem um mecanismo adaptativo para o coeficiente de ponderação $\lambda$ na perda total $L = L_{task} + \lambda \cdot L_{FD}$.

• Em vez de um $\lambda$ fixo, o sistema visa uma razão de contribuição fixa $\rho$ da perda de distilação em relação à perda total.
• $\lambda$ é atualizado via média móvel exponencial (EMA) para satisfazer $\rho = \frac{\lambda L_{FD}}{\lambda L_{FD} + L_{task}}$.
• Comportamento: $\lambda$ começa alto quando uma nova tarefa começa (quando $L_{task}$ é grande e o risco de esquecimento é alto) e relaxa à medida que o modelo converge, prevenindo restrições excessivas durante o ajuste fino.

Contribuições Principais

1. Framework SAE-FD: O primeiro método de aprendizado contínuo a realizar regularização no espaço de características esparsas e desacopladas de SAEs, abordando o problema de superposição de características inerente a espaços densos.
2. Função de Perda de Distilação Consciente de Esparsidade: Uma função de perda novel que combina preservação de direção cosseno com correspondência seletiva de magnitude em características ativas, aproveitando a esparsidade das representações do SAE para minimizar interferência com novas tarefas.
3. Regularização Adaptativa: Um mecanismo adaptativo de $\lambda$ que ajusta dinamicamente a força da regularização com base na magnitude relativa das perdas de tarefa e distilação, garantindo proteção forte durante períodos de alto risco sem impedir a convergência.

Resultados Experimentais

Os autores avaliaram o SAE-FD em dois benchmarks através de três arquiteturas de modelo (LLaMA-2-7B, Vicuna-7B, Mistral-7B).

• Benchmark TRACE (8 Tarefas Heterogêneas):
  • No LLaMA-2-7B, o SAE-FD alcançou 52,70% de Precisão Média (AA) com apenas −0,46% de Transferência Reversa (BWT).
  • Isso supera a melhor linha de base baseada em regularização publicada (GORP) em +2,3% de AA e reduz o esquecimento em 34% (−0,46 vs −0,70 de BWT).
  • Supera significativamente outros métodos baseados em LoRA (EWC, O-LoRA, TreeLoRA), que mostraram BWT pior que −3%.
• Generalização:
  • O SAE-FD consistentemente alcançou o melhor AA e BWT em todas as três famílias de modelos.
  • No Vicuna-7B, superou o O-LoRA em +10,1% de AA e alcançou o melhor BWT (−0,69%) entre todos os métodos, incluindo linhas de base de ajuste fino completo.
• Estudos de Ablação:
  • Componentes de Perda: Tanto a perda de direção cosseno quanto a perda de magnitude ativa são necessárias; usar qualquer uma delas isoladamente resultou em AA menor ou esquecimento maior.
  • $\lambda$ Adaptativo: Substituir o mecanismo adaptativo por um $\lambda=1,0$ fixo degradou o BWT de −0,46% para −2,38%.
  • SAE vs. MSE Bruto: A distilação no espaço de características do SAE superou a distilação MSE bruta em ativações densas, confirmando o benefício do espaço desacoplado.
• Eficiência: O custo computacional é aproximadamente 1,85× o do SeqLoRA, com um custo constante por etapa independentemente do número de tarefas anteriores. O armazenamento requer ~0,8 GB por tarefa para ativações de âncora.

Significado e Alegações

O artigo afirma que o SAE-FD conecta com sucesso a interpretabilidade mecânica (via SAEs) com o aprendizado contínuo. Ao utilizar representações esparsas e desentrelaçadas originalmente desenvolvidas para entender redes neurais, o método fornece um viés indutivo eficaz para controlá-las.

Os autores afirmam que operar em um espaço de características esparsas permite uma preservação mais direcionada do conhecimento, resolvendo o problema de "superposição de características" que aflige a regularização em espaços densos. O método alcança um compromisso superior entre plasticidade e estabilidade, permitindo que LLMs se adaptem a tarefas em evolução sem retreinamento do zero, mantendo alto desempenho em tarefas previamente aprendidas.

Limitações notadas pelos autores:

• Requer um custo de treinamento único para um SAE específico do modelo.
• Os requisitos de armazenamento para ativações de âncora crescem linearmente com o número de tarefas (embora compressíveis).
• Os experimentos foram conduzidos em uma ordem fixa de tarefas para o benchmark TRACE, embora a consistência entre múltiplos modelos sugira robustez.
• As comparações com algumas linhas de base devem ser interpretadas com cautela devido a diferenças nas configurações experimentais (por exemplo, ajuste fino completo vs. LoRA, contagem de épocas).