Resumo Técnico: IG-Lens

Definição do Problema

Transformers do tipo decoder-only são frequentemente concebidos como processos de inferência iterativos, onde o fluxo residual refina uma estimativa contínua do próximo token. Uma questão de interpretabilidade crítica surge: em qual camada específica a probabilidade de um token previsto é realmente decidida?

Ferramentas existentes para leitura camada a camada (layer-wise readout) falham em responder a isso precisamente no espaço de probabilidade:

• Logit Lens & Tuned Lens: Estes relatam o nível de probabilidade em cada camada, mas não fornecem uma decomposição aditiva. Suas estimativas são enviesadas e não monótonas ao longo da profundidade, o que significa que a soma das probabilidades por camada não possui significado matemático.
• Atribuição Direta de Logit (DLA) & Métodos de Embedding: Estes são aditivos, mas operam no espaço de logits. Como a função softmax é altamente não linear, uma decomposição aditiva de logits não se traduz em uma decomposição aditiva de probabilidades, que é a grandeza de interesse prático.
• Condutância de Camada & Gradiente Integrado de Camada (IG): Estes métodos integram gradientes por camada, mas atribuem cada camada ao seu próprio baseline independente. Consequentemente, os termos por camada não "telescopam" (somam-se) para a mudança total na probabilidade de previsão.

Metodologia: IG-Lens

Os autores introduzem o IG-Lens, um método que fornece uma decomposição aditiva exata da probabilidade através das camadas do transformer.

Conceito Central

O IG-Lens aplica uma aplicação telescópica de Gradientes Integrados (IG) ao longo de um único caminho através dos estados ocultos de um baseline até a camada final.

1. Definição do Caminho: O método seleciona uma sequência de camadas $0 = L_0 < L_1 < \dots < L_k = n$ e define um caminho de estados ocultos $h_{L_0} \to h_{L_1} \to \dots \to h_{final}$.
2. Leitura de Cabeça Única (Head-Once Readout): Crucialmente, a função de leitura $f(h)$ é definida como o unembedding final aplicado uma única vez a um único vetor oculto: $f(h) = \text{softmax}(\text{head}(\text{norm}(h)))_{y_t}$. Esta função não contém atenção ou MLP; toda a mistura entre tokens é assumida como já ocorrida no forward pass que produziu os estados ocultos destacados.
3. Atribuição Telescópica: A contribuição da camada $L_j$ é definida como o segmento IG de $f$ ao longo da linha reta de $h_{L_{j-1}}$ para $h_{L_j}$. Pelo teorema do gradiente, a integral do gradiente sobre um segmento é igual à diferença dos valores dos pontos finais. Somar esses segmentos resulta em uma soma telescópica:
   $$\sum_{j=1}^k IGL_j = f(h_{final}) - f(h_{base}) = \Delta p$$
   Isso garante que a soma das atribuições seja exatamente a mudança na probabilidade alvo, com o softmax e o LayerNorm integrados em vez de linearizados ou congelados.

Estimador Consciente da Predição

Embora a integral seja exata, estimadores de passos finitos geralmente introduzem erro de discretização. O IG-Lens adota um estimador consciente da predição (inspirado no IDGI) como seu padrão:

• Em vez de creditar passos baseando-se em gradientes brutos, ele credita cada passo por sua mudança observada na probabilidade alvo ($f(h^{(s)}) - f(h^{(s-1)})$).
• Para uma leitura de probabilidade escalar, isso colapsa cada segmento em uma soma telescópica de valores de pontos finais.
• Resultado: A completude mantém-se à precisão de ponto flutuante em qualquer contagem de passos ($m \ge 1$). Passos onde o modelo mostra sensibilidade ao gradiente, mas nenhuma mudança na saída, são naturalmente filtrados, eliminando o erro de discretização de Riemann.

Implementação

O método utiliza um algoritmo em lote de passagem única que computa o mapa completo de token por camada sem exigir um backward pass (autograd).

• Ao empilhar todos os pontos de interpolação em um único tensor e avaliar a função pontual $f$ em um único forward pass, o método evita o overhead do cálculo de gradiente.
• Esta eficiência é possível especificamente devido ao design de "cabeça única", que garante a independência entre os pontos de tokens no lote.

Principais Contribuições

1. Decomposição Aditiva Exata: O IG-Lens é o primeiro método a fornecer uma decomposição aditiva exata de probabilidade (não de logits) através das camadas, onde a soma das atribuições é igual a $\Delta p$.
2. Identidade Telescópica: O artigo fornece uma prova formal (Proposição 1) de que as integrais de segmento telescopam para a mudança total na predição, mantendo o softmax dentro do caminho de integração.
3. Estimativa Consciente da Predição: O estimador padrão garante a completude de precisão de máquina, independentemente do número de passos, e filtra passos de "sensibilidade sem resposta".
4. Implementação Eficiente: Uma avaliação de passagem única computa o mapa completo de atribuição sem chamadas de backward, tornando-o computacionalmente eficiente.
5. Posicionamento: Os autores delineiam claramente por que os métodos existentes falham no espaço de probabilidade (aditividade de espaço de logit, viés ou falta de telescopia) e posicionam o IG-Lens como o contraparte de espaço de probabilidade para métodos de espaço de logit.

Resultados e Plano de Avaliação

• Verificação Numérica: Em testes controlados e em um modelo real (Llama-3.2-1B-Instruct), o estimador padrão reproduz as diferenças de pontos finais por segmento para $\sim 10^{-7}$ (precisão float32). A soma telescópica mantém-se exata conforme o ponto flutuante, independente do número de passos $m$.
• Saída Ilustrativa: Em experimentos com o prompt "What is the capital of Vietnam?", o método concentra com sucesso a massa de atribuição nos segmentos de camada específicos onde o valor do logit-lens transita de baixo para alto. A soma das atribuições para cada token coincide com $\Delta p$ até a precisão de ponto flutuante.
• Estratégia de Avaliação: Os autores reconhecem que o "início" (onset) não possui um ground truth intrínseco. Portanto, propõem validar o método via ablação causal: se o IG-Lens identificar a camada $L$ como o início, ablar a camada $L$ deve deprimir a probabilidade alvo mais do que ablar camadas com pontuações de início menores. Este teste é projetado para ser barato de executar e capaz de falsificar a afirmação.

Significância e Alegações

O artigo posiciona o IG-Lens como um refinamento quantitativo da visão de inferência iterativa. Diferente do Logit Lens, que pergunta "qual é a predição em execução nesta camada?", o IG-Lens pergunta "quanto da probabilidade final cada camada é responsável?" e garante que essas responsabilidades somem o todo.

Alegações Modestas e Limitações:

• Marginal Condicional: Os autores declaram explicitamente que o IG-Lens mede uma marginal condicional: a probabilidade adicional que o norm+head pode ler de $h_L$ dado $h_{L-1}$. Ele descarta o efeito de $h_L$ roteado através de blocos de atenção e MLP superiores.
• Compromisso (Trade-off): O método escolhe "soma = $\Delta p$" em vez de "efeito causal total via o forward pass completo". Para efeitos causais totais, a decomposição de componentes (como DLA) continua sendo a ferramenta apropriada.
• Não-Monotonicidade: O método lida com trajetórias de probabilidade não monótonas (onde a probabilidade é construída e parcialmente desfeita) usando atribuições com sinal, oferecendo uma visão mais fiel das contribuições de camada do que métodos baseados em valor absoluto.

A questão central permanece empírica: se o início do IG-Lens prevê a camada na qual a ablação causal mais deprime a predição, particularmente para tokens com trajetórias não monótonas.