Além da Self-Attention: O Que Vem Depois dos Transformers

A arquitetura transformer está por trás da IA há oito anos. Todos os grandes modelos de linguagem, a maioria dos sistemas de geração de imagens e um número crescente de modelos de áudio e vídeo são construídos sobre o mecanismo de self-attention apresentado no artigo 'Attention Is All You Need'. Mas a self-attention tem um problema fundamental: seu custo de computação e memória cresce quadraticamente com o tamanho da sequência. Dobre o tamanho da entrada e o custo quadruplica.

Para sequências curtas, isso não faz diferença. Mas para as janelas de contexto de 128K tokens que estamos buscando — e os milhões de tokens que as pessoas querem — isso vira um gargalo sério. Uma onda de pesquisas explora alternativas: residuais de atenção que reaproveitam computação entre camadas, variantes de atenção linear que eliminam o custo quadrático e arquiteturas híbridas que misturam atenção com mecanismos mais baratos. O transformer não está desaparecendo, mas está sendo remodelado.

Por Que a Self-Attention É Cara

Para entender as alternativas, é preciso entender o que a self-attention realmente calcula. Dada uma sequência de N tokens, a self-attention calcula um score de relevância entre cada par de tokens. Token 1 vs. token 2, token 1 vs. token 3, ..., token 1 vs. token N, depois token 2 vs. todos os outros, e assim por diante. Isso dá N² pares.

import torch
import torch.nn.functional as F
def self_attention(Q, K, V):
"""
Standard self-attention.
Q, K, V: (batch, seq_len, d_model)
The attention matrix is seq_len × seq_len.
For seq_len = 1024:   ~1M entries   (manageable)
For seq_len = 32768:  ~1B entries   (expensive)
For seq_len = 131072: ~17B entries  (very expensive)
"""
d_k = Q.size(-1)
# This matmul creates the N×N attention matrix
scores = torch.matmul(Q, K.transpose(-2, -1)) / (d_k ** 0.5)
weights = F.softmax(scores, dim=-1)
return torch.matmul(weights, V)

Com 4K tokens, a matriz de atenção tem 16 milhões de entradas — sem problema para GPUs modernas. Com 128K tokens, são 16 bilhões de entradas. Com 1 milhão de tokens, passa de um trilhão. Mesmo com o Flash Attention (que não reduz a computação, mas melhora muito os padrões de acesso à memória), o crescimento quadrático acaba vencendo.

É por isso que os primeiros transformers ficavam limitados a 512 ou 1024 tokens. Cada geração de hardware e otimização empurrou esse teto mais para cima, mas estamos batendo numa parede matemática. Escala linear (O(N)) seria fundamentalmente melhor que escala quadrática (O(N²)), e é isso que a maioria das arquiteturas alternativas busca.

Residuais de Atenção: Reaproveitando o Que Já Foi Calculado

Uma das abordagens mais pragmáticas para reduzir o custo da atenção não substitui a atenção — ela torna cada camada de atenção mais barata, reaproveitando a computação das camadas anteriores.

A observação: em um transformer profundo (digamos, 32 camadas), os padrões de atenção de camadas adjacentes costumam ser notavelmente parecidos. A camada 15 e a camada 16 tendem a prestar atenção nas mesmas posições, com pequenos ajustes. Calcular a matriz N² completa do zero em cada camada é redundante — boa parte do trabalho já foi feita uma camada antes.

Os residuais de atenção exploram isso calculando um padrão de atenção 'residual': a diferença entre aquilo em que esta camada quer prestar atenção e o que a camada anterior calculou. Se a diferença for pequena (o que geralmente acontece nas camadas do meio), a computação fica mais barata. O padrão de atenção completo é a soma do padrão da camada anterior com o residual da camada atual.

Isso é análogo a como a compressão de vídeo funciona: em vez de armazenar cada quadro de forma independente, você guarda um keyframe e depois uma série de diferenças (resíduos) em relação a ele. As diferenças costumam ser bem menores que o quadro completo, então a compressão fica muito melhor.

Na prática, os residuais de atenção reduzem o custo de computação da atenção em 30-50% nas camadas do meio de modelos profundos, com impacto mínimo na qualidade. As primeiras e últimas camadas ainda precisam de atenção completa (seus padrões são mais distintos), mas as camadas do meio — que são a maioria — ganham speedups significativos.

Atenção Linear: Eliminando o Custo Quadrático

As variantes de atenção linear tentam reformular a atenção para que escale como O(N) em vez de O(N²). A ideia geral: em vez de calcular explicitamente a matriz de atenção N×N, encontrar uma forma de obter a mesma saída (ou aproximadamente a mesma) usando operações lineares.

O truque matemático depende da decomposição por kernel do softmax. A atenção padrão calcula softmax(QK^T)V. Se você substituir o softmax por uma função kernel diferente que possa ser decomposta como φ(Q) · φ(K)^T, pode reorganizar a ordem do cálculo: em vez de (φ(Q) · φ(K)^T) · V (que tem um intermediário N×N), calcule φ(Q) · (φ(K)^T · V) (que tem um intermediário d×d, onde d é a dimensão do modelo). Como d << N para sequências longas, isso fica muito mais barato.

def linear_attention(Q, K, V, feature_map=None):
"""
Linear attention via kernel feature maps.
Cost: O(N * d^2) instead of O(N^2 * d)
"""
if feature_map is None:
# ELU+1 is a common choice (from Katharopoulos et al.)
feature_map = lambda x: F.elu(x) + 1
Q = feature_map(Q)  # (batch, seq_len, d)
K = feature_map(K)  # (batch, seq_len, d)
# Key insight: compute K^T @ V first (d × d matrix)
# instead of Q @ K^T first (N × N matrix)
KV = torch.einsum('bnd,bnm->bdm', K, V)  # (batch, d, d)
# Then multiply by Q
output = torch.einsum('bnd,bdm->bnm', Q, KV)  # (batch, N, d)
# Normalize
Z = torch.einsum('bnd,bd->bn', Q, K.sum(dim=1))  # normalization
output = output / Z.unsqueeze(-1)
return output

O porém: trocar o softmax por outra função kernel altera a distribuição da atenção, e modelos treinados com atenção softmax não necessariamente se adaptam bem à atenção linear. A diferença de qualidade diminuiu bastante — variantes recentes de atenção linear atingem 95-98% da qualidade da atenção softmax —, mas ela persiste, especialmente em tarefas que exigem recuperação precisa de informações de longo alcance.

Modelos de Espaço de Estados: Um Paradigma Diferente

Os modelos de espaço de estados (SSMs), como o Mamba, adotam uma abordagem fundamentalmente diferente. Em vez de calcular relações entre pares de tokens, processam a sequência por meio de uma recorrência — mantendo um estado oculto de tamanho fixo que é atualizado a cada token. Isso é inerentemente O(N): processar o dobro de tokens leva o dobro do tempo, não quatro vezes mais.

A inovação dos SSMs modernos é fazer com que os parâmetros da recorrência dependam da entrada (selective state spaces). Isso dá ao modelo uma forma de atenção baseada em conteúdo — ele pode 'escolher' quais informações lembrar e quais esquecer — sem o custo quadrático. Modelos no estilo Mamba igualam a qualidade dos transformers em muitos benchmarks e são significativamente mais rápidos em sequências longas.

O porém: SSMs processam tokens sequencialmente, o que dificulta a paralelização durante o treinamento em comparação com transformers (que processam todos os tokens simultaneamente). A eficiência de treinamento importa — um modelo que é 2x mais rápido na inferência, mas 3x mais lento para treinar, não é necessariamente uma vitória, já que a maior parte da computação total vai para o treinamento.

Arquiteturas Híbridas: O Caminho Pragmático

A tendência atual nos modelos em produção são as arquiteturas híbridas, que combinam diferentes mecanismos de atenção. A lógica é simples: partes diferentes de um modelo se beneficiam de tipos diferentes de computação.

• Atenção completa para raciocínio global. Algumas camadas precisam prestar atenção em toda a sequência — encontrar contexto relevante a milhares de tokens de distância. Essas camadas usam self-attention padrão (possivelmente otimizada com Flash).
• Atenção local para o contexto próximo. Muitas camadas prestam atenção principalmente em tokens próximos (atenção em janela deslizante). Usar uma janela fixa de 256-1024 tokens reduz o custo para O(N·W), onde W é o tamanho da janela.
• Atenção linear para contexto amplo. Algumas camadas precisam agregar informações ao longo da sequência, mas não precisam de pesos de atenção precisos. A atenção linear oferece isso com custo O(N).
• Camadas SSM para processamento sequencial. Camadas no estilo Mamba conseguem processar dependências sequenciais de forma eficiente, sem nenhum cálculo de atenção.

Modelos como o Jamba (AI21) e várias arquiteturas de pesquisa alternam entre esses mecanismos conforme o papel da camada. As camadas iniciais usam atenção local (processando sintaxe e padrões locais). As camadas do meio usam atenção linear ou SSMs (construindo representações mais amplas). Algumas poucas camadas estratégicas usam atenção completa (raciocínio global e recuperação). Isso resulta em escala quase linear no geral, preservando a qualidade do modelo que exige alguma atenção completa.

O Que os Desenvolvedores Devem Acompanhar

Se você está construindo aplicações sobre modelos de linguagem, as mudanças arquiteturais que acontecem por baixo afetam seu trabalho de formas concretas.

• As janelas de contexto vão continuar crescendo. À medida que o custo da atenção cai, as janelas de contexto se expandem. Isso muda a arquitetura da aplicação: em vez de construir pipelines RAG complexos para encaixar o contexto relevante em uma janela de 4K, você pode simplesmente jogar tudo em um prompt de 1M tokens. A simplicidade é atraente, mas as implicações de latência e custo variam entre arquiteturas.
• Os perfis de latência mudam. Transformers têm latência relativamente estável até certo ponto, e depois ela cresce quadraticamente. Modelos com atenção linear e SSMs têm aumentos de latência mais graduais e lineares. Para aplicações em que o tempo de resposta importa, entender o comportamento de escala do seu modelo é fundamental.
• As diferenças de qualidade dependem da tarefa. Modelos com atenção linear podem ter desempenho levemente inferior em tarefas que exigem recuperação precisa de posições específicas em contextos longos ('qual era o terceiro item da lista na página 47?'). Em tarefas que exigem compreensão geral, eles se saem igualmente bem. Conheça o seu caso de uso.
• A otimização de inferência importa cada vez mais. À medida que os modelos ficam mais complexos na arquitetura (misturando diferentes tipos de atenção), os motores de inferência precisam lidar com computação heterogênea de forma eficiente. vLLM, TensorRT-LLM e frameworks semelhantes estão se adaptando, mas arquiteturas customizadas podem não ter suporte imediato.

O transformer não está sendo substituído — está evoluindo. A self-attention continua sendo o mecanismo mais expressivo que temos para modelar relações entre tokens. Mas ele não precisa ser usado em todo lugar, em toda camada, com custo N² total. Os modelos dos próximos anos vão usar a atenção de forma cirúrgica — precisão total onde mais importa, alternativas mais baratas em todo o resto. O resultado serão modelos mais rápidos, que lidam com contextos mais longos e custam menos para rodar, igualando ou superando a qualidade atual. Vale prestar atenção nisso.