Efficient Autoregressive Inference for Transformer Probabilistic Models

1. O Problema

Os modelos baseados em transformadores para inferência probabilística amortizada e meta-aprendizado (como Processos Neurais, Redes Ajustadas a Priori e Modelos de Fundação Tabulares) são excelentes para prever distribuições marginais em uma única passagem. No entanto, muitas aplicações práticas exigem distribuições conjuntas coerentes sobre múltiplas previsões (ex: interpolação de sinais, modelagem de dados comportamentais, previsão de colunas em tabelas).

A abordagem padrão para obter distribuições conjuntas nesses modelos é o uso autoregressivo (AR): gerar previsões sequencialmente, onde cada nova previsão condiciona-se ao contexto original mais todas as previsões anteriores.

• O Gargalo: Como os codificadores baseados em conjuntos (set-based) usam atenção bidirecional para manter a invariância de permutação, adicionar uma nova previsão ao conjunto de condicionamento força o re-codificação completa de todo o conjunto aumentado a cada passo.
• Complexidade: Isso resulta em uma complexidade computacional de $O(K(N+K)^2)$, onde $N$ é o tamanho do contexto e $K$ é o número de previsões. Para contextos grandes ou sequências longas, isso torna a amostragem e a avaliação de densidade conjunta extremamente lentas e custosas em termos de memória.

2. Metodologia: O Buffer Autoregressivo Causal

Os autores propõem uma nova arquitetura que combina as vantagens dos modelos baseados em conjuntos (flexibilidade e invariância de permutação) com a eficiência das arquiteturas puramente autoregressivas (caching de chaves/valores). A solução central é o Buffer Autoregressivo Causal.

Mecanismo Principal

O modelo desacopla a codificação do contexto estático da geração sequencial de alvos:

1. Codificação Única do Contexto: O conjunto de contexto $C$ é codificado uma única vez, e suas representações (chaves e valores) são cacheadas e tratadas como imutáveis (somente leitura).
2. Buffer Causal: As previsões geradas são armazenadas em um "buffer" separado. Este buffer utiliza um padrão de atenção causal estrito, onde cada novo token no buffer só pode atender aos tokens anteriores no buffer e ao contexto cacheado.
3. Decodificação Eficiente: Para prever o próximo alvo, o decodificador atende ao contexto cacheado (fresco) e ao prefixo visível do buffer. Isso evita a necessidade de re-codificar o contexto a cada passo.

Complexidade Computacional

Ao separar o contexto do histórico autoregressivo, a complexidade de atenção é reduzida drasticamente:

• Método Antigo (Re-codificação): $O(K(N+K)^2)$
• Método Proposto (Buffer): $O(N^2 + NK + K^2)$
  • $O(N^2)$: Codificação única do contexto.
  • $O(NK)$: Leitura cruzada do contexto cacheado para cada alvo.
  • $O(K^2)$: Atenção causal dentro do buffer.

Treinamento Unificado

O modelo é treinado com uma estratégia unificada que utiliza máscaras de atenção estruturadas e um "currículo" de tamanho de buffer:

• Durante o treinamento, os alvos podem atender apenas ao contexto (50% dos casos) ou ao contexto + um prefixo variável do buffer (50% dos casos).
• Isso permite que um único modelo aprenda tanto a previsão marginal (quando o buffer está vazio) quanto a dependência autoregressiva, sem custo adicional significativo de treinamento.

Avaliação de Densidade Conjunta

O método permite a avaliação da densidade logarítmica conjunta de $K$ alvos em uma única passagem forward (one-pass), empacotando tokens de buffer e tokens de consulta com uma máscara causal, em vez de exigir $K$ passagens sequenciais.

3. Principais Contribuições

1. Introdução do Buffer Autoregressivo Causal: Um mecanismo arquitetônico que reduz o custo da inferência conjunta autoregressiva de $O(K(N+K)^2)$ para $O(N^2 + NK + K^2)$, permitindo amostragem eficiente e avaliação de densidade em uma única passagem.
2. Estratégia de Treinamento Unificada: Combina atenção mascarada com um currículo de tamanho de buffer, permitindo que um único modelo domine tanto a previsão independente quanto a dependente com overhead mínimo.
3. Generalidade e Desempenho: Demonstrado em diversos modelos (TNPs, PFNs e Modelos de Fundação Tabulares como TabICL), alcançando até 20x mais rápido na amostragem conjunta e 7x menos uso de memória (VRAM) em comparação com métodos autoregressivos padrão, mantendo a precisão preditiva.

4. Resultados Experimentais

Os autores avaliaram o método em quatro domínios distintos:

• Funções Sintéticas e Séries Temporais EEG:
  • O método com buffer ($K=16$) alcançou desempenho quase idêntico ao método autoregressivo completo (TNP-D-AR) em tarefas de interpolação e previsão.
  • Redução de tempo de amostragem de 3x a 20x e redução de uso de memória de 6x a 7x para contextos grandes ($N=1024$).
• Comparação de Modelos Bayesianos (Inferência Causal Multissensorial):
  • O método foi usado para calcular a Verossimilhança Marginal (LML) para seleção de modelos.
  • Os resultados mostraram uma correlação quase perfeita ($R^2 \approx 1.0$) com a "verdade fundamental" (ground-truth) obtida via métodos variacionais caros, superando todas as outras baselines em precisão de comparação de modelos.
• Modelos de Fundação Tabulares (TabICL):
  • Integrado ao TabICL, o método com buffer manteve a precisão da previsão autoregressiva padrão em tarefas de interpolação e previsão em 6 conjuntos de dados reais (UCI/Kaggle), superando significativamente a previsão independente.
• Eficiência de Treinamento:
  • O tempo de treinamento por passo foi comparável aos métodos mais rápidos (TNP-D) e 4x a 12x mais rápido que o TNP-A (que sofre com re-codificação e duplicação de tokens).

5. Significado e Impacto

Este trabalho resolve um gargalo fundamental na aplicação de modelos probabilísticos baseados em transformadores para tarefas que exigem modelagem de dependências conjuntas.

• Viabilidade Prática: Torna viável o uso de inferência autoregressiva em contextos grandes e em lotes (batched), algo que era proibitivo devido ao custo computacional e de memória.
• Flexibilidade: Oferece um "meio-termo" ideal entre a eficiência da previsão independente e a expressividade da previsão conjunta, sem sacrificar a invariância de permutação do contexto.
• Aplicabilidade Geral: A abordagem é modular e pode ser aplicada a qualquer arquitetura baseada em transformadores que utilize condicionamento por conjuntos, incluindo modelos de fundação tabulares emergentes, abrindo caminho para aplicações em ciência de dados, neurociência computacional e otimização bayesiana em tempo real.

Em resumo, o artigo apresenta uma inovação arquitetônica que permite que modelos probabilísticos modernos escalem para cenários de inferência conjunta complexa, oferecendo ganhos de desempenho de ordem de magnitude sem perda de qualidade preditiva.