Resumo Técnico: LaneRoPE

Declaração do Problema

Técnicas de escalonamento paralelo em tempo de teste, como amostragem best-of-N, melhoram a precisão de Modelos de Linguagem de Grande Escala (LLMs) gerando múltiplas sequências condicionadas ao mesmo prompt de entrada. Embora esses métodos aproveitem a eficiência computacional da inferência em lote, tradicionalmente operam com paralelismo não coordenado. Nas abordagens padrão, cada sequência em um lote é gerada independentemente, o que significa que elas não reutilizam gerações intermediárias, cálculos ou observações de outras sequências. Essa independência proíbe a exploração de estruturas de problemas decomponíveis, limita o compartilhamento de saídas intermediárias e restringe a geração de respostas diversas e colaborativas.

Abordagens "colaborativas" existentes tentam resolver isso permitindo visibilidade compartilhada sobre gerações intermediárias. No entanto, muitos métodos anteriores dependem de reparametrizar posições para ordenar virtualmente múltiplas sequências em uma única sequência longa. Os autores argumentam que essa abordagem é inerentemente um problema de geração de múltiplas sequências, onde a inferência deve refletir com precisão as posições ao longo tanto da dimensão de token quanto da dimensão de sequência (ou "faixa").

Metodologia: LaneRoPE

O artigo propõe LaneRoPE, um método para permitir coordenação e colaboração entre $N > 1$ sequências no momento da geração com alterações mínimas na arquitetura subjacente do LLM. A abordagem baseia-se em duas inovações técnicas principais:

1. Mascaramento de Atenção Inter-Sequência:
   LaneRoPE modifica o mecanismo de atenção causal padrão para permitir que a amostragem de sequências seja dependente umas das outras. Em vez de condicionar o próximo token $x^{(n)}_{i+1}$ apenas ao seu próprio histórico $x^{(n)}_{1:i}$, o modelo o condiciona aos históricos de todas as sequências $\{x^{(m)}_{1:i}\}_{m=1}^N$. Isso é alcançado por meio de máscaras de atenção causal entre sequências, tornando os traços de raciocínio intermediários visíveis em todas as faixas paralelas.

2. Codificação Posicional Consciente de Sequência (LaneRoPE):
   A Codificação Posicional Rotativa (RoPE) padrão não consegue distinguir entre tokens no mesmo passo de tempo em sequências diferentes, pois compartilham índices posicionais idênticos. LaneRoPE estende a RoPE introduzindo uma segunda matriz de rotação em blocos diagonais, $R_{\Omega, m}$, que codifica o índice da sequência $m$ (a "faixa").

   • Mecanismo: Os vetores de consulta e chave são rotacionados tanto pelo índice do token $i$ quanto pelo índice da faixa $m$: $f_{\{q,k\}}(x_i, i, m) = R_{\Omega, m} R_{\Theta, i} (W x_i + b)$.
   • Propriedades: Essa formulação trata a grade conjunta (token, sequência) como uma base de Fourier 2D. Crucialmente, como a matriz de rotação da faixa é ortogonal, os pesos de atenção entre tokens dentro da mesma sequência permanecem inalterados, garantindo compatibilidade retroativa. Para $N=1$, LaneRoPE simplifica-se para a RoPE padrão.
   • Estratégias de Inicialização: Os autores fornecem estratégias para inicializar LaneRoPE a partir de modelos pré-treinados para imitar comportamentos existentes:
     • GroupThink: Ao definir frequências de faixa $\Omega = K\Theta$ (onde $K$ é um intervalo de sequência), LaneRoPE recupera a "ordenação virtual" do GroupThink.
     • Amostragem Independente: Ao aumentar as camadas de projeção com um viés de Fourier, o modelo pode ser inicializado para suprimir a atenção entre faixas, simulando efetivamente a amostragem independente.
     • Correção Sensível a NTK: Para mitigar problemas com posições relativas virtuais negativas na inicialização do GroupThink, os autores introduzem uma função de escala que decai rotações de faixa de baixa frequência, preservando a capacidade do modelo de codificar posições absolutas.

Pipeline de Treinamento

Para elicitar interações entre faixas, os autores propõem duas estratégias de ajuste fino:

1. Ajuste Fino Supervisionado (SFT): Um conjunto de dados é sintetizado onde um LLM desempenha o papel de $N$ assistentes diferentes colaborando em um problema. Os assistentes são instruídos a ver os pensamentos parciais uns dos outros e evitar trabalho redundante. O conjunto de dados final consiste em traços individuais dessas conversas, forçando o modelo a aprender dependências entre sequências.
2. Otimização Kahneman-Tversky (KTO): O modelo é treinado em lotes de traços de raciocínio (alguns corretos, alguns incorretos) gerados independentemente pelo modelo base. A perda KTO maximiza a probabilidade de conclusões corretas e minimiza a probabilidade de conclusões incorretas, incentivando o modelo a aprender uma distribuição conjunta não trivial onde faixas paralelas podem corrigir ou complementar umas às outras.

Resultados

Os autores avaliam LaneRoPE em benchmarks de raciocínio matemático (MATH500, AMC 23, AIME 24/25) usando modelos DeepSeek-R1-Distill-Qwen (1,5B e 7B parâmetros).

• Desempenho: LaneRoPE demonstra melhorias confiáveis de precisão sobre baselines sequenciais e amostragem independente. Os modelos de 7B mostram ganhos consistentes, particularmente ao usar treinamento KTO com inicialização sensível a NTK e frequências aprendíveis.
• Comparação: LaneRoPE geralmente supera a baseline "Bridge" (outro método de geração interdependente) em modelos de 7B, atribuído à maior capacidade do LaneRoPE de compartilhar informações e à sua arquitetura mais eficiente.
• Modelos Pequenos: O modelo de 1,5B mostra resultados mistos, beneficiando-se menos da inicialização GroupThink não treinada e mostrando escalabilidade mais fraca com o aumento de faixas, sugerindo uma limitação do tamanho do modelo base em vez da arquitetura.
• Eficiência: LaneRoPE introduz sobrecarga negligenciável. Medições de tempo de inferência mostram apenas um aumento de ~6% em comparação com transformers padrão, principalmente devido ao pré/pós-processamento e não ao mecanismo de atenção em si. Em contraste, a baseline Bridge incorre em um aumento de até 25% no tempo de inferência devido a camadas adicionais.

Significado e Alegações

O artigo afirma que LaneRoPE representa um passo promissor rumo à habilitação do raciocínio colaborativo paralelo explorando fundamentalmente as posições relativas de tokens tanto dentro quanto fora de uma sequência.

• Eficiência Arquitetural: Diferentemente de trabalhos anteriores que exigem engines de tempo de execução especializados ou mecanismos de ramificação complexos, LaneRoPE requer alterações mínimas na arquitetura e introduz sobrecarga de inferência negligenciável, tornando-o adequado para incorporação rápida em pipelines de LLM existentes.
• Flexibilidade: A formulação é expressiva o suficiente para abranger técnicas existentes (como amostragem independente e GroupThink) enquanto permite ajustes finos adicionais para descobrir novas estratégias colaborativas.
• Eficiência de Dados: O método pode ser desbloqueado com muito poucos parâmetros adicionais (< 0,5% sobre modelos base) usando SFT simples ou KTO com dados sintéticos.

Os autores reconhecem limitações, observando que ainda não empregam um esquema de treinamento dedicado para mesclar saídas de múltiplas faixas e que seus experimentos atuais focam em tamanhos de lote pequenos ($N \le 4$). Eles sugerem que trabalhos futuros poderiam explorar Aprendizado por Reforço com Feedback Verificável (RLVF) para fortalecer ainda mais comportamentos colaborativos.