Resumo Técnico: Cubit – Misturador de Tokens com Regressão de Ridge com Kernel

Declaração do Problema

Desde a introdução da arquitetura Transformer em 2017, ela tornou-se o paradigma dominante em aprendizado profundo para processamento de linguagem natural, visão computacional e tarefas multimodais. No entanto, o mecanismo central de mistura de tokens nos Transformers baseia-se em autoatenção, que pode ser teoricamente interpretado como realizando regressão Nadaraya–Watson (NW) com um kernel exponencial (via softmax) e normalização L1.

Embora eficaz, essa abordagem baseada em NW possui limitações inerentes:

1. Complexidade Computacional: A atenção padrão incorre em custos computacionais quadráticos ($O(N^2)$) em relação ao comprimento da sequência, limitando a escalabilidade para sequências longas.
2. Restrições Teóricas: A regressão NW carece de regularização explícita para o compromisso viés-variância e pode ser sensível a ruídos e efeitos de fronteira.
3. Escalabilidade de Melhorias: Embora numerosas modificações (atenção esparsa, atenção linear, mistura de especialistas) tenham sido propostas, o paradigma computacional fundamental permanece enraizado na regressão NW.

Os autores postulam que substituir a fundação da regressão NW por uma estrutura estatística mais robusta poderia produzir arquiteturas com fundamentos matemáticos mais sólidos e capacidades aprimoradas de modelagem de sequências longas.

Metodologia: A Arquitetura Cubit

O artigo propõe o Cubit, uma arquitetura inovadora que substitui a regressão NW do módulo de atenção por Regressão de Ridge com Kernel (KRR).

1. Fundamentação Teórica: De NW para KRR

• Nadaraya–Watson (Transformer): Na atenção padrão, o token de saída $z_i$ é uma soma ponderada dos vetores de valor $v_j$, onde os pesos são determinados por similaridades normalizadas por softmax entre consultas e chaves. Isso corresponde à regressão NW, que estima um valor de função com base em uma média ponderada local sem resolver um sistema global.
• Regressão de Ridge com Kernel (Cubit): A KRR resolve um problema de minimização de risco empírico regularizado. A previsão para uma nova entrada envolve o inverso da matriz de kernel mais um termo de regularização: $(K + \lambda I)^{-1}$.
  • Diferença Chave: Diferentemente da NW, a KRR introduz um parâmetro de regularização explícito $\lambda$ para controlar o compromisso viés-variância e utiliza o inverso da matriz de Gram para acoplar informações entre todas as amostras de dados, oferecendo maior robustez a ruídos e efeitos de fronteira.

2. Design Arquitetural

O Cubit modifica o cálculo clássico de atenção incorporando a solução de forma fechada da KRR:

• Construção do Kernel: A matriz de kernel é construída usando mapas de características derivados de incorporações de chave. Os autores parametrizam o mapa de características $\psi(x)$ usando uma matriz de projeção treinável $W_K$ e uma função de ativação prescrita $\phi$.
• Normalização via Matriz Inversa: Em vez de normalização por softmax linha a linha, o Cubit computa uma matriz de normalização $\Sigma^{(h)} = (K^{(h)}K^{(h)\top} + \lambda I)^{-1}$. Esta matriz atua como um operador de acoplamento global.
• Formulação de Saída: A saída é computada como $Z^{(h)} = A^{(h)} \Sigma^{(h)} V^{(h)}$, onde $A^{(h)}$ representa a matriz de similaridade de kernel (interações consulta-chave) e $\Sigma^{(h)}$ é a matriz de normalização inspirada em KRR.
• Restrição Causal: Para lidar com a geração autoregressiva, a construção do kernel é adaptada para que o $i$-ésimo token atenda apenas a si mesmo e aos tokens anteriores, garantindo que a regressão seja realizada no contexto disponível.

3. Redimensionamento de Alcance Limitado (LRR)

Para abordar problemas de estabilidade de treinamento inerentes à inversão de grandes matrizes, os autores introduzem o Redimensionamento de Alcance Limitado (LRR):

• Um vetor de escala treinável $s^{(h)}$ é aplicado para ajustar adaptativamente a precisão do kernel.
• Os fatores de escala são restringidos dentro de um intervalo limitado $(\beta, \alpha + \beta)$ usando uma função sigmoide.
• Isso garante a invertibilidade da matriz de escala e limita os elementos inversos, prevenindo instabilidade numérica durante a solução do sistema linear.

4. Distinção em Relação a Trabalhos Relacionados

O artigo distingue explicitamente o Cubit do DeltaFormer (que também utiliza matrizes inversas):

• Motivação: O DeltaFormer combina DeltaNet e Transformers; o Cubit é motivado pela mudança teórica de NW para KRR.
• Implementação: O DeltaFormer usa duas incorporações diferentes para construir o inverso e força valores diagonais para 1. O Cubit usa uma única incorporação específica (por exemplo, apenas $K$) para garantir invertibilidade naturalmente, sem forçar valores diagonais para 1.
• Estabilidade: O Cubit propõe LRR para melhorar a estabilidade, um componente ausente no DeltaFormer.

Contribuições Principais

1. Estrutura Teórica Unificada: Os autores estabelecem uma conexão entre mecanismos de mistura de tokens e métodos clássicos de regressão, analisando sistematicamente a regressão NW e a KRR para fornecer uma nova perspectiva para o desenvolvimento de arquiteturas de modelos.
2. Arquitetura Cubit: A proposta de um misturador de tokens baseado em KRR que substitui o estimador NW nos Transformers, apresentando o mecanismo Redimensionamento de Alcance Limitado (LRR) para aprimorar a estabilidade de treinamento.
3. Validação Empírica: Experimentos abrangentes demonstrando que o Cubit exibe capacidades de modelagem de sequências longas mais fortes do que Transformers padrão e DeltaFormer, com ganhos de desempenho aumentando conforme o comprimento da sequência cresce.

Resultados Experimentais

Os autores avaliaram o Cubit contra Transformers e DeltaFormer em diversos conjuntos de dados (Arxiv, Books3, FinWeb-Edu), comprimentos de sequência (512 a 8192) e tamanhos de modelo (125M a 1,3B parâmetros).

• Desempenho em Modelagem de Linguagem: O Cubit consistentemente alcançou menor perplexidade (perda de validação) do que as linhas de base. Por exemplo, no conjunto de dados Arxiv (modelo 125M, 512 tokens), o Cubit alcançou uma perda de 1,8802 comparado a 1,9003 para o Transformer.
• Modelagem de Sequências Longas: A lacuna de desempenho entre Cubit e Transformer alargou-se conforme o comprimento da sequência de treinamento aumentou. Em 8192 tokens, a lacuna de perda estabilizou-se em 0,05–0,06, sugerindo que o mecanismo baseado em kernel do Cubit captura dependências de longo alcance mais efetivamente.
• Propriedades de Escala: O Cubit demonstrou comportamento de escala favorável. Enquanto a vantagem do DeltaFormer sobre o Transformer diminuiu conforme o tamanho do modelo aumentou (de 350M para 1,3B), os ganhos de desempenho do Cubit persistiram e melhoraram.
• Tarefas Downstream: Em avaliações zero-shot (ARC, HellaSwag, PIQA, etc.), o Cubit consistentemente superou tanto Transformer quanto DeltaFormer em diferentes tamanhos de dados (10B a 50B tokens) e escalas de modelo.
• Estudos de Ablação:
  • Design Arquitetural: Os ganhos de desempenho foram atribuídos ao design estrutural (formulação KRR) e não à expansão de parâmetros.
  • Eficácia do LRR: A inclusão do LRR forneceu melhorias mensuráveis, com "Cubit-NoLRR" ainda superando o Transformer padrão, mas "Cubit" (com LRR) alcançando os melhores resultados.

Significado e Alegações

O artigo alega que o Cubit oferece uma fundação matemática mais forte do que o Transformer vanilla ao aproveitar os benefícios teóricos da Regressão de Ridge com Kernel, incluindo regularização explícita, taxas de convergência mais rápidas em Espaços de Hilbert de Kernel Reprodutor (RKHS) e robustez a ruídos.

Os autores sugerem que o Cubit pode representar uma arquitetura de modelo fundamental de próxima geração, particularmente para tarefas que exigem compreensão de contexto longo. Eles notam que os ganhos de desempenho não são meramente devidos ao aumento de tokens de treinamento, mas derivam da eficácia arquitetural em sequências longas. Embora o artigo reconheça que o potencial para abuso de modelos mais poderosos exista, ele enfatiza os benefícios sociais potenciais de capacidades de modelagem aprimoradas.

O trabalho também aponta para direções futuras, como explorar Regressão Linear Local (que aborda viés de fronteira, mas é computacionalmente mais lenta) e outras estruturas baseadas em regressão para mistura de tokens, indicando que a perspectiva baseada em regressão sobre atenção é um terreno fértil para inovação arquitetural adicional.