Resumo Técnico: Um Único Vetor é Tudo o que Você Precisa para a Atenção O(1): O Estado Oceano (The Ocean State)

1. Problema

Os Transformers padrão sofrem de uma limitação estrutural: a autoatenção incorre em um custo computacional de $O(N^2)$ e requer um cache de memória (cache de Chave-Valor) de $O(N)$ que cresce linearmente com o comprimento do contexto. Embora trabalhos recentes (RetNet, RWKV, Mamba, Titans) tentem mitigar isso introduzindo nova maquinaria, como atenção linear, portões seletivos (gating) ou módulos de memória externa, eles alteram fundamentalmente a arquitetura. Além disso, os Transformers padrão descartam o estado oculto final após cada passagem direta (forward pass), apesar de ele conter uma absorção de todo o contexto.

O artigo postula que o custo $O(N^2)$ é estrutural à implementação atual da atenção, não inerente ao conceito de modelagem de contexto. Ele questiona: E se o estado oculto, que absorve todo o contexto, não fosse descartado, mas retido e usado como o único objeto de atenção para o próximo token?

2. Metodologia: O Estado Oceano (The Ocean State)

A solução proposta, "O Estado Oceano", envolve dois reposicionamentos específicos da arquitetura Transformer padrão sem adicionar novos componentes (por exemplo, sem gating ou memória externa).

2.1 Mecanismos Principais

1. O Estado Oceano (Reposicionamento 1):

   • Em vez de descartar o estado oculto após uma passagem direta, a saída do processamento do token atual substitui o estado oculto anterior, que é denominado "estado oceano".
   • Mecanismo de Atenção: A atenção é computada estritamente entre o token atual e o estado oceano persistente (um único vetor denso).
   • Consequências: Isso elimina a necessidade de caches de Chave-Valor (KV), codificações posicionais e máscaras causais. O estado é atualizado via um "revezamento" (substituição 1:1) em vez de acumulação, garantindo que a informação seja rearranjada, mas não comprimida ou descartada.
   • Somente Leitura: O estado oceano é lido pelo token atual, mas nunca processado através da Rede Feed-Forward (FFN) ou da pilha de camadas; apenas o token passa pela pilha.

2. Ocean BPTT (Reposicionamento 2):

   • Como o estado oceano é uma variável recorrente, os gradientes devem ser desenrolados através do tempo (Backpropagation Through Time, BPTT).
   • Fluxo de Gradiente: Durante a passagem reversa (backward pass), o gradiente do estado oceano no passo $t$ é computado com base no token no passo $t$. Esse gradiente é então adicionado ao gradiente do estado oculto do passo $t-1$.
   • Operador de Caminho de Leitura: Um operador específico mapeia o gradiente do token de volta para o estado oceano usando os pesos de atenção salvos da passagem direta. Quanto mais um token "lê" o oceano, mais forte é o sinal de gradiente enviado de volta para atualizar a representação do oceano.

2.2 Implementação Algorítmica

O loop de treinamento consiste em:

• Passagem Direta (Forward Pass): O token $t$ atenta para $ocean_{t-1}$. O token passa pela pilha de camadas (Atenção + FFN). A saída substitui $ocean_{t-1}$ para se tornar $ocean_t$.
• Passagem Reversa (Backward Pass): Os gradientes fluem do erro no passo $t$ de volta através das camadas. Crucialmente, o gradiente do estado oceano é acumulado ao longo dos passos temporais, permitindo que erros de tokens futuros atualizem a representação de contextos passados.

2.3 Variantes de Engenharia

O artigo explora três regimes de treinamento para equilibrar custo, tempo e qualidade:

• Serial Ocean BPTT: Desenrolamento completo da cadeia temporal. Maior ganho de qualidade, menor custo computacional, mas tempo de execução (wall-clock time) serial.
• Ocean on Framework: Executa a cadeia completa em frameworks de autograd padrão. Mantém o ganho total de qualidade, mas melhora o rendimento (throughput) via processamento em lote (batching).
• Ocean Detach: Corta a cadeia de gradiente temporal (sem BPTT entre passos). Permite o treinamento paralelo total (como os Transformers padrão), mas sacrifica o ganho de desempenho derivado dos gradientes temporais.

3. Principais Contribuições

• Reposicionamento Arquitetural: Demonstra que a atenção $O(1)$ e a memória $O(1)$ são alcançáveis reutilizando o estado oculto descartado como um vetor "oceano" persistente, eliminando a necessidade de caches KV e codificações posicionais.
• Capacidade Emergente: Mostra que um Transformer pode aprender a ler e extrair informações de um único vetor denso sem instrução explícita, baseando-se puramente no treinamento de ponta a ponta.
• Escalabilidade para Contextos Extremos: Prova que a arquitetura pode lidar com comprimentos de contexto de até 10 milhões de passos sem degradação, um feito impossível para Transformers padrão devido às restrições de memória e computação.
• Desenrolamento de Gradiente: Introduz o "Ocean BPTT", aplicando o BPTT padrão à saída da pilha de camadas completa como o estado recorrente, identificando isso como a principal fonte de ganhos de desempenho.

4. Resultados Experimentais

Experimentos foram conduzidos no corpus de Shakespeare (níveis de caractere e palavra) usando uma implementação customizada em NumPy (float64) e verificada com PyTorch.

4.1 Desempenho e Dinâmica de Treinamento

• Superioridade: Em todas as configurações testadas (larguras $d \in \{32, 64, 128, 256\}$, profundidades $L \in \{2, 4, 6\}$ e vários números de cabeças), o regime "Ocean Full" consistentemente superou o Transformer padrão.
• Ganho Estável: A diferença de desempenho ($\Delta = Loss_{std} - Loss_{full}$) foi positiva e estável (aprox. +0,15 de redução de perda) através da grade de testes.
• Fonte do Ganho: O ganho é primariamente estrutural. Quando a cadeia temporal é desconectada (Ocean Detach), o desempenho cai significativamente, indicando que o desenrolamento do gradiente temporal é o fator crítico, não apenas o mecanismo de atenção em si.

4.2 Longo Contexto e Memória

• Sem Esquecimento: Em um teste de prefill de 10 milhões de passos, a perda permaneceu constante (1,532 no passo 1M vs. 1,531 no passo 10M). O estado oceano não mostrou sinais de saturação de capacidade ou esquecimento.
• Tarefa de Eco (Echo Task): O modelo alcançou perda próxima de zero (0,0006) em uma tarefa de eco com um atraso de 2 milhões de passos, enquanto o Transformer padrão falhou completamente.
• Eficiência de Memória: O estado oceano requer apenas $O(1)$ de memória (um único vetor de dimensão $d$), independente do comprimento do contexto. Em contraste, os Transformers padrão exigem $O(N)$ de memória para caches KV, o que se torna inviável em escalas como 1 milhão de tokens (exigindo ~1 TB para um modelo de 13B).

4.3 Inferência e Geração

• Sem Extrapolação: O oceano gera texto coerente sem precisar de extrapolação de codificação posicional (ex: RoPE, ALiBi) porque o comprimento do contexto não é uma variável de entrada para o mecanismo de atenção.
• Tempo de Execução (Wall-Clock Time): Embora a natureza serial da recorrência torne o treinamento mais lento por passo do que os Transformers paralelos padrão, o artigo observa que o pipelining por camadas (Ocean A-2) pode alcançar um aumento de 1,9x no tempo de execução, e o paralelismo aninhado (Ocean A-3) oferece ganhos teóricos de $10^3$. O autor enfatiza que o treinamento envolve um triângulo de compensação explícito entre custo, tempo e qualidade, onde a implementação prática exige a seleção de regimes específicos (ex: detach para escala industrial, full chain para qualidade de pesquisa).

5. Significância e Alegações

O artigo reivindica três pontos principais, que define como falsificáveis:

1. Consistência de Treinamento: O estado Oceano treina com a mesma forma de curva de aprendizado que o Transformer padrão, confirmando que é um reposicionamento válido e não um paradigma de aprendizado fundamentalmente diferente.
2. Ganho Universal: Para qualquer arquitetura Transformer sob condições idênticas, o estado Oceano produz um ganho de desempenho positivo ($\Delta > 0$).
3. Capacidade Infinita: Com dimensão suficiente ($d$), o estado Oceano exibe nenhuma limitação de capacidade, nenhum esquecimento e nenhuma falha de recuperação em qualquer escala observável.

O autor conclui que o "Estado Oceano" aborda o princípio de inferência $O(1)$ e treinamento para comprimentos de contexto arbitrários, afirmando: "acreditamos que o princípio está resolvido". Ao confiar na habilidade emergente do Transformer de organizar informações dentro de um único vetor, o método remove os gargalos estruturais da atenção quadrática e do crescimento linear de memória, potencialmente permitindo Transformers "pequenos e pessoais" e robôs com estados cognitivos contínuos e humanos. O trabalho enfatiza que esses resultados derivam de experimentação de "e se" em componentes existentes, respeitando as propriedades emergentes da arquitetura original, ao mesmo tempo em que reconhece questões abertas sobre configurações específicas (ex: por que o detach supera o padrão em alguns casos).