Resumo Técnico: RADLADS

Declaração do Problema

Embora variantes de atenção linear (como RWKV e Mamba) ofereçam vantagens significativas sobre os transformadores de atenção softmax tradicional — especificamente tempo de inferência $O(1)$ por token e uso constante de memória ao evitar caches de Chave-Valor (KV cache) — o treinamento de modelos lineares em larga escala do zero permanece proibitivamente caro. Modelos transformer de estado da arte (SoTA) frequentemente exigem treinamento em mais de dez trilhões de tokens, um custo inacessível para a maioria das organizações. Métodos existentes para converter modelos transformer softmax pré-treinados em modelos lineares ou recorrentes (ex: T2R, SUPRA, LoLCats, MOHAWK) historicamente exigiram contagens massivas de tokens (variando de 20B a 100B+ tokens) ou resultaram em baixo desempenho em benchmarks, particularmente em métricas como MMLU. Além disso, muitas tentativas de conversão dependem de arquiteturas híbridas (mantendo parte da atenção softmax) ou falham em igualar a qualidade dos modelos professores originais.

Metodologia

Os autores apresentam o RADLADS (Rapid Attention Distillation to Linear Attention Decoders at Scale), um protocolo projetado para converter modelos transformer de atenção softmax em modelos de decodificador de atenção linear com custo computacional mínimo. O processo envolve três estágios primários e adaptações arquiteturais específicas:

1. O Protocolo RADLADS

A conversão é um processo de três etapas que requer apenas 350–700 milhões de tokens (menos de 0,005% dos dados de pré-treinamento do professor):

• Configuração (Transferência de Pesos de Atenção): Os pesos relacionados à atenção ($W_q, W_k, W_v, W_o$) do modelo professor são transferidos diretamente para as camadas de mistura de sequência do modelo estudante. Outros pesos são inicializados via métodos padrão de pré-treinamento ou configurados para mimetizar o comportamento do professor sem efeito imediato.
• Etapa 1: Alinhamento do Estado Oculto da Atenção: As camadas de mistura de sequência do modelo estudante são treinadas para aproximar os outputs dos estados ocultos das camadas de atenção correspondentes do professor. Isso é realizado usando um objetivo de distância L2 (ou MSE). Os autores descobriram que o uso de um kernel de Atenção Linear com Gate (removendo termos de decaimento off-by-one e bônus) permite um ajuste mais próximo aos estados ocultos do professor do que o RWKV-6 padrão.
• Etapa 2: Destilação de Conhecimento: O modelo estudante completo é treinado para aproximar os logits de saída do professor usando perda de divergência de Kullback-Leibler (KL). Os autores hipotetizam que o conhecimento factual reside primariamente nos MLPs e embeddings do professor; portanto, as taxas de aprendizado para esses componentes são mantidas baixas ou fixas para evitar o esquecimento catastrófico, enquanto o mixer de sequência é treinado de forma mais agressiva.
• Etapa 3: Extensão do Comprimento de Contexto (Opcional): O modelo é ajustado em sequências mais longas (até 16k tokens) usando perda de entropia cruzada padrão sem um modelo professor para aumentar as capacidades de contexto longo. Alternativamente, os autores propõem a Etapa 2a, que estende o comprimento da sequência para 4096 durante a própria fase de destilação, potencialmente eliminando a necessidade de uma Etapa 3 separada.

2. Novas Arquiteturas

Os autores identificaram que as arquiteturas padrão do RWKV nem sempre eram ideais para a conversão. Eles introduziram duas novas variantes:

• RAD-RWKV6 ("RADFinch"): Uma modificação do RWKV-6 usando um kernel de Atenção Linear com Gate e técnicas de balanceamento de estado para melhorar a estabilidade e o ajuste.
• RAD-RWKV7 ("RADGoose"): Uma modificação do RWKV-7 que remove o mecanismo de "tokenshift" (que não ofereceu benefício neste contexto) e aplica Rotary Positional Embeddings (RoPE) diretamente. Esta arquitetura demonstrou convergência mais rápida e menor perda de destilação em comparação com o RWKV-7 não modificado.

3. Hiperparâmetros e Dados

• Dataset: Os autores estabeleceram o DCLM (DataComp-LM) para todas as etapas de conversão, considerando-o superior ao FineWeb ou FineWeb-Edu para converter modelos Qwen.
• Taxas de Aprendizado: Um cronograma de cosine annealing é usado na Etapa 1, começando alto ($10^{-3}$) para alinhar os estados ocultos e terminando próximo à taxa de aprendizado final de pré-treinamento do professor ($10^{-5}$). As Etapas 2 e 3 utilizam uma taxa de aprendizado constante.

Principais Contribuições

1. Receita de Destilação RADLADS: Um protocolo detalhado passo a passo, incluindo hiperparâmetros específicos, contagem de tokens e escolhas de datasets que permite uma conversão de alta qualidade com dados mínimos.
2. Novas Arquiteturas: A introdução do RAD-RWKV6 e RAD-RWKV7, que são otimizados para o processo de conversão, proporcionando inferência mais rápida e melhor alinhamento com os modelos professores do que seus predecessores não modificados.
3. Modelos de Larga Escala: A conversão bem-sucedida de modelos populares de código aberto Qwen2.5 em variantes de atenção linear nas escalas de 7B, 32B e 72B de parâmetros.
4. Lançamento Open Source: O lançamento do código e dos modelos convertidos (QRWKV6/7-7B/32B/72B) sob a licença Apache 2.0 (com restrições de Licença Qwen para o modelo de 72B), permitindo que outros repliquem o processo.

Resultados

Os modelos convertidos alcançam o desempenho de estado da arte entre modelos puramente recorrentes de seu tamanho:

• Eficiência: Converter um modelo de 72B custa menos de US$ 2.000 e requer apenas ~700M de tokens.
• Desempenho: Em benchmarks padrão (Lambada, MMLU, ARC, etc.), os modelos RADLADS superam consistentemente outros métodos de conversão (ex: SUPRA, LoLCats, MOHAWK, ARWKV).
  • O QRWKV7-7B-Instruct alcança uma pontuação MMLU relativa de 92,4% em relação ao seu professor, superando significativamente outras conversões de 7B.
  • O QRWKV6-72B-Instruct alcança uma pontuação MMLU relativa de 89,9%, estabelecendo um novo SoTA para modelos de linguagem puramente RNN nesta escala.
• Velocidade de Inferência: Devido ao mecanismo de atenção linear, os modelos convertidos mostram acelerações significativas sobre os modelos professores conforme o contexto aumenta. Por exemplo, com 8k tokens de entrada e 256 de saída, o modelo 32B QRWKV7 é 1,61x mais rápido que o professor Qwen3-32B; com 6k de entrada e 2k de saída, a aceleração chega a 3,41x.

Limitações e Alegações

O artigo reconhece modestamente diversas limitações:

• Raciocínio e Contexto Longo: Embora o desempenho seja forte em benchmarks padrão, os modelos mostram limitações em tarefas de raciocínio complexo (ex: Minerva Math) e tarefas de contexto muito longo (ex: RULER), onde não melhoram tão efetivamente com tokens de saída adicionais quanto os modelos professores.
• Sensibilidade da Arquitetura: Cada novo design de arquitetura exige testes meticulosos para garantir compatibilidade com o protocolo RADLADS. Por exemplo, o uso de GroupNorm/LayerNorm, comum em algumas variantes de atenção linear, causou instabilidade no treinamento em escalas de 14B+ e teve que ser substituído por técnicas de pré-escalonamento.
• Alinhamento de Dataset: Os autores observam que datasets de raciocínio podem precisar de um alinhamento mais próximo com a distribuição do modelo professor para evitar comportamentos de repetição em loop, um desafio que abordaram parcialmente misturando DCLM com OpenThoughts.

Significância

O artigo afirma que o RADLADS fornece um caminho de baixo custo para democratizar o acesso a modelos de atenção linear de larga escala. Ao reduzir o requisito de treinamento de trilhões para centenas de milhões de tokens, ele permite que pesquisadores e organizações menores testem, treinem e implantem novas variedades de arquiteturas RNN em escala sem os custos extremos associados ao pré-treinamento. Os autores posicionam isso como uma ferramenta para acelerar o desenvolvimento da próxima geração de variantes de atenção de estado compressivo.