OS-Pruner: Pruning Chains-of-Thought of Reasoning Models via Optimal Stopping

Resumo Técnico: OS-Pruner

Declaração do Problema

Grandes Modelos de Linguagem (LLMs) que utilizam prompting de Cadeia de Pensamento (Chain-of-Thought - CoT) alcançaram um sucesso significativo em tarefas de raciocínio complexo. No entanto, esses modelos frequentemente exibem "sobrepensamento computacional" (computational overthinking), gerando etapas de raciocínio redundantes que aumentam a latência e os custos de inferência sem melhorar a precisão da resposta final. Esse comportamento inclui a reiteração de argumentos, a verificação de conclusões já estabelecidas ou o seguimento de cálculos de baixo valor. Embora estudos recentes sugiram que as trajetórias de CoT podem ser podadas, as soluções existentes sofrem de limitações específicas:

• Orçamentos Estáticos: Métodos como a imposição de orçamento (budget forcing) impõem um limite fixo de pensamento externamente, ignorando o estado evolutivo do traço de raciocínio.
• Compressão do Lado do Modelo: Técnicas como ajuste fino com penalização de comprimento ou salto de tokens exigem o retreinamento custoso do modelo base e podem degradar o desempenho em tarefas que não são de raciocínio.
• Parada Antecipada Heurística: Métodos dinâmicos atuais frequentemente dependem de classificação ou limiares (por exemplo, parar quando a confiança excede um valor fixo), o que falha em considerar o equilíbrio entre o custo de tokens adicionais e a melhoria esperada na qualidade da resposta.

Os autores argumentam que a poda de CoT é fundamentalmente um problema de decisão sequencial, que se enquadra melhor como um problema de parada ótima (optimal stopping problem), onde o sistema deve decidir em cada etapa se o ganho esperado em precisão justifica o custo de gerar outro token.

Metodologia: OS-Pruner

O OS-Pruner é um framework leve e plug-in projetado para formular a poda de CoT como um problema de parada ótima. Ele opera sem modificar o modelo de raciocínio base congelado.

1. Formulação do Problema

Dado um problema de entrada $x$ e um prefixo de CoT parcialmente gerado $y_{\le i}$, o OS-Pruner aprende uma política de parada paramétrica $\pi_\theta$ que retorna a probabilidade de parada. O objetivo é maximizar uma função de utilidade explícita que equilibra a precisão da resposta final contra o comprimento gerado:
$$r(y_{\le i}|x) = A(y_{\le i}|x) - \lambda L(y_{\le i})$$
Onde:

• $A(y_{\le i}|x)$ é a precisão da resposta derivada do prefixo (1 se correta, 0 caso contrário).
• $L(y_{\le i})$ é o número de tokens gerados até o passo $i$.
• $\lambda$ é um parâmetro escalar que controla a penalidade de comprimento, permitindo que os usuários selecionem um ponto operacional específico na fronteira de precisão-eficiência.

2. Vantagem Teórica

O artigo distingue o OS-Pruner dos métodos de limiar de confiança através de uma Visão de Função de Valor.

• Métodos baseados em classificação param se $A(y_{\le i}|x) \ge \gamma$ (um limiar fixo).
• Parada Ótima (OS-Pruner) para se $A(y_{\le i}|x) \ge C_\lambda(x, y_{\le i})$.
  O valor de continuação $C_\lambda$ representa a recompensa futura esperada menos o custo de tokens adicionais. Os autores provam (Teorema 1) que classificadores de limiar fixo podem ter uma lacuna de desempenho arbitrariamente grande em comparação com a política de parada ótima, pois falham em considerar o valor dependente do estado dos caminhos de raciocínio futuros.

3. Implementação e Arquitetura

• Pré-processamento de Dados: O modelo base gera traços de raciocínio completos. Para o treinamento, o sistema intervém após cada etapa de raciocínio, força o modelo a fornecer uma resposta final (impedindo o "pensamento" na seção de resposta) e a avalia contra a verdade fundamental (ground truth) usando o Math-Verify. Isso cria uma supervisão granular para cada prefixo.
• Arquitetura da Política: A política é uma cabeça linear leve ($\sigma(W^\top h_l + b)$) aplicada ao último estado oculto $h_l$ do prefixo. Para aumentar o desempenho, as últimas $n$ (tipicamente 2) camadas de autoatenção do modelo base são ajustadas (fine-tuned), enquanto o restante do modelo permanece congelado.
• Eficiência de Treinamento: Como o modelo base é causal, o estado oculto de um parágrafo não é afetado por tokens subsequentes. Isso permite que todo o traço de raciocínio seja processado em uma única passagem direta (forward pass), com a camada linear aplicada em paralelo aos estados ocultos relevantes.
• Treinamento Progressivo: Para evitar ficar preso em mínimos locais onde a política nunca para, o treinamento começa com um $\lambda$ alto (alto custo de comprimento) e o reduz progressivamente (annealing) para valores mais baixos.

4. Inferência

Durante a inferência, a política é invocada apenas nos limites das etapas de raciocínio (ex: fim de um parágrafo). O sistema verifica a probabilidade de parada; se a decisão for parar, ele solicita a resposta final imediatamente. Este design garante alta eficiência em sistemas modernos de serviço de LLM (ex: vLLM) ao aproveitar o processamento paralelo de preenchimento (pre-fill) e minimizar a recomputação do cache KV.

Principais Contribuições

1. Formulação Rigorosa: O artigo fornece uma formalização da poda de CoT como um problema de parada ótima com uma função de recompensa explícita que equilibra precisão versus comprimento.
2. Plug-in Leve: O OS-Pruner é uma política de parada baseada em treinamento que opera sobre prefixos de raciocínio sem modificar o modelo de raciocínio base congelado, exigindo apenas o ajuste fino de uma pequena cabeça de política e de algumas camadas do transformer.
3. Desempenho Empírico: Experimentos demonstram que o OS-Pruner reduz substancialmente o comprimento do raciocínio preservando a precisão em diversos modelos e benchmarks.
4. Complementaridade: O método mostra prover ganhos adicionais mesmo quando aplicado a modelos já otimizados para brevidade via métodos do lado do modelo.

Resultados Experimentais

Os autores avaliaram o OS-Pruner em três modelos (DeepSeek-R1-Distill-Qwen-7B, GPT-OSS-20B e DRPO-7B) em benchmarks de diferentes dificuldades (GSM8K, MATH-500, AIME).

• Redução de Comprimento: O OS-Pruner alcançou uma redução de 20 a 60% no comprimento da geração com sacrifício mínimo de precisão.
  • No DeepSeek-R1-Distill-Qwen-7B, reduziu o comprimento em 59,3% no GSM8K e 52,8% no MATH-500, com mudanças de precisão de apenas -0,7% e +2,7%, respectivamente.
  • No GPT-OSS-20B (um modelo já otimizado para concisão), ainda alcançou reduções significativas (ex: 21% no MATH-500) mantendo uma precisão competitiva.
• Comparação com Baselines: O OS-Pruner superou consistentemente os baselines baseados em treinamento (Classification Probe, Answer Convergence, FlashThink) na Pontuação de Precisão-Eficiência (AES). Ele situa-se na ou próxima da fronteira de Pareto, enquanto os baselines de limiar fixo degradaram de forma mais acentuada à medida que os orçamentos de comprimento eram reduzidos.
• Robustez: O método funcionou bem em diferentes níveis de dificuldade, sendo mais agressivo em tarefas fáceis (onde o sobrepensamento é prevalente) e mais conservador em tarefas difíceis (AIME), onde o cálculo adicional é valioso.

Significância e Alegações

O artigo afirma que o OS-Pruner oferece uma abordagem mais natural e ótima para a poda de CoT do que os métodos heurísticos ou estáticos existentes. Ao enquadrar o problema como um processo de decisão sequencial, ele avalia dinamicamente o ponto suficiente de terminação com base no estado de raciocínio específico, em vez de um limiar global.

Os autores enfatizam que o OS-Pruner:

• Reduz o Custo Computacional: Ao cortar tokens desnecessários, diminui a latência de inferência e a pegada energética.
• Fornece Controle Granular: O parâmetro $\lambda$ permite que os usuários ajustem explicitamente o equilíbrio entre o esforço de raciocínio e a precisão.
• Funciona com Modelos Existentes: Não requer o retreinamento do gerador base, tornando-o um "plug-in" prático para modelos de raciocínio existentes.
• Complementa Métodos do Lado do Modelo: Mesmo modelos treinados para brevidade (como o DRPO) beneficiam-se da poda adicional, sugerindo que a compressão do lado do modelo e a parada ótima abordam aspectos diferentes de eficiência.

O trabalho conclui que, embora foque no raciocínio matemático, o framework é aplicável a contextos mais amplos, embora trabalhos futuros sejam necessários para explorar tarefas de codificação e raciocínio científico e a escala para modelos de fronteira.