SAE no Gemma-2-2B: fração de features correspondentes sobe de 2% para 29% com 50M de tokens

Os sparse autoencoders (SAE) se tornaram uma das principais ferramentas de interpretabilidade de modelos de linguagem e, em geral, são treinados com centenas de milhões ou bilhões de tokens de ativações. Quase ninguém verificou se essa escala é realmente necessária. Os autores de um novo trabalho publicado no Hugging Face Daily Papers fizeram esse teste: usaram quatro arquiteturas de SAE, três modelos de linguagem e orçamentos de 5M de tokens até a referência de 500M. A conclusão principal: a resposta depende do que exatamente você mede. No Gemma-2-2B com BatchTopK, a fração de features com correspondência estrita sobe de 2% com 5M de tokens para 29% com 50M.

Duas medidas de convergência

Os pesquisadores separaram dois conceitos. O primeiro é o comportamento funcional: a qualidade da reconstrução e os resultados nos benchmarks padrão de SAE. O segundo é a reprodutibilidade "um para um": um SAE treinado de forma independente encontra as mesmas features que o SAE de referência?

A diferença foi expressiva. A reconstrução frequentemente chega ao nível da referência completa com orçamentos bem menores. A reprodutibilidade das features, por outro lado, continua melhorando conforme os dados aumentam. Um SAE pode parecer "pronto" pelas métricas, mas o dicionário de features dentro dele continua instável.

Principais observações do trabalho:

• Escala do experimento: 4 arquiteturas de SAE, 3 modelos de linguagem, orçamentos de 5M a 500M de tokens.
• Gemma-2-2B, BatchTopK: 2% de features com correspondência estrita com 5M de tokens e 29% com 50M.
• A necessidade de dados depende muito da métrica: a reconstrução satura antes; a reprodutibilidade, não.
• A amostragem aleatória preserva melhor a distribuição original das ativações e costuma vencer em reconstrução.
• A amostragem por diversidade frequentemente melhora o resultado no AutoInterp.

Como escolher os dados

Importa não só o volume, mas também o método de seleção. A amostragem aleatória reproduz com mais fidelidade a distribuição das ativações e, por isso, costuma ser a melhor opção quando o objetivo é uma boa reconstrução. A amostragem baseada em diversidade cobre mais comportamentos raros do modelo, o que muitas vezes eleva as notas no AutoInterp, a interpretação automática de features. Não existe um único amostrador "melhor": a escolha depende do objetivo.

A seção teórica do trabalho explica por que as duas métricas divergem. A consistência estatística do treinamento em um subconjunto é uma condição. A recuperação de features um para um exige condições mais rígidas de identificabilidade e de otimização. Em outras palavras, mesmo que o subconjunto produza uma boa aproximação do funcional-alvo, isso não garante que execuções diferentes cheguem ao mesmo conjunto de features.

O que fazer na prática

A conclusão para equipes que treinam SAEs é bastante pé no chão. Se o objetivo é obter rapidamente um SAE funcional com boa reconstrução, um orçamento de centenas de milhões de tokens, pelos resultados, costuma ser excessivo. Isso barateia os experimentos, sobretudo na busca por arquiteturas e hiperparâmetros.

Já se você publica afirmações sobre features específicas, como "este latente codifica tal conceito", as métricas de reconstrução sozinhas não bastam. É preciso uma verificação explícita de reprodutibilidade: várias execuções independentes e o pareamento das features entre elas. Uma fração baixa de correspondências com orçamento pequeno, como os mesmos 2% com 5M de tokens, é um alerta de que features individuais podem ser artefatos de um treinamento específico.

Por fim, vale escolher o amostrador conforme o objetivo. Para reconstrução, comece pela amostragem aleatória; para interpretabilidade medida pelo AutoInterp, teste a seleção por diversidade. Os autores ressaltam que, com objetivos e orçamentos bem alinhados, a pesquisa com SAEs pode ficar mais barata e mais confiável. Ressalva: os números valem para modelos e arquiteturas específicos, e a extrapolação para outras escalas exige verificação própria.