Resumo Técnico: Rumo a Autoencoders Esparsos Identificáveis

Declaração do Problema

Os Autoencoders Esparsos (SAEs) tornaram-se uma ferramenta padrão para interpretar e intervir nas representações internas de redes neurais, particularmente em Grandes Modelos de Linguagem (LLMs). Eles decompõem ativações de alta dimensão em combinações lineares esparsas de "conceitos" aprendidos (átomos do dicionário). No entanto, o artigo identifica uma falha crítica nas implementações atuais de SAEs: a instabilidade. O folclore empírico e trabalhos teóricos recentes sugerem que diferentes execuções de treinamento no mesmo conjunto de dados frequentemente resultam em diferentes dicionários de conceitos e códigos esparsos. Essa não-identificabilidade compromete a confiabilidade dos SAEs para interpretabilidade mecanística, direcionamento (steering) e análise científica, pois os "recursos" (features) extraídos podem ser artefatos do processo de treinamento, em vez de estruturas estáveis nos dados.

Os autores argumentam que, embora a teoria de aprendizado de dicionário ofereça condições para identificabilidade, as arquiteturas de SAE padrão falham em satisfazê-las devido a três características específicas, denominadas "tríade sombria":

1. Recursos Bidirecionais: As funções de ativação padrão (como ReLU ou TopK) restringem os coeficientes a serem positivos, forçando o dicionário a aprender conceitos opostos (ex: $d$ e $-d$) para representar sinais bidirecionais, o que viola as condições para recuperação única.
2. Condicionamento do Dicionário: Os SAEs padrão não impõem condições (como a Isometria Restrita) ao dicionário aprendido, levando a cenários onde múltiplos códigos esparsos podem reconstruir o mesmo dado com igual eficiência.
3. Expressividade do Encoder: Os encoders quase-lineares tipicamente usados em SAEs podem carecer da capacidade de mapear observações para os códigos esparsos corretos, mesmo que o dicionário seja bem condicionado.

Metodologia

O artigo propõe uma abordagem unificada para preencher a lacuna entre a escalabilidade dos SAEs e as garantias teóricas do aprendizado de dicionário. Os autores introduzem o Autoencoder Esparso Identificável (iSAE), que modifica a arquitetura e o procedimento de treinamento padrão do TopK SAE através de três intervenções específicas:

1. Recursos Bidirecionais (AbsTopK)

Para abordar a restrição a coeficientes positivos, os autores substituem a ativação TopK padrão por AbsTopK. Esta função seleciona os $k$ coeficientes com os maiores valores absolutos, permitindo tanto cargas positivas quanto negativas. Isso efetivamente dobra a capacidade de representação do dicionário sem aumentar a contagem de parâmetros e alinha o SAE com formulações clássicas de codificação esparsa que permitem coeficientes negativos.

2. Condicionamento do Dicionário (Regularização aRIP)

Os autores introduzem uma relaxação de dados da Propriedade de Isometria Restrita (RIP), denominada RIP aproximada (aRIP).

• Teoria: A RIP clássica exige que qualquer combinação de $k$ átomos preserve normas, uma condição combinatoriamente difícil de verificar para dicionários aprendidos e sobrecompletos. Os autores propõem impor essa propriedade apenas na união dos padrões de esparsidade realmente observados nos dados.
• Implementação: Eles derivam um regularizador baseado nos autovalores da submatriz de Gram dos suportes observados. Este regularizador é adicionado à perda de reconstrução padrão para incentivar o dicionário aprendido a satisfazer a condição aRIP, garantindo assim que os códigos esparsos sejam unicamente recuperáveis (salvo permutação e sinal).

3. Expressividade do Encoder (Codificação de Múltiplas Etapas)

Para abordar a limitação dos encoders quase-lineares, os autores incorporam um esquema de codificação de múltiplas etapas inspirado em LISTA (Gregor & LeCun, 2010). Em vez de uma única passagem direta, o encoder refina iterativamente o código esparso $z$ ao longo de $T$ etapas (tipicamente $T=5$) usando uma matriz de tamanho de passo aprendida. Isso aumenta a expressividade do mapeamento do dado de entrada $x$ para o código $z$, permitindo que o encoder amortizado aproxime melhor a solução do problema de codificação esparsa.

Contribuições Principais

Contribuições Teóricas

• Definição de Identificabilidade Estatística: O artigo define a identificabilidade do SAE não como a recuperação de um mecanismo de "verdade fundamental" (que é frequentemente desconhecido), mas como estabilidade entre execuções. Um SAE é identificável se execuções de treinamento independentes no mesmo conjunto de distribuição produzirem aproximadamente o mesmo dicionário e códigos esparsos (salvo equivalências triviais como permutação e sinal).
• Teorema da Impossibilidade (Teorema 3.2): Os autores provam que aprender um dicionário estável não é suficiente para aprender códigos estáveis. Mesmo com um dicionário perfeitamente estável, se o dicionário não satisfizer propriedades de condicionamento específicas, o problema de codificação esparsa permanece não-identificável.
• Teorema da Quase-Identificabilidade (Teorema 3.6): Eles provam que, se um SAE atingir baixo erro de reconstrução e seu dicionário satisfizer a condição aRIP (com riqueza e diversidade suficientes na distribuição de dados), o dicionário e os códigos esparsos aprendidos são quase-identificáveis.

Contribuições Empíricas

• Novas Métricas: Os autores propõem medir a identificabilidade do código diretamente, usando Interseção sobre União (IoU) e erro $\ell_2$ em padrões esparsos, além da similaridade de cosseno de dicionário (DCS) padrão. Eles demonstram que um alto DCS não garante alta identificabilidade de código.
• Variantes de Modelo: Apresentam duas variantes:
  • iSAE: Usa AbsTopK + regularização aRIP.
  • iSAE-ME: Usa AbsTopK + regularização aRIP + Codificação de Múltiplas Etapas.

Resultados

Dados Sintéticos

• Reconstrução: A transição de TopK para AbsTopK reduziu o erro de reconstrução em 30–40%.
• Identificabilidade: A variante iSAE-ME alcançou identificabilidade quase perfeita em configurações sintéticas. O encoder de múltiplas etapas foi crucial aqui, permitindo que o encoder amortizado igualasse o desempenho de um solver "oráculo" (Orthogonal Matching Pursuit) que possui acesso ao verdadeiro dicionário.
• Condicionamento: O regularizador aRIP melhorou com sucesso o condicionamento do dicionário, o que era necessário para que o oráculo alcançasse alta identificabilidade.

Dados do Mundo Real (LLM e Visão)

• Reconstrução: Em ativações de Pythia-160M (LLM) e DINOv2 (Visão), o encoder de múltiplas etapas (iSAE-ME) alcançou o menor erro de reconstrução (MSE), superando significativamente os baselines.
• Identificabilidade:
  • Dicionário: O AbsTopK melhorou a estabilidade do dicionário (DCS) comparado ao TopK.
  • Códigos: Os resultados foram mistos. Embora o solver oráculo tenha mostrado melhor identificabilidade de código com o dicionário condicionado, o encoder amortizado do SAE nem sempre explorou isso plenamente. Em alguns casos reais, o encoder de múltiplas etapas na verdade piorou a identificabilidade do código, apesar de uma melhor reconstrução, sugerindo que a dinâmica de otimização em ambientes reais é complexa e o encoder pode ainda não estar convergindo para a solução ótima.
• Tarefas Downstream (SAEBench): Em ativações de LLM, o iSAE-ME apresentou o melhor desempenho na minimização da perda de entropia cruzada (fidelidade de reconstrução). No entanto, para "Remoção de Correlação Espúria" (SCR), o baseline TopK teve o melhor desempenho, provavelmente porque zerar um recurso TopK possui uma interpretação de "ausência" mais clara. O iSAE-ME teve o melhor desempenho em "Perturbação de Sonda Alvo" (TPP).

Significância e Alegações

O artigo alega que a identificabilidade é um requisito de "nível mínimo" para SAEs usados em aplicações de alto risco, como direcionamento de modelos e supervisão. Se dois SAEs treinados nos mesmos dados produzem recursos diferentes, a validade científica desses recursos é comprometida.

Os autores posicionam seu trabalho como um passo para tornar os SAEs estatisticamente confiáveis. Eles não alegam ter resolvido o problema de recuperar os conceitos de "verdade fundamental" originais, nem alegam que seus modelos generalizam a identificabilidade para distribuições fora dos dados de treinamento. Em vez disso, eles demonstram que:

1. SAEs TopK padrão são inerentemente instáveis e não-identificáveis.
2. Ao abordar explicitamente o condicionamento do dicionário (via aRIP) e a expressividade do encoder, é possível construir SAEs que são estatisticamente identificáveis (estáveis entre execuções) e alcançam reconstrução superior.
3. A lacuna entre o desempenho sintético e o do mundo real destaca que, embora o framework teórico seja válido, os desafios de otimização em ativações de redes neurais reais permanecem uma barreira significativa para a identificabilidade perfeita.

O artigo conclui que essas mudanças arquiteturais e de regularização resultam em modelos que são mais robustos e teoricamente fundamentados, fornecendo uma base mais sólida para a interpretabilidade de redes neurais de larga escala.