Beyond Rejection Sampling: Trajectory Fusion for Scaling Mathematical Reasoning

Resumo Técnico: TrajFusion para Escalonamento de Raciocínio Matemático

Declaração do Problema

Os Grandes Modelos de Linguagem (LLMs) alcançaram progressos significativos no raciocínio matemático, em grande parte através do Ajuste Fino Supervisionado (SFT) em trajetórias de solução de alta qualidade. O paradigma dominante para a construção desses dados é o Ajuste Fino por Rejeição de Amostragem (RFT). No RFT, um modelo professor gera múltiplos caminhos de raciocínio candidatos para um determinado problema, e apenas aqueles com respostas finais corretas são retidos para treinamento; todas as trajetórias incorretas são descartadas.

Embora eficaz, essa abordagem trata a supervisão como um filtro binário. Ela exclui sistematicamente os erros gerados pelo professor, criando uma lacuna na forma como as falhas de raciocínio são modeladas durante o treinamento. Os autores argumentam que as trajetórias incorretas contêm informações diagnósticas valiosas — como deslizes computacionais comuns, suposições ausentes e caminhos de argumento logicamente plausíveis, mas inválidos — que poderiam ajudar os modelos a reconhecer e evitar modos de falha específicos, particularmente em tarefas de raciocínio de formato longo ou ambíguas. Os métodos atuais tratam esses erros como ruído a ser descartado, em vez de supervisão estruturada a ser aproveitada.

Metodologia: TrajFusion

O artigo propõe o TrajFusion, uma estratégia de ajuste fino que reformula o RFT de um mecanismo de filtragem binária para um processo de construção de supervisão estruturada. O TrajFusion integra trajetórias de raciocínio incorretas selecionadas com as corretas em uma única amostra de treinamento fundida, modelando explicitamente um processo de raciocínio de tentativa e erro.

Componentes Principais

1. Amostragem e Análise: Para cada problema $x$, o modelo professor amostra $K$ trajetórias. Estas são particionadas em conjuntos corretos ($Y^+$) e incorretos ($Y^-$) usando um verificador automático.
2. Seleção Adaptativa: O método determina adaptativamente quantos trajetórias incorretas incluir com base em duas estatísticas:
   • Taxa de Erro ($r(x)$): A frequência de tentativas incorretas.
   • Diversidade de Erro ($u(x)$): O número de respostas finais únicas entre as tentativas incorretas (medido pela entropia de Shannon ou contagem de respostas únicas).
     O número de trajetórias incorretas retidas $k(x)$ é calculado como:
     $$k(x) = \min(k_{max}, \lfloor \alpha \cdot r(x) \cdot u(x) \rfloor)$$
     Isso garante que problemas com erros diversos e frequentes recebam uma supervisão mais rica, enquanto problemas com erros esparsos ou pouco informativos degeneram naturalmente para o RFT convencional.
3. Construção da Amostra Fundida: As trajetórias incorretas selecionadas são intercaladas com a trajetória correta usando prompts de reflexão fixos (ex: "Espere... isso não parece certo. Vamos tentar novamente."). A estrutura final é:
   $$T(x) = [y_{e1}, \rho_1, \dots, y_{ek}, \rho_k, y^*]$$
   onde $y_{ei}$ são caminhos incorretos, $\rho_i$ são prompts de reflexão e $y^*$ é a solução correta.
4. Treinamento: As sequências fundidas são treinadas usando a predição do próximo token padrão (SFT) sem alterações na arquitetura do modelo ou no objetivo de treinamento.

Contribuições Principais

• Formalização de uma Lacuna de Supervisão: Os autores identificam e formalizam a exclusão sistemática de trajetórias incorretas geradas pelo professor no RFT, argumentando que esses erros fornecem valor diagnóstico para a modelagem de falhas de raciocínio.
• Framework TrajFusion: Um framework de supervisão estruturada que funde adaptativamente trajetórias incorretas e corretas. Ele trata a diversidade de erro como um sinal de controle explícito, permitindo que o método escale de forma robusta através de problemas de variadas dificuldades sem exigir modelos auxiliares, geradores de crítica ou rótulos de dificuldade externos.
• Validação Empírica: O método demonstra melhorias consistentes sobre o RFT convencional em múltiplos modelos base (LLaMA3-8B, DeepSeekMath-7B), regimes de dados (15K a 100K exemplos) e comprimentos de contexto (até 32K).

Resultados Experimentais

Os autores avaliaram o TrajFusion em seis benchmarks de raciocínio matemático (GSM8K, MATH, CollegeMath, DeepMind-Mathematics, OlympiadBench, TheoremQA) e conjuntos de dados de nível de competição (AIME24, AIME25, MATH-500).

• Ganhos de Desempenho: O TrajFusion superou consistentemente o RFT convencional. No conjunto de dados MATH com LLaMA3-8B (100K dados), o TrajFusion alcançou 51,4% de acurácia comparado a 48,5% do RFT. No DeepSeekMath-7B, atingiu 52,3% vs. 50,2%.
• Impacto na Complexidade: As melhorias foram mais pronunciadas em benchmarks de raciocínio de formato longo e desafiadores (MATH, OlympiadBench, TheoremQA), onde os caminhos de raciocínio são mais longos e os modos de falha são mais diversos. Os ganhos em conjuntos de dados mais simples como o GSM8K foram menores, mas ainda positivos.
• Eficiência de Dados: O TrajFusion alcançou resultados competitivos com apenas 15K exemplos de treinamento, superando baselines que utilizaram conjuntos de dados significativamente maiores (ex: MMIQC com 2,3M exemplos).
• Escalabilidade para Contexto Longo: Quando aplicado a uma versão de 32K de contexto do Qwen2.5-Math-7B usando dados de formato longo destilados do DeepSeek-R1, o TrajFusion melhorou o Pass@1 no AIME24 de 40,0% (RFT) para 43,3%.
• Estudos de Ablação:
  • Aumentar o número de erros fundidos ($k_{max}$) melhorou o desempenho até certo ponto (retornos decrescentes observados em $k_{max}=6$).
  • Selecionar erros baseando-se em diversidade (HighDiv) rendeu melhores resultados do que selecionar baseando-se apenas em frequência (LowDiv), sugerindo que modos de falha estruturalmente distintos são mais informativos.
  • Modelos treinados com TrajFusion inicialmente geraram sequências mais longas, mas tornaram-se gradualmente mais concisos enquanto aumentavam a frequência de tokens de reflexão, indicando um uso mais eficiente dos padrões estruturais.

Significância e Alegações

O artigo afirma que o TrajFusion oferece uma maneira simples e eficaz de escalar o raciocínio matemático ao alavancar sinais de supervisão que são tipicamente descartados. Ao reformular a amostragem de rejeição como um processo de construção de trajetórias estruturadas de tentativa e erro, o método fornece uma supervisão mais rica sem modificações arquiteturais ou pipelines complexos de aprendizado por reforço.

Os autores enfatizam que o TrajFusion é particularmente valioso para problemas de raciocínio desafiadores e de formato longo, onde distinguir o raciocínio válido do falho é essencial. Eles afirmam que a natureza adaptativa do TrajFusion garante robustez: ele enriquece a supervisão quando os erros são informativos e retorna com segurança ao comportamento padrão do RFT quando estes não o são. Essa abordagem preenche a lacuna entre a filtragem binária de correção e a realidade matizada das falhas de raciocínio, permitindo que os modelos aprendam com seus próprios erros em um ambiente supervisionado.