Resumo Técnico: Profundidade Computacional Adaptativa via Roteamento de Tokens Aprendido em Transformers

Declaração do Problema

As arquiteturas Transformer padrão aplicam um número fixo de camadas a cada token em uma sequência, independentemente da dificuldade contextual ou da densidade de informação daquele token. Este design prioriza a simplicidade arquitetural sobre a adaptabilidade por token, resultando em ineficiência computacional significativa. Na prática, tokens comuns em contextos previsíveis requerem muito menos processamento do que tokens raros em construções novas, mas ambos recebem recursos computacionais idênticos durante cada passagem direta. Esta ineficiência é particularmente custosa na escala de inferência, onde o custo dominante é a passagem direta completa através de todas as camadas para todos os tokens. Soluções existentes, como Tempo Computacional Adaptativo (ACT) para RNNs, Transformers Universais, Mistura de Profundidades (MoD) e mecanismos de Saída Antecipada, tentam abordar isso, mas frequentemente dependem de seleção rígida, orçamentos acumulados ou limiares de confiança post-hoc que podem não ser totalmente diferenciáveis ou ótimos.

Metodologia: Atenção Seletiva por Token (TSA)

O artigo propõe a Atenção Seletiva por Token (TSA), um mecanismo que introduz um portão aprendido, por token, nas atualizações residuais entre blocos de transformadores consecutivos.

Arquitetura e Roteamento

A TSA insere um roteador leve $r_l$ após cada bloco $f_l$ (para $l = 0, \dots, L-2$) em um transformador apenas decodificador com pré-normalização. O bloco inicial $f_0$ executa sempre incondicionalmente.

• Design do Roteador: Cada roteador é um Perceptron Multicamada (MLP) de duas camadas com saída sigmoide, produzindo uma probabilidade de parada contínua $p_l \in (0, 1)$ para cada token. A dimensão oculta é $d/4$ (arredondada para baixo em 16), adicionando aproximadamente 1,7% de sobrecarga de parâmetros ao modelo base.
• Atualização Portada: As atualizações residuais para o bloco subsequente $f_{l+1}$ são escaladas por $(1 - p_l)$.
  • Se $p_l = 0$, a atualização é padrão.
  • Se $p_l = 1$, o estado permanece inalterado, efetivamente pulando o bloco.
  • A interpolação é suave, garantindo que o mecanismo seja totalmente diferenciável sem exigir estimadores straight-through, amostragem Gumbel ou aprendizado por reforço.

Treinamento e Regularização

Durante o treinamento, todas as camadas executam totalmente; o portão escala as atualizações residuais, mas não pula o cálculo. A perda total é definida como $L = L_{task} + L_{depth}$.

• Regularização de Profundidade: Um termo $L_{depth} = \lambda \cdot \frac{1}{L-1} \sum (1 - p_l)$ é adicionado para incentivar gentilmente a parada antecipada.
• Comportamento Emergente: Notavelmente, o artigo demonstra que o mecanismo de roteamento aprende a pular operações mesmo quando $\lambda = 0$ (sem pressão explícita de profundidade), impulsionado exclusivamente pelo gradiente da perda da tarefa.

Inferência

Na inferência, o mecanismo muda para um modo "esparso". Tokens com $p_l > 0.5$ são excluídos da Rede Feed-Forward (FFN) via operações de coleta/dispersão (gather/scatter), enquanto a atenção permanece densa para preservar a semântica exata chave-valor. Isso gera economias reais de FLOPs e acelerações no tempo de parede.

Contribuições Principais

1. Roteamento de Token Diferenciável: Um portão suave simples e contínuo nas atualizações residuais que permite roteamento por token e por camada sem alterações arquitetônicas no Transformer base.
2. Emergência Impulsionada por Gradiente: Evidência de que o roteamento proporcional à dificuldade emerge naturalmente apenas do gradiente da perda da tarefa. Mesmo sem regularização de profundidade ($\lambda=0$), o roteador aprende a pular ~20% das operações token-camada.
3. Compromisso Eficiência-Qualidade: Validação em modelagem de linguagem em nível de caractere mostrando economias de 14–23% em operações token-camada (TLOps) com perda de qualidade inferior a 0,5%.
4. Superioridade sobre Saída Antecipada: Com eficiência igualada, a TSA alcança uma perda de validação 0,71% menor do que as linhas de base de Saída Antecipada, sugerindo que decisões de roteamento aprendidas end-to-end são superiores à aplicação de limiares de confiança no momento da inferência.
5. Aceleração no Mundo Real: Demonstração de aceleração real no tempo de parede (até 2,3%) em hardware comercial (Apple M1 Pro) via inferência esparsa.

Resultados Experimentais

Tarefas Sintéticas

Em tarefas de cópia e ordenação (sequências de comprimento 10):

• Tarefa de Cópia: O roteador aprendeu que a tarefa era um mapeamento identidade, pulando 54,9% das TLOps ($\alpha = 0,341$) sem perda de precisão.
• Tarefa de Ordenação: O roteador alocou mais computação, pulando 22,5% das TLOps ($\alpha = 0,730$), refletindo a maior dificuldade.

Modelagem de Linguagem em Nível de Caractere

• Tiny-Shakespeare: A TSA economizou 22,8% das TLOps ($\alpha = 0,726$) com um aumento na perda de validação de apenas 0,006 nats (+0,4%).
• enwik8: Em um corpus mais diversificado (Wikipedia), a TSA economizou 13,9% das TLOps ($\alpha = 0,833$) sem degradação de qualidade (perda de validação marginalmente melhor).
• Robustez: O mecanismo é robusto ao hiperparâmetro de regularização $\lambda$ em duas ordens de grandeza ($[0, 0,1]$). Em $\lambda=0,05$, alcançou economias de 50,4% de TLOps com custo de qualidade negligenciável.

Comparação com Saída Antecipada

Com fração ativa igualada ($\alpha \approx 0,726$), a TSA alcançou uma perda de validação 0,71% menor do que uma linha de base de Saída Antecipada. Isso indica que aprender decisões de roteamento end-to-end produz decisões melhores do que aplicar um limiar de confiança no momento da inferência.

Vazão no Tempo de Parede

• Soft-TSA (Treinamento): Adicionou ~1% de sobrecarga independentemente de $\alpha$.
• Sparse-TSA (Inferência): Alcançou uma aceleração de 2,3% no Tiny-Shakespeare ($\alpha=0,726$) e desempenho de equilíbrio (break-even) no enwik8 ($\alpha=0,833$). As acelerações exigiram um tamanho de lote de $\ge 64$ para superar a sobrecarga de sincronização CPU-GPU.

Significado e Alegações

O artigo alega que a TSA oferece uma solução mínima e arquiteturalmente simples para a ineficiência dos Transformers de profundidade fixa. Seu significado principal reside em demonstrar que a profundidade computacional adaptativa pode ser aprendida diretamente a partir do objetivo da tarefa sem supervisão explícita regarding dificuldade ou sinais complexos de aprendizado por reforço.

Os autores enfatizam que o roteador atua como um "portão de ruído aprendido", atenuando atualizações quando as contribuições residuais são redundantes. Isso permite que o modelo alocie computação dinamicamente com base na dificuldade do token, alcançando ganhos significativos de eficiência (até 55% em tarefas sintéticas, 14–23% em modelagem de linguagem) com compromissos de qualidade mínimos. O trabalho sugere que o roteamento diferenciável end-to-end é uma estratégia mais eficaz para profundidade adaptativa do que seleção discreta ou limiarização post-hoc, abrindo caminho para inferência mais eficiente em hardware comercial.

O artigo permanece modesto quanto à escala, observando que todos os experimentos foram conduzidos em modelos com aproximadamente 5–6 milhões de parâmetros, e o comportamento de escalabilidade para modelos maiores (10M–100M+) permanece uma questão em aberto. Além disso, as acelerações atuais dependem do tamanho do lote, e a esparsidade de atenção ainda não foi implementada.