Resumo Técnico: Escalonamento de Teste de Modelo de Linguagem Eficiente e Treinável via Roteamento de Ramificação Local

Declaração do Problema

Avanços recentes no raciocínio de modelos de linguagem (LLMs) demonstram que alocar computação adicional no tempo de inferência (escalonamento de tempo de teste) pode melhorar significativamente o desempenho. No entanto, as abordagens existentes enfrentam um compromisso crítico:

• Amostragem de Cadeia de Pensamento (CoT): Embora eficaz, a CoT padrão permanece de thread única, comprometendo-se com um único token antes de observar as consequências das alternativas.
• Métodos Baseados em Busca: Abordagens como Árvore de Pensamentos (Tree of Thoughts) ou Busca em Árvore de Monte Carlo (MCTS) introduzem largura ao explorar árvores de pensamentos intermediários. Contudo, estas são frequentemente de grão grosso, computacionalmente caras e difíceis de otimizar como uma política de decodificação treinável de ponta a ponta.
• Métodos de Tokens Suaves (Soft-Token): Técnicas como Soft Thinking ou Multiplex Thinking oferecem largura ao nível do token através da fusão de embeddings candidatos em misturas contínuas. Embora treináveis, estes métodos obscurecem a identidade discreta das ramificações candidatas, podendo ocultar informações específicas de conceitos.

O artigo postula que os modelos de linguagem poderiam beneficiar-se de um mecanismo de escalonamento de tempo de teste por nível de token, leve, que permita que as decisões sejam informadas pelos estados ocultos de candidatos locais futuros sem incorrer no custo de uma busca de solução completa ou na ambiguidade de misturas suaves.

Metodologia: Roteamento de Ramificação Local (LBR)

Os autores propõem o Roteamento de Ramificação Local (LBR), um framework de decodificação de ramificação e roteamento local, treinável e ao nível do token. O LBR aumenta a geração autoregressiva padrão com uma árvore de olhar à frente (lookahead) local e rotativa.

Pipeline de Decodificação

Em cada etapa de decodificação $t$, dado um prefixo comprometido $x_{<t}$, o LBR executa um processo de quatro estágios "podar–deslocar–cultivar" (prune–shift–grow):

1. Cultivar (Grow): Em vez de amostrar um único próximo token, o LBR expande uma árvore local $\mathcal{A}_t$ de largura $K$ e profundidade $L$ enraizada em $x_{<t}$. Cada nó amostrado é encaminhado através do modelo de linguagem para computar seu estado oculto $h_t(v)$. Isso cria uma árvore de estados ocultos locais em vez de apenas um conjunto de strings candidatas.
2. Rotear (Route): Um roteador leve $\rho_\phi$ observa a árvore encaminhada e seleciona uma subárvore de profundidade 1 para comprometer. O roteador utiliza um mecanismo de atenção de conjunto (set-attention) para codificar cada subárvore candidata (por exemplo, um único token para $L=1$, ou um token mais suas continuações para $L>1$) em um vetor. Esses vetores são comparados via atenção entre subárvores para produzir uma pontuação relativa, permitindo que o roteador tome uma decisão contrastiva entre futuros irmãos.
3. Podar (Prune): O roteador seleciona a melhor subárvore (índice $k^*_t$). Todas as outras subárvores são descartadas. O token raiz da subárvore selecionada torna-se o próximo token comprometido $x_t$.
4. Deslocar e Recultivar (Shift and Regrow): A subárvore selecionada é deslocada para frente (sua raiz é removida e seus descendentes tornam-se o novo prefixo para o próximo passo). Uma nova camada é cultivada para restaurar a profundidade da árvore $L$, e o ciclo se repete.

Objetivo de Treinamento

O LBR define uma verossimilhança de trajetória de árvore tratável que permite o aprendizado por reforço de ponta a ponta com recompensas verificáveis (RLVR).

• Estocasticidade: Apenas duas operações são estocásticas: a amostragem dos novos nós $G_t$ crescidos a partir do modelo base e a amostragem da decisão do roteador $k^*_t$. A poda, o deslocamento e o reuso de nós são determinísticos.
• Verossimilhança: O log-probabilidade de uma trajetória $F$ é a soma dos log-probabilidades dos novos tokens crescidos e das probabilidades de seleção do roteador:
  $$\log p_{\theta,\phi}(F | q) = \sum_{t=1}^T \sum_{v \in G_t} \log \tilde{\pi}_\theta(\text{tok}(v) | \text{ctx}(v)) + \sum_{t=1}^T \log \rho_\phi(k^*_t | x_{<t}, \mathcal{A}_t; \theta)$$
• Otimização: O framework é treinado usando um objetivo de razão de verossimilhança (ex: GRPO) onde a vantagem $A(F)$ de um verificador recompensa tanto a geração de tokens do modelo base quanto as decisões de seleção do roteador.

Principais Contribuições

1. Ramificação ao Nível do Token com Identidade Discreta: Ao contrário dos métodos de tokens suaves que fundem candidatos em misturas contínuas, o LBR preserva os candidatos como ramificações encaminhadas discretas. Isso permite que o roteador explore os estados ocultos pós-candidato — informação revelada apenas após um token candidato ser gerado e processado pelo modelo.
2. Treinamento de Ponta a Ponta Tratável: Ao definir uma verossimilhança específica para a trajetória da árvore, o LBR permite que o modelo base e o roteador sejam otimizados conjuntamente sob o mesmo princípio de RLVR usado em sistemas de raciocínio de tokens discretos.
3. Escalonamento de Teste Eficiente: O framework introduz largura ao nível do token com um olhar à frente raso ($L$ é pequeno, ex: 1 ou de 2), evitando o overhead computacional de uma busca de profundidade total, enquanto ainda aproveita evidências além da distribuição do próximo token da raiz.
4. Roteador de Atenção de Conjunto: Os autores introduzem uma arquitetura de roteador que usa atenção entre subárvores para comparar candidatos irmãos, demonstrando que o roteamento é uma decisão contrastiva e não um reranking independente.

Resultados Experimentais

Planejamento Hierárquico Sintético

Em uma tarefa sintética envolvendo alcançabilidade de grafo com tradução de raio:

• Desempenho: O LBR alcançou a maior precisão alvo (0.967), superando a CoT Discreta (0.914) e o Soft Thinking (0.845).
• Análise de Mecanismo: Experimentos de sondagem revelaram que os estados ocultos pós-candidato no LBR são significativamente mais preditivos do alvo final do que os estados pré-ramificação. Isso confirma que o encaminhamento de candidatos permite que o modelo acesse informações a jusante que não estão disponíveis para a decodificação autoregressiva padrão.
• Representação: Diferente do Soft Thinking, que borrou identidades de conceitos via misturas contínuas, o LBR preservou identidades de ramificações discretas, tornando o caminho correto mais recuperável.

Benchmarks de Raciocínio Matemático

Avaliado em modelos DeepSeek-R1-Distill-Qwen (1.5B e 7B) em seis benchmarks (Minerva, AIME, MATH500, etc.):

• Ganhos de Desempenho: O LBR melhorou tanto o Pass@1 quanto o Pass@32 em relação a:
  • Cadeia de Pensamento Discreta.
  • RLVR de token discreto vanilla.
  • Baselines de ramificação de token suave compatíveis com RL (ex: Multiplex Thinking).
• Ablação de Profundidade: Aumentar a profundidade do olhar à frente de $L=1$ para $L=2$ melhorou ainda mais o desempenho, particularmente no backbone de 7B, sugerando que continuações locais mais profundas fornecem mais evidências de roteamento.
• Ablação do Roteador: Remover a atenção entre subárvores (usando um roteador independente) resultou em menor desempenho, validando que a comparação contrastiva de ramos irmãos é crucial.
• Dinâmica de Treinamento: Durante o treinamento RLVR, a entropia do roteador diminuiu conforme o roteador se tornava mais decisivo, enquanto a entropia do modelo base nas posições roteadas também diminuiu, indicando uma co-adaptação onde o modelo aprende a gerar ramificações locais mais distinguíveis.

Significância e Alegações

O artigo afirma que o Roteamento de Ramificação Local oferece uma forma eficiente, treinável e discreta de escalonamento de tempo de teste para modelos de linguagem. Sua principal significância reside em demonstrar que:

1. Evidência pós-candidato é valiosa: As decisões de roteamento beneficiam-se dos estados ocultos induzidos pelos tokens candidatos, não apenas da distribuição pré-ramificação.
2. Ramificação discreta é superior à fusão suave: Preservar identidades de ramificação discretas permite uma tomada de decisão mais clara em comparação com misturas de tokens suaves, que podem obscurecer a identidade do conceito.
3. Escalabilidade: O framework fornece um meio-termo entre a CoT de thread única e a busca de árvore completa dispendiosa, permitindo um escalonamento eficaz através de ramificação local leve.

Os autores reconhecem a limitação de localidade: o LBR é restrito pela janela de olhar à frente rasa e não realiza uma busca global sobre soluções completas. Sugere-se como trabalho futuro combinar o LBR com profundidade adaptativa ou busca global seletiva.