Training Large Language Models To Reason In Parallel With Global Forking Tokens

O artigo propõe o método SSFT, que utiliza uma perda baseada em conjuntos para preservar modos de raciocínio únicos e gerar tokens de bifurcação globais, permitindo que o modelo GFPO subsequente supere abordagens existentes em tarefas complexas de matemática e geração de código.

O essencial

Imagine que você tem um gênio da matemática (um Modelo de Linguagem Grande, ou LLM) que está tentando resolver um problema muito difícil.

O problema tradicional é o seguinte: quando o gênio começa a pensar, ele muitas vezes fica preso em um único caminho de raciocínio. Se esse caminho estiver errado, ele continua andando na direção errada até o fim, desperdiçando tempo e energia. Para tentar evitar isso, os pesquisadores costumam pedir ao gênio para pensar "mais rápido" ou "mais devagar" (ajustando a temperatura), mas isso é como tentar adivinhar qual caminho é o certo jogando dados: às vezes funciona, mas muitas vezes você perde a precisão em troca de tentar muitas coisas.

Este paper, chamado "Treinando Grandes Modelos de Linguagem a Raciocinar em Paralelo com Tokens de Garfo Global", propõe uma solução inteligente e elegante. Vamos usar uma analogia de cozinha e receitas para entender como funciona.

1. O Problema: A Cozinha Caótica

Imagine que você é o chefe de uma cozinha e precisa preparar um prato complexo. Você tem vários cozinheiros (o modelo) e várias receitas corretas (trilhas de raciocínio) que levam ao mesmo prato final.

• O jeito antigo (SFT comum): Você pega todas as receitas corretas e joga na mesma panela, misturando tudo. O cozinheiro tenta aprender, mas acaba criando uma "sopa" confusa. Ele não sabe qual receita seguir, então ele faz um meio-termo que não é nem uma coisa, nem outra. O resultado é um prato medíocre.
• O jeito do "Garfo" (Forking Tokens): Em vez de misturar tudo, você dá ao cozinheiro etiquetas especiais (os Tokens de Garfo Global). Imagine que você tem 6 etiquetas coloridas: "Pense como um Matemático Clássico", "Pense como um Lógico", "Pense como um Criativo", etc.

2. A Solução: O Sistema de Etiquetas Inteligentes (SSFT)

A grande inovação deste trabalho é ensinar o cozinheiro a associar automaticamente cada etiqueta a uma receita específica, sem que você precise dizer qual é qual a cada vez.

Eles criaram um método chamado SSFT (Ajuste Fino Supervisionado por Conjunto). Funciona assim:

1. O Jogo de Emparelhamento: O modelo recebe um problema e várias soluções corretas (trilhas de raciocínio).
2. O "Match" Perfeito: Em vez de forçar o modelo a aprender tudo ao mesmo tempo, o algoritmo faz um "jogo de emparelhamento" (como um casamento perfeito). Ele pergunta: "Qual etiqueta de pensamento (Garfo) combina melhor com qual solução correta?"
   • Se a "Etiqueta Azul" combina com a "Receita A", o modelo aprende isso.
   • Se a "Etiqueta Vermelha" combina com a "Receita B", o modelo aprende isso.
3. O Resultado: Ao final do treinamento, o modelo não é mais uma "sopa". Ele se tornou um mestre que sabe exatamente qual "modo de pensamento" ativar dependendo da etiqueta que você der a ele.
   • Se você der a etiqueta <think 1>, ele ativa um modo de raciocínio profundo e lógico.
   • Se você der a etiqueta <think 2>, ele ativa um modo criativo e lateral.

3. Por que isso é melhor? (A Analogia do Garfo)

Imagine que o raciocínio é uma árvore. No começo, todas as ideias são iguais. Mas, em algum ponto, a árvore precisa se dividir (furar) para explorar caminhos diferentes.

• Sem o método: O modelo tenta adivinhar onde deve "furar" a árvore enquanto está pensando. É difícil e muitas vezes ele erra o momento.
• Com o método: Você segura a árvore e diz: "Neste ponto exato, use o Garfo Azul para ir para a esquerda, e o Garfo Vermelho para ir para a direita". O modelo aprende que esses "Garfos" (tokens) são os interruptores que ativam modos de pensamento distintos e precisos.

4. O Poder do "Garfo" (GFPO)

Depois de treinar o modelo com essas etiquetas, eles fazem um segundo passo chamado GFPO. É como se eles ensinassem o modelo a escolher sozinho qual etiqueta usar para um problema novo.

• Se o problema é muito difícil, o modelo aprende a dizer: "Ah, esse precisa do Garfo Azul!" e ativa o modo de pensamento mais complexo automaticamente.

5. Os Resultados na Prática

Os autores testaram isso em matemática (como o AIME, uma olimpíada difícil) e em programação.

• O que aconteceu? O modelo treinado com esse método conseguiu resolver problemas muito mais difíceis do que os modelos comuns.
• A mágica: Quando você pede para o modelo pensar em paralelo (gerar 6 respostas ao mesmo tempo usando as 6 etiquetas diferentes), a chance de acertar a resposta correta dispara. É como ter 6 especialistas diferentes trabalhando no mesmo problema ao mesmo tempo, em vez de ter 6 cópias do mesmo especialista tentando adivinhar.

Resumo em uma frase

Em vez de pedir ao cérebro artificial para "adivinhar" como pensar de várias formas, os pesquisadores ensinaram o modelo a ter interruptores específicos (os Garfos Globais) que ativam modos de raciocínio diferentes e precisos, garantindo que ele explore todas as possibilidades corretas sem se confundir.

É como transformar um cozinheiro que mistura tudo em uma panela em uma equipe de especialistas, onde cada um sabe exatamente qual receita seguir quando você acende a luz certa na cozinha.

Resumo técnico

Título: Treinamento de Grandes Modelos de Linguagem para Raciocinar em Paralelo com Tokens de Forking Globais

1. O Problema

Os Grandes Modelos de Linguagem (LLMs) têm melhorado seu desempenho em tarefas de raciocínio ao alocar mais computação no momento do teste (test-time compute), seja gerando sequências mais longas (escalonamento sequencial) ou amostrando múltiplos caminhos de raciocínio em paralelo (escalonamento paralelo). No entanto, existem desafios críticos:

• Overthinking: O escalonamento sequencial excessivo pode levar a um desempenho degradado após certo comprimento de sequência.
• Diversidade vs. Precisão: Para problemas complexos, os tokens que iniciam caminhos de raciocínio corretos e diversos (chamados de forking tokens) geralmente estão profundamente na árvore de amostragem. Estratégias comuns para aumentar a diversidade, como o aumento da temperatura (temperature scaling), criam um trade-off prejudicial: aumentam a diversidade, mas reduzem drasticamente a precisão, pois o modelo pode explorar caminhos incorretos.
• Colapso de Modos: Ao treinar modelos com múltiplos traços de raciocínio (trilhas) usando Supervised Fine-Tuning (SFT) padrão, o modelo tende a "colapsar" esses modos diversos em um único padrão de raciocínio, perdendo a capacidade de gerar soluções alternativas corretas.

2. Metodologia Proposta

Os autores propõem tratar o raciocínio paralelo como um problema de predição de um conjunto de próximos tokens e introduzem duas técnicas principais:

A. Tokens de Forking Globais (Global Forking Tokens)
Em vez de depender da amostragem aleatória para encontrar o ponto de divergência no meio da geração, o modelo é condicionado a um conjunto reservado de tokens especiais (ex: <think 1>, <think 2>, ..., <think N>) no início da geração. Cada token deve ativar um modo de raciocínio distinto e correto.

B. Set Supervised Fine-Tuning (SSFT)
O SSFT é uma nova abordagem de ajuste fino que utiliza uma perda baseada em conjuntos (set-based loss) para aprender a associar tokens de forking específicos a traços de raciocínio distintos.

• Mecanismo: Para cada pergunta, o modelo recebe $M$ traços de raciocínio corretos (ground-truth) e $N$ tokens de forking globais (onde $N \ge M$).
• Correspondência Bipartida Ótima: O algoritmo calcula a perda de predição de próximo token (NTP) para todas as combinações possíveis entre os $N$ tokens e os $M$ traços. Em seguida, utiliza o Algoritmo Húngaro para encontrar a correspondência bipartida de custo mínimo que associa cada traço de raciocínio ao token de forking que melhor o prediz.
• Função de Perda: A perda total é a soma das perdas NTP sob essa correspondência ótima. Isso força o modelo a aprender que o token <think i> deve gerar especificamente o traço de raciocínio $r_j$ correspondente, preservando a diversidade sem colapso.
• Vantagem: Diferente do SFT padrão, onde os traços são tratados como dados independentes, o SSFT aprende a estrutura de controle global, garantindo que cada token de forking inicie um modo de raciocínio único.

C. Global Forking Policy Optimization (GFPO)
Após o SSFT, os autores aplicam uma etapa leve de Otimização de Política (RL) focada apenas na seleção do token de forking ideal ($g(i)$) para uma dada pergunta. Isso afina a distribuição de saída do modelo para escolher o token que maximiza a complexidade e a precisão do raciocínio, sem re-treinar toda a rede neural.

3. Contribuições Principais

1. Introdução de Tokens de Forking Globais: Um mecanismo para controlar explicitamente a diversidade do raciocínio via tokens de entrada reservados, eliminando a dependência de encontrar pontos de divergência aleatórios durante a geração.
2. Framework SSFT: Uma nova função de perda que utiliza correspondência bipartida ótima para treinar modelos a mapear tokens de controle para traços de raciocínio diversos, prevenindo o colapso de modos observado no SFT tradicional.
3. GFPO: Uma técnica eficiente de RL que otimiza a seleção do token de forking, permitindo que o modelo escolha dinamicamente o melhor modo de raciocínio para cada problema.
4. Evidência Empírica: Demonstração de que o SSFT preserva modos de raciocínio distintos (visualizado através de diferentes comprimentos de "pensamento" e estratégias) e supera métodos baseados em SFT e RL padrão.

4. Resultados Experimentais

Os modelos foram avaliados em benchmarks de raciocínio matemático (AIME 2024/25, MATH-500, GPQA-Diamond) e geração de código (LiveCodeBench).

• Desempenho em Pass@1 (Geração Única): O modelo SSFT-32B superou consistentemente os baselines (incluindo SFT com múltiplos traços e modelos como o Multiverse-32B).
  • No AIME 2024, atingiu 64.06% (vs. 58.23% do SFT misto).
  • No AIME 2025, atingiu 58.13% (vs. 51.96% do SFT misto).
• Desempenho em Cons@k (Consistência com Paralelismo): Ao usar 6 ou 32 gerações paralelas (uma para cada token de forking) e votação majoritária, o SSFT demonstrou uma cobertura superior.
  • No AIME 2025, o Cons@6 do SSFT foi de 73.94%, enquanto o SFT misto precisou aumentar a temperatura para 1.0 (perdendo precisão em Pass@1) para tentar igualar a cobertura.
• Robustez e Generalização:
  • O método funcionou bem em diferentes escalas de modelo (Qwen2.5-32B, Qwen3-4B, Llama3.1-8B).
  • Mostrou forte generalização out-of-distribution (OOD), especialmente em tarefas de codificação (LiveCodeBench) quando treinado apenas com dados matemáticos, e vice-versa.
  • A ablação mostrou que a correspondência ótima (Hungarian) é crucial; o SSFT com correspondência aleatória não superou significativamente os baselines.

5. Significado e Impacto

Este trabalho representa um avanço significativo na forma como os LLMs utilizam a computação no momento do teste (test-time compute):

• Superação do Trade-off Diversidade-Precisão: Ao aprender tokens de forking globais, o modelo consegue gerar múltiplos caminhos de raciocínio corretos e distintos sem sacrificar a precisão individual de cada caminho, algo que o aumento de temperatura falha em fazer.
• Eficiência de Treinamento: O SSFT permite que modelos menores aprendam a "pensar" de múltiplas formas complexas a partir de dados de destilação de modelos maiores, sem a necessidade de RL pesado durante todo o treinamento.
• Controle Explícito: Oferece uma alavanca controlável para a diversidade do raciocínio, permitindo que sistemas de IA explorem diferentes estratégias (ex: álgebra vs. geometria, ou diferentes abordagens de código) de forma determinística baseada no token de entrada.

Em resumo, o SSFT transforma o raciocínio paralelo de uma busca estocástica ineficiente em um processo estruturado e aprendível, onde o modelo internaliza "modos de pensamento" distintos e acessíveis via tokens de controle, resultando em modelos mais robustos e capazes de resolver problemas complexos com maior consistência.