Resumo Técnico: Autoencoders Esparsos Temporais Residualizados para Interpretação de Modelos de Difusão

Declaração do Problema

Modelos de difusão texto-para-imagem geram imagens por meio de um processo iterativo de remoção de ruído, fazendo com que as camadas neurais internas produzam trajetórias de ativações em vez de representações estáticas únicas. Embora Autoencoders Esparsos (SAEs) tenham sido recentemente aplicados para decompor ativações de difusão em direções de características interpretáveis, as abordagens existentes geralmente analisam ativações em instantes de tempo isolados, condicionam-se ao tempo ou aplicam SAEs a arquiteturas como Diffusion Transformers.

Uma limitação crítica desses métodos existentes é que eles não levam explicitamente em conta a forte correlação temporal inerente às ativações de difusão. As ativações em instantes de tempo adjacentes são altamente correlacionadas, com grande parte do sinal em um passo sendo linearmente previsível a partir do passo anterior. Treinar um SAE diretamente sobre ativações brutas concatenadas temporalmente arrisca desperdiçar a capacidade do modelo no aprendizado dessas dinâmicas redundantes e linearmente previsíveis, em vez de capturar a estrutura esparsa subjacente das características.

Metodologia

Os autores propõem SAEs Temporais Residualizados, um framework projetado para aprender características esparsas a partir de trajetórias de ativação, separando explicitamente dinâmicas linearmente previsíveis da estrutura residual.

1. Coleta e Pré-processamento de Dados

• Fonte: As ativações são coletadas do ramo condicional do U-Net no Stable Diffusion 1.5 durante a orientação sem classificador. Especificamente, a contribuição residual da camada de atenção cruzada do bloco médio é extraída.
• Construção de Trajetória: Para uma imagem gerada, as ativações são subamostradas a cada $\ell$ passos de remoção de ruído (por exemplo, $\ell=10$) para formar uma sequência de tensores de ativação $\{a(\tau_i)\}$.

2. Residualização Linear

Para remover redundância temporal, o método ajusta preditores lineares entre blocos consecutivos de ativação normalizados:

• Normalização: As ativações são normalizadas por instante de tempo usando estatísticas do conjunto de treinamento.
• Regressão Ridge: Para cada transição adjacente $i-1 \to i$, um regressor ridge $(W_i, b_i)$ é ajustado para prever a ativação atual $\bar{a}_{q,i}$ a partir da anterior $\bar{a}_{q,i-1}$.
• Cálculo do Resíduo: O resíduo $r_{q,i} = \bar{a}_{q,i} - (W_i \bar{a}_{q,i-1} + b_i)$ captura o componente da ativação não linearmente previsível a partir do passo anterior.

3. Treinamento do SAE em Trajetórias Residualizadas

A entrada para o SAE é um vetor concatenado $z_q$ consistindo na ativação inicial e nos resíduos subsequentes:
$$z_q = [\bar{a}_{q,0}, r_{q,1}, \dots, r_{q,T-1}]$$
Um BatchTopK SAE é treinado sobre essa representação residualizada. O objetivo é reconstruir a entrada normalizada enquanto impõe esparsidade (retendo apenas os $K$ latentes ativos superiores por lote). Isso incentiva os latentes esparsos a capturar estrutura que existe além da deriva temporal linearmente previsível.

4. Interpretação no Espaço de Ativação

Um componente chave do método é mapear as características aprendidas de volta ao espaço de ativação original para analisá-las como trajetórias:

1. Direções do Decodificador: Os vetores decodificadores do SAE estão inicialmente no espaço residualizado normalizado.
2. Desnormalização: Escalas são aplicadas para retornar ao espaço residual normalizado.
3. Integração de Dinâmicas Lineares: As direções residuais são mapeadas de volta através dos preditores lineares ajustados ($W_i$) para reconstruir a direção da trajetória completa no espaço de ativação normalizado.
4. Mapeamento Final: O resultado é desnormalizado para obter a trajetória do decodificador $\{\phi_{k,i}\}$ nas unidades de ativação originais. Isso permite que cada latente seja analisado como uma trajetória de característica evoluindo ao longo do tempo de remoção de ruído.

Principais Contribuições

1. SAEs Temporais Residualizados: A introdução de um método que separa trajetórias de ativação de difusão em componentes linearmente previsíveis e resíduos, treinando SAEs especificamente sobre estes últimos para evitar modelar dinâmicas temporais redundantes.
2. Interpretação do Decodificador no Espaço de Ativação: Uma derivação que mapeia direções decodificadoras residualizadas de volta ao espaço de ativação original, permitindo a análise de cada latente como uma trajetória de característica espaço-temporal.
3. Caracterização Empírica: Uma avaliação abrangente no Stable Diffusion 1.5 demonstrando que essas representações capturam direções estruturadas temporalmente e localizadas espacialmente com efeitos generativos semanticamente significativos.

Resultados e Avaliação

Desempenho de Reconstrução

• Superioridade da Residualização: Sob um orçamento fixo de esparsidade ao nível da trajetória, as variantes residualizadas superam consistentemente as contrapartes não residualizadas (tanto concatenadas quanto por instante de tempo) no Erro Quadrático Médio (MSE) e Variância Explicada (EV).
• Eficiência: O modelo Concatenado Residualizado alcança desempenho de reconstrução competitivo enquanto usa um único código esparso para toda a trajetória, whereas as baselines por instante de tempo exigem códigos distribuídos.
• Comparação com Baseline: O método supera uma baseline de SAE Matryoshka (que usa trajetórias normalizadas por componente, não residualizadas) na qualidade de reconstrução.

Análise de Estrutura Espaço-Temporal

• Especialização Temporal: Os latentes exibem perfis temporais distintos. Características "iniciais" tendem a ter suporte temporal amplo e alta auto-similaridade entre instantes de tempo, sugerindo que capturam conceitos estruturais persistentes. Características "finais" são mais específicas ao instante de tempo e menos estáveis temporalmente.
• Localização Espacial:
  • Características iniciais mostram suporte espacial positivo mais amplo (maior entropia espacial), frequentemente cobrindo regiões difusas ou semelhantes a bordas.
  • Características médias e finais são mais concentradas espacialmente (menor entropia), localizando-se em áreas específicas semelhantes a objetos ou regiões salientes.

Experimentos de Direcionamento (Steering)

• Direcionamento de Característica Única: A intervenção com latentes específicos produz efeitos generativos visíveis. Características iniciais afetam a estrutura global, enquanto características médias podem localizar mudanças (por exemplo, adicionar joias a um torso). Características finais mostraram efeitos semânticos menos claros em alguns exemplos.
• Transferência de Característica: Usando o conjunto de dados RIEBench, os autores realizaram transferência de característica de origem para destino. O modelo Concatenado Residualizado alcançou o melhor equilíbrio entre eficiência de edição (melhorando a similaridade CLIP com o prompt de destino) e preservação da estrutura da imagem de origem (minimizando a distância LPIPS). Superou variantes não residualizadas e por instante de tempo na eficiência de edição.

Significado e Alegações

O artigo afirma que os SAEs Temporais Residualizados fornecem um framework útil para estudar ativações de difusão estruturadas temporalmente. Ao remover explicitamente a deriva linearmente previsível, o método permite que latentes esparsos se concentrem na estrutura não redundante e semanticamente significativa do processo de geração.

Os autores enfatizam que sua abordagem difere de trabalhos anteriores ao aprender dicionários esparsos sobre trajetórias de ativação residualizadas em vez de instantes de tempo isolados ou representações temporais não residualizadas. A capacidade de mapear esses latentes de volta ao espaço de ativação como trajetórias permite uma compreensão mais profunda de quando e onde características específicas são expressas durante o processo de remoção de ruído.

Limitações notadas pelos autores:

• Os experimentos focam em uma única fonte de ativação do U-Net (atenção cruzada do bloco médio) no Stable Diffusion 1.5.
• Os experimentos de direcionamento são qualitativos; o artigo não alega desempenho otimizado de edição de imagem, mas sim usa o direcionamento para validar a interpretabilidade e a influência causal das características aprendidas.