Resumo Técnico: AMOR (Roteador de Saída Metacognitivo Adaptativo)

Declaração do Problema

Modelos híbridos de atenção recorrente visam combinar a eficiência de inferência dos modelos de espaço de estado recorrentes lineares (SSMs) com o poder expressivo dos mecanismos de atenção. No entanto, arquiteturas híbridas existentes geralmente dependem de estratégias de cronograma fixo, onde a atenção é aplicada uniformemente em todos os tokens ou em camadas predeterminadas (por exemplo, a cada $N$-ésima camada). Essa abordagem incorre no custo quadrático total da atenção ($O(N^2)$) mesmo quando o estado recorrente sozinho é suficiente para prever o próximo token com alta confiança. Consequentemente, esses modelos não aproveitam totalmente os ganhos de eficiência dos SSMs, particularmente em cenários de contexto longo, onde a sobrecarga computacional de blocos de atenção desnecessários torna-se um gargalo. Além disso, híbridos fixos frequentemente degradam-se sob mudanças de distribuição (por exemplo, extrapolação de comprimento) porque forçam mecanismos de atenção que podem não estar calibrados para a dificuldade específica da entrada.

Metodologia: Arquitetura AMOR

Os autores propõem o AMOR (Roteador de Saída Metacognitivo Adaptativo), uma arquitetura híbrida post-hoc que aloca atenção dinamicamente com base na incerteza preditiva em vez de posições fixas.

Arquitetura Central

O AMOR é construído sobre uma espinha dorsal recorrente padrão (testada com Mamba2 e Gated DeltaNet) e anexa $K$ (tipicamente 3) blocos de atenção com portão de entropia após as camadas recorrentes, mas antes da cabeça final do LM.

• Espinha Dorsal: O misturador recorrente processa toda a sequência, mantendo um estado de tamanho fixo.
• Blocos Post-Hoc: Diferentemente de híbridos em série que substituem camadas recorrentes, os blocos do AMOR são aditivos. Eles operam na saída normalizada da espinha dorsal.
• Assimetria: O primeiro bloco AMOR (Bloco 0) opera diretamente no resíduo final normalizado da espinha dorsal (pós-norm), enquanto os blocos subsequentes (1 e 2) utilizam uma formulação pré-norm padrão. Esse design preserva a espinha dorsal como um modelo de linguagem válido e autônomo, tratando os blocos AMOR como refinamentos condicionais.

O Mecanismo de Portão de Entropia

A inovação central é um mecanismo de portão rígido, sem gradiente, inspirado na distinção entre Sistema 1 e Sistema 2 na ciência cognitiva (Kahneman, 2011).

1. Estimativa de Incerteza: Para cada token, a cabeça compartilhada do LM computa logits. A entropia de saída normalizada ($H_t$) é calculada a partir da distribuição softmax.
2. Limiar Adaptativo: Um limiar dinâmico $\tau$ é computado como uma média móvel exponencial (EMA) da mediana de entropia do lote ($\mu$) mais um desvio padrão escalado ($\kappa\sigma$):
   $$\tau = \mu + \kappa\sigma$$
   Isso permite que o limiar acompanhe as mudanças na distribuição de incerteza do modelo durante o treinamento, sem exigir parâmetros de roteador aprendidos.
3. Execução Condicional:
   • Se $H_t > \tau$, o portão dispara ($g_t=1$): O token sofre atenção causal (SDPA) com RoPE, e a saída é adicionada ao resíduo com um fator de escala aprendido por canal $\alpha$.
   • Se $H_t \le \tau$, o portão permanece inativo ($g_t=0$): O cálculo de atenção é pulado, e o token passa inalterado.
4. Fluxo de Gradiente: O portão é não diferenciável. Os gradientes fluem através dos pesos de atenção ($W_Q, W_K, W_V, W_O$) apenas nas posições de disparo. Crucialmente, $W_K$ e $W_V$ recebem gradientes de futuras posições de disparo para garantir que o cache KV permaneça válido para recuperação, mesmo que o token atual tenha pulado a atenção.

Treinamento e Inferência

• Treinamento: O AMOR é treinado em um modo "denso com máscara" usando FlashAttention-2 para eficiência, matematicamente equivalente à execução esparsa verdadeira. O limiar $\tau$ é congelado na inferência.
• Inferência: O modelo realiza um pulo condicional. Quando o portão está desligado, a projeção $Q$ cara, o cálculo de atenção e a projeção $O$ são contornados, economizando computação. No entanto, as projeções $K$ e $V$ e as atualizações de cache sempre ocorrem para manter a continuidade do contexto.

Contribuições Principais

1. Roteamento Impulsionado por Incerteza: O AMOR introduz um mecanismo simples e post-hoc que invoca atenção seletivamente apenas quando a espinha dorsal recorrente exibe alta incerteza preditiva (alta entropia), em vez de depender de cronogramas de camadas fixos.
2. Hibridização Eficiente em Parâmetros: O método adiciona sobrecarga mínima de parâmetros (aprox. <5% para os blocos de portão e atenção) enquanto alcança desempenho comparável ou superior aos híbridos de cronograma fixo.
3. Robustez a Mudanças de Distribuição: Ao confiar na espinha dorsal recorrente para tokens de baixa incerteza, o AMOR mitiga a degradação de desempenho observada em híbridos fixos durante a extrapolação de comprimento, onde os mecanismos de atenção frequentemente sofrem com deriva posicional.
4. Portão Sem Gradiente: A abordagem evita a complexidade de roteadores aprendidos ou aprendizado por reforço, utilizando um limiar estatístico derivado da própria distribuição de saída do modelo.

Resultados Experimentais

Os autores avaliaram o AMOR em três escalas (180M, 440M, 1,5B de parâmetros) usando espinhas dorsais Mamba2 e Gated DeltaNet, comparando com modelos recorrentes puros, Transformers puros e híbridos em série/fundidos de cronograma fixo.

• Raciocínio de Senso Comum: O AMOR consistentemente igualou ou superou tanto as espinhas dorsais recorrentes puras quanto os híbridos de cronograma fixo em todas as escalas e benchmarks (LAMBADA, HellaSwag, PIQA, ARC, etc.).
• Recuperação em Contexto: Enquanto modelos recorrentes puros lutam com recuperação de cópia exata, o AMOR melhorou substancialmente sobre sua espinha dorsal, alcançando resultados competitivos com híbridos fixos enquanto ativava atenção em apenas ~22% dos tokens.
• Desempenho em Contexto Longo: No benchmark LongBench, o AMOR manteve desempenho estável à medida que o comprimento do contexto aumentava. Em contraste, híbridos de cronograma fixo e Transformers puros degradaram-se significativamente devido a mudanças na distribuição do RoPE e restrições de memória. A capacidade do AMOR de pular a atenção em tokens previsíveis permitiu que ele escalasse de forma mais graciosa.
• Eficiência de Inferência: O AMOR demonstrou melhor escalabilidade de taxa de transferência de decodificação em regimes de contexto longo. Análise teórica e benchmarks empíricos (até 448K de contexto) mostraram que as economias de pular a atenção superam a sobrecarga do cálculo de entropia assim que o comprimento do contexto excede um certo ponto de cruzamento (aprox. 256K para modelos de 440M).

Significado e Alegações

O artigo argumenta que quando a atenção é aplicada é tão crítico quanto quanto é aplicada. Os autores afirmam que:

1. Eficiência e Expressividade não são uma troca: Ao tratar a alocação de computação como um problema de roteamento baseado na dificuldade da entrada, os modelos podem alcançar alta expressividade sem incorrer no custo total da atenção.
2. Simplicidade é eficaz: Roteadores aprendidos complexos ou objetivos auxiliares não são necessários; um simples portão de entropia orientado à saída pode coordenar efetivamente arquiteturas híbridas.
3. Computação Adaptativa: O AMOR aponta para uma classe de modelos que correspondem dinamicamente a computação à dificuldade da entrada, preservando as forças de espinhas dorsais de espaço de estado leves enquanto implantam seletivamente mecanismos pesados de atenção apenas quando o modelo está incerto.

Os resultados sugerem que a alocação seletiva de atenção baseada na incerteza preditiva oferece uma alternativa robusta a estratégias de roteamento uniformes ou fixas, particularmente para tarefas de geração e recuperação de contexto longo.