Resumo Técnico: MESANET: MODELAGEM DE SEQUÊNCIAS POR TREINAMENTO DE TESTE LOCALMENTE ÓTIMO

Declaração do Problema

A modelagem de sequências é atualmente dominada por arquiteturas transformer causais que utilizam autoatenção softmax. Embora eficazes, esses modelos sofrem com o escalonamento linear nos requisitos de memória e computação em relação ao comprimento da sequência durante a inferência. Esforços recentes buscaram abordar isso linearizando a operação softmax, resultando em modelos de redes neurais recorrentes (RNNs) com custos constantes de memória e computação (ex: DeltaNet, Mamba, xLSTM). Esses modelos existentes são frequentemente unificados pela perspectiva de que sua dinâmica recorrente deriva de um objetivo de regressão em contexto, otimizado aproximadamente via regras de aprendizado online (ex: passos de gradiente descendente). No entanto, essas abordagens existentes dependem tipicamente de atualizações incrementais (métodos de primeira ordem) em vez de resolver para os pesos ótimos em cada ponto temporal, o que pode limitar sua capacidade de se adaptar rapidamente a novas evidências enquanto retêm informações passadas. Além disso, iterações anteriores de camadas de treinamento de teste otimais (como a camada Mesa original) foram limitadas pela instabilidade numérica ao usar portões de esquecimento (forget gates) e pela falta de paralelizabilidade, impedindo o treinamento escalável.

Metodologia

Os autores introduzem o MesaNet, uma arquitetura de modelo de linguagem centrada em uma nova camada Mesa, numericamente estável e paralelizável por blocos (chunkwise). Esta camada é projetada para realizar um treinamento de teste localmente ótimo em cada ponto temporal.

A Camada Mesa

Ao contrário das RNNs lineares padrão que atualizam seu estado via um único passo de gradiente descendente (ex: DeltaNet), a camada Mesa define o mapa linear $\Phi_t$ como a solução de um problema de otimização de teste-tempo. Especificamente, ela minimiza uma perda de erro quadrático acumulada e regularizada sobre todos os dados vistos até o momento ($e_{1:t}$):
$$\hat{\Phi}^{mesa}_t = \arg \min_{\Phi} \left( \frac{1}{2} \sum_{t'=1}^t \zeta_{tt'} \|v_{t'} - \Phi k_{t'}\|^2 + \frac{1}{2} \text{Tr}(\Phi \Lambda \Phi^\top) \right)$$
onde $\zeta_{tt'}$ representa fatores de esquecimento cumulativos, e $\Lambda$ é um regularizador simétrico definido positivo.

A saída é computada como $\Delta^{mesa}_t = \hat{\Phi}^{mesa}_t q_t$. Crucialmente, em vez de usar mínimos quadrados recursivos (que sofrem de instabilidade numérica com portões de esquecimento e não são paralelizáveis), os autores resolvem para a resposta de consulta ótima $q^*_t = (H_t + \Lambda)^{-1} q_t$ usando o método de Gradiente Conjugado (CG).

• Paralelização: Os autores aproveitam a equivalência entre o método CG sobre múltiplos passos temporais e a Atenção Linear com Portão (Gated Linear Attention - GLA). Ao manter variáveis de estado $G_t$ e $H_t$ e resolver o sistema linear em paralelo sobre blocos da sequência, a camada alcança complexidade $O(T)$ durante o treinamento e $O(1)$ durante a inferência.
• Computação Dinâmica: O número de iterações de CG ($k$) necessárias para atingir um critério de parada varia por sequência e por token. Isso permite que o modelo aloque dinamicamente o computação de teste-tempo, trocando o custo de inferência por desempenho.

Arquitetura (MesaNet)

O MesaNet adota um backbone de transformer padrão apenas-decodificador (semelhante ao Llama 2), mas substitui o componente de mistura de sequência (tipicamente Multi-Head Attention) pela camada Mesa. A arquitetura inclui:

• Blocos Residuais: Empilhamento de $N$ blocos contendo uma camada de mistura de sequência (Mesa) e uma camada de mistura de canal (MLP SwiGLU).
• Pré-processamento: Chaves (keys) e consultas (queries) são normalizadas e convoluídas com uma pequena janela antes da camada Mesa.
• Variante Híbrida: Um modelo "Hawk-Mesa" também é explorado, alternando entre camadas LRU (Hawk) e camadas Mesa para combinar as forças de modelagem de sequências curtas e longas.

Principais Contribuições

1. Uma Camada Mesa Paralelizável e Estável: O artigo apresenta uma versão da camada Mesa que supera o gargalo sequencial e a instabilidade numérica da formulação original. Utiliza um solver CG com paralelização por blocos, permitindo o treinamento eficiente em hardware moderno (GPUs/TPUs) enquanto mantém os benefícios teóricos da regressão de teste-tempo ótima.
2. Alocação de Computação de Teste-Tempo Dinâmica: A camada invoca explicitamente um solver externo onde o número de iterações depende da dificuldade do sistema linear para um dado sequência. Isso fornece um mecanismo principiado para variar a computação de inferência baseada na complexidade do input.
3. Avaliação Empírica Abrangente: Os autores treinam MesaNets nas escalas de 140M, 440M e 1B de parâmetros no dataset SlimPajama. Eles fornecem uma comparação rigorosa, "maçãs com maçãs", contra Mamba2, xLSTM, DeltaNet, Gated DeltaNet e Transformers, controlando arquitetura, tokenizer e hiperparâmetros.

Resultados

• Desempenho de Modelagem de Linguagem: Os MesaNets alcançam uma perplexidade de validação (PPL) inferior a outras baselines de RNN (Mamba2, xLSTM, DeltaNet) em todas as escalas. Na escala de 1B de parâmetros, o MesaNet iguala ou supera ligeiramente o desempenho de uma baseline Transformer com a mesma arquitetura.
• Análise de Posição de Token: Uma análise detalhada da Perplexidade Negativa de Log-Verossimilhança (NLL) em relação à posição do token revela um padrão distinto:
  • Tokens Iniciais: A maioria das RNNs, incluindo o MesaNet, supera os Transformers nos tokens iniciais de uma sequência.
  • Contextos Longos: Enquanto os Transformers se destacam em tokens posteriores, o MesaNet e o Hawk-Mesa estendem sua vantagem significativamente mais (além de 512 tokens) em comparação com outras RNNs, embora eventualmente fiquem atrás dos Transformers em profundidades muito longas.
• Benchmarks Downstream:
  • Raciocínio Global: O MesaNet supera todas as outras RNNs em benchmarks que exigem dependências de longo alcance (ex: Lambada, HellaSwag, RACE), embora ainda fique atrás da baseline Transformer.
  • Recall em Contexto (In-Context Recall): O MesaNet supera outras RNNs e até mesmo um modelo de Atenção de Janela Deslizante (SWA) com uma janela de 1024 tokens, mas permanece atrás do Transformer de atenção total.
  • Aprendizado Few-Shot: O MesaNet demonstra forte desempenho em tarefas de embaralhamento de palavras (word scrambling) e tradução, frequentemente excedendo os Transformers em tarefas de embaralhamento.
• Eficiência: Apesar de exigir FLOPs adicionais para o solver CG durante a inferência, o MesaNet mantém um throughput de tokens competitivo em GPUs H100 e TPUv5s, particularmente quando se utiliza um número moderado de passos de CG (ex: 10-30).

Significância e Alegações

O artigo posiciona o MesaNet como um passo significativo em direção ao treinamento de teste-tempo ótimo dentro de redes neurais. Os autores alegam que, ao resolver problemas de otimização sequencial (especificamente, encontrar os melhores pesos de memória associativa linear) durante o tempo de inferência, o modelo pode se adaptar mais efetivamente a novas evidências do que modelos que usam simples atualizações de gradiente.

O trabalho destaca um compromisso (trade-off): os ganhos de desempenho em perplexidade e raciocínio downstream vêm ao custo de maior computação de inferência (FLOPs). Os autores sugerem que isso se alinha com as tendências recentes de aumentar a computação de teste-tempo para melhorar o desempenho. Eles argumentam que a camada Mesa oferece um mecanismo flexível para alocar essa computação dinamicamente, potencialmente resolvendo problemas sequenciais complexos dentro da própria rede.

O artigo conclui que, embora os MesaNets enfrentem atualmente um overhead de computação, sua capacidade de superar outras RNNs em compreensão de contexto longo e tarefas de raciocínio global sugere que o treinamento de teste-tempo localmente ótimo é uma direção promissora para a modelagem de sequências eficiente, potencialmente preenchendo a lacuna entre a eficiência das RNNs e o desempenho dos Transformers. Os autores reconhecem que o custo de computação da implementação atual é uma deficiência e sugerem que trabalhos futuros poderiam envolver estratégias de "warm-start" para otimizações ou estratégias híbridas de inferência para mitigá-lo.