Resumo Técnico: Autoencoders Esparsos Equivariantes

Declaração do Problema

Modelos de aprendizado de máquina, particularmente em domínios científicos, operam frequentemente em dados com simetrias inerentes (ex: rotações em dobramento de proteínas, imagens de galáxias ou microscopia de células). Embora os Autoencoders Esparsos (SAEs) tenham emergido como uma ferramenta primária para interpretabilidade mecanística ao decompor ativações neurais densas em características esparsas e interpretáveis, eles enfrentam uma limitação crítica quando aplicados a dados simétricos: identificabilidade.

SAEs padrão assumem uma Hipótese de Representação Linear (LRH) onde as ativações são combinações lineares esparsas de características. No entanto, quando os dados de entrada possuem simetrias (ações de grupo $G$), um SAE vanilla deve aprender direções de características distintas para o mesmo conceito semântico sob diferentes transformações (ex: uma característica de "gato" para cada rotação possível). Isso leva a:

1. Emaranhamento de Características: O modelo pode aprender características não ortogonais e não interpretáveis para reconstruir entradas díspares.
2. Não Identificabilidade: Diferentes explicações podem se ajustar aos dados igualmente bem, mas apenas uma pode alinhar-se com a verdadeira estrutura subjacente.
3. Desalinhamento de Métricas: A qualidade da reconstrução (MSE) frequentemente falha em correlacionar-se com a utilidade da característica, pois os SAEs podem aprender características específicas de reconstrução que não correspondem aos conceitos reais subjacentes.

O artigo postula que os SAEs atuais falham ao não considerar essas simetrias, levando a uma interpretabilidade subótima em domínios científicos onde tais simetrias são prevalentes.

Metodologia: SAEs Equivariantes

Os autores propõem Autoencoders Esparsos Equivariantes, que estendem a Hipótese de Representação Linear para contabilizar explicitamente as simetrias de grupo.

1. Hipótese de Representação Linear Estendida

Os autores reformulam a LRH para afirmar que, para um grupo $G$ atuando sobre entradas $X$, as ativações da rede neural transformam-se de forma previsível. Especificamente, se $x$ é transformado por $g \in G$, as ativações transformam-se via uma aplicação linear $\rho(g)$:
$$\psi(gx) = \rho(g) \sum_{i=1}^C z(x)_i D_i = \sum_{i=1}^C z(x)_i (\rho(g) D_i)$$
Aqui, as magnitudes das características $z(x)$ permanecem invariantes (coeficientes esparsos), enquanto os vetores de características $D_i$ transformam-se de acordo com a ação do grupo.

2. Design da Arquitetura

O SAE Equivariante desacopla o aprendizado de características invariantes do aprendizado de regras de transformação:

• SAE Invariante: Um par codificador-decodificador treinado para mapear todas as ativações dentro de uma órbita (ex: todas as rotações de uma imagem) para uma representação latente canônica única. Isso aprende as características invariantes.
• Módulo de Transformação Linear: Um conjunto separado de matrizes aprendíveis ($M, R$) que modela como a reconstrução canônica se transforma de volta para o espaço de entrada original. Para um grupo gerado por $g, h$, o modelo aprende matrizes tais que:
  $$M_p R_q D(E(\psi(g^p h^q x))) = \psi(g^p h^q x)$$
  Isso permite que o SAE capture a estrutura de simetria sem forçar o dicionário a ser equivariante, o que exigiria um número impraticável de parâmetros.

3. Estratégia de Treinamento

O treinamento ocorre em duas etapas independentes (que podem ser paralelizadas):

1. Treinamento do SAE Invariante: O codificador e o decodificador são otimizados para minimizar a perda de reconstrução em entradas canônicas enquanto impõem invariância. Os autores experimentam três objetivos de perda:
   • $L_{Canonical}$: Reconstrói uma amostra canônica de qualquer membro da órbita.
   • $L_{Latent}$: Impõe invariância diretamente nos códigos latentes ($E(\psi(x)) \approx E(\psi(gx))$).
   • $L_{Output}$: Impõe invariância nas saídas do SAE.
2. Aprendizado de Transformação: As matrizes $M$ e $R$ são otimizadas para minimizar o erro entre as ativações transformadas previstas e as ativações transformadas reais do modelo base.

Principais Contribuições

1. Extensão Teórica: O artigo estende a Hipótese de Representação Linear para dados simétricos, definindo formalmente como os coeficientes e vetores de características se comportam sob ações de grupo.
2. Inovação Algorítmica: O design de SAEs Equivariantes que separa o aprendizado de características invariantes do modelagem de transformação, evitando a explosão de parâmetros de dicionários grupo-equivariantes.
3. Validação Empírica: Demonstração através de conjuntos de dados sintéticos e do mundo real (formas, imagens de galáxias, células sanguíneas) que os SAEs Equivariantes superam os baselines de SAE vanilla, Archetypal e Group Cross-coder em recuperação de características e tarefas de sondagem (probing) subsequentes.
4. Insight de Métrica: A descoberta de que a qualidade da reconstrução é inversamente correlacionada com a utilidade da característica em configurações simétricas. SAEs Equivariantes frequentemente alcançam uma qualidade de reconstrução inferior (maior MSE) do que SAEs vanilla, mas geram características significativamente mais interpretáveis e úteis.

Resultados Experimentais

Os autores avaliaram seu método em:

• Modelos de Brinquedo (Toy Models): Dados sintéticos com características equivariantes de verdade fundamental conhecidas.
• Dataset de Formas: Imagens sintéticas com simetrias $C_4$ (rotações) e $D_4$ (rotações + reflexões), usando modelos base MLP, CNN e ViT.
• GalaxyMNIST: Imagens reais de galáxias com simetrias $C_4/D_4$.
• MLL23: Imagens de células sanguíneas periféricas utilizando um modelo de fundação DinoBloom.

Principais Descobertas:

• Recuperação de Características: Em modelos de brinquedo, os SAEs Equivariantes alcançaram uma recuperação de características superior (Similaridade de Cosseno Média) e detecção (F1 Score) em comparação com os baselines, especialmente em regimes de sobrecompletude onde a interferência é alta.
• Sondagem Subsequente (Downstream Probing): Em todos os datasets, sondas treinadas nos latentes do SAE Equivariante alcançaram maior precisão na classificação de propriedades de entrada (ex: presença de forma, orientação, tipo de célula) do que sondas treinadas nos latentes de SAEs vanilla. Em alguns casos (ex: GalaxyMNIST), eles igualaram o desempenho de sondas treinadas diretamente nas ativações do modelo base.
• Reconstrução vs. Utilidade: Uma tendência consistente surgiu onde os SAEs Equivariantes apresentaram maior erro de reconstrução (pior MSE) do que os SAEs vanilla. No entanto, essa reconstrução "pior" correspondeu a características que estavam mais alinhadas com a verdadeira estrutura semântica e eram mais úteis para tarefas subsequentes.
• Modelagem de Transformação: As matrizes lineares aprendidas ($M, R$) explicaram com sucesso mais de 97% da variância nas transformações de ativação para a maioria dos modelos, validando a hipótese de que as simetrias de entrada induzem transformações lineares no espaço de ativação.

Significância e Alegações

O artigo alega que alinhar os priors de métodos de interpretabilidade com a estrutura conhecida dos dados (especificamente simetrias) é crucial para obter explicações fiéis. Os autores destacam duas implicações principais:

1. Reconstrução não é Interpretabilidade: Depender exclusivamente de métricas de reconstrução (como MSE) para avaliar SAEs é perigoso, particularmente para dados simétricos. Um modelo que reconstrói entradas de forma pobre ainda pode descobrir as características mais úteis e interpretáveis.
2. Interpretabilidade de Domínio Específico: Para aplicações científicas (ex: biologia, astronomia) onde as simetrias de dados são fundamentais, os SAEs padrão são insuficientes. Os SAEs Equivariantes fornecem um caminho para descobrir características que são robustas a essas simetrias, potencialmente levando a novos insights científicos e sistemas de IA mais confiáveis.

O trabalho sugere que a pesquisa futura em interpretabilidade deve ir além da codificação esparsa genérica e incorporar priors estruturais de domínio específico para evitar as armadilhas da não identificabilidade.