Do LLMs Encode Functional Importance of Reasoning Tokens?

Título: Do LLMs Encode Functional Importance of Reasoning Tokens?

Autores: Janvijay Singh e Dilek Hakkani-Tür (University of Illinois Urbana Champaign)

1. Problema e Motivação

Os Grandes Modelos de Linguagem (LLMs) dependem de longas cadeias de raciocínio explícito (Chain-of-Thought) para resolver tarefas complexas em matemática, ciências e tomada de decisão. Embora essas cadeias longas melhorem a precisão, elas impõem custos significativos:

• Maior latência de inferência e requisitos de memória.
• Dificuldade em isolar quais partes do raciocínio são funcionalmente responsáveis pela resposta correta.

Trabalhos anteriores sobre "raciocínio compacto" focam em encurtar essas cadeias usando amostragem probabilística, heurísticas ou supervisão de modelos de ponta (frontier models). No entanto, essas abordagens oferecem pouca insight sobre se o modelo internamente codifica a importância funcional de cada token para a geração da resposta. A questão central deste trabalho é: Os LLMs codificam internamente a importância funcional dos tokens de raciocínio para a geração de respostas?

2. Metodologia: "Greedy Pruning" (Poda Gananciosa)

Os autores propõem uma nova abordagem diagnóstica chamada Greedy Pruning, um procedimento de deleção que preserva a verossimilhança (likelihood).

• Conceito Central: Em vez de gerar novas cadeias curtas, o método começa com uma cadeia de raciocínio completa gerada por um modelo "Professor" e remove iterativamente tokens cuja exclusão causa a menor degradação possível na verossimilhança do modelo sob um objetivo específico.
• O Algoritmo:
  1. Começa com a cadeia completa de $N$ tokens.
  2. Para cada passo, avalia a remoção de cada token candidato.
  3. Calcula a verossimilhança pós-deleção ($L_{del}$).
  4. Remove o token que minimiza a queda na verossimilhança (ou seja, o token "menos importante" para a probabilidade do modelo).
  5. Repete o processo até atingir uma fração de retenção ($\rho$) desejada.
• Objetivos de Poda:
  1. Preservação da Resposta ($L_{ANS}$): Mantém apenas a probabilidade da resposta final.
  2. Preservação Conjunta ($L_{JOINT}$): Mantém a probabilidade tanto do raciocínio quanto da resposta, preservando a trajetória de raciocínio do modelo.
• Ranking de Poda: O processo gera uma ordem global de poda ($\pi$), onde tokens removidos mais cedo são considerados funcionalmente menos importantes para a geração da resposta sob a distribuição do próprio modelo.

3. Contribuições Principais

1. Método Diagnóstico: Introduz a poda gananciosa como uma ferramenta para revelar a estrutura funcional interna dos tokens de raciocínio, sem depender de supervisão externa (como GPT-4) para rotular a importância.
2. Framework de Distilação: Propõe um ciclo "Professor-Podador-Aluno", onde modelos menores (alunos) são treinados em cadeias de raciocínio podadas pelo método.
3. Análise de Estrutura Funcional: Demonstra que a poda revela padrões sistemáticos: tokens de cálculo simbólico (SYMBMATH) são preservados desproporcionalmente, enquanto tokens de referência (COREF), gramática e narrativa (METADISC) são removidos precocemente.
4. Previsibilidade via Atenção: Mostra que os rankings de poda podem ser previstos com alta correlação (Pearson 0.88) apenas a partir dos scores de atenção do modelo, sugerindo que a importância funcional é acessível através dos mecanismos internos do modelo.

4. Resultados Experimentais

Os experimentos foram realizados em benchmarks de raciocínio (GSM8K, MATH, MMLU-Pro) utilizando modelos como Llama-3.1, Qwen-2.5, Llama-2 e Mistral.

• Desempenho na Distilação:
  • Alunos treinados em cadeias podadas pelo método Greedy superaram consistentemente várias linhas de base (baselines), incluindo TokenSkip (que usa supervisão de modelos de ponta), H2O (baseado em atenção acumulada) e Surprisal.
  • Em comprimentos de raciocínio equivalentes, a poda gananciosa resultou em maior precisão do aluno, indicando que ela preserva melhor os tokens funcionalmente críticos.
• Estrutura Funcional:
  • A análise de retenção por categoria mostrou que a poda gananciosa cria uma ordem funcional clara: SYMBMATH (equações, números) é mantido até o final, enquanto COREF (pronomes) e GRAMMAR são os primeiros a ser removidos.
  • O uso do objetivo JOINT (preservar raciocínio + resposta) gerou uma estrutura funcional mais nítida e melhor desempenho de distilação do que o objetivo apenas de resposta (ANS).
  • A força do modelo "Podador" afeta a qualidade da estrutura: podadores mais fortes geram rankings de importância mais confiáveis.
• Dinâmica de Importância:
  • A importância dos tokens não é estática; ela é reavaliada dinamicamente à medida que o contexto encolhe. Rankings locais calculados em estágios intermediários de poda são mais preditivos do que rankings globais fixos.
• Integridade Semântica:
  • Mesmo com poda agressiva (redução de ~30% dos tokens), as cadeias mantêm a coerência semântica e o estado matemático, embora sofram degradação superficial (pontuação, palavras de ligação).

5. Significado e Implicações

• Validação de Importância Interna: O trabalho prova que LLMs não apenas geram raciocínio longo, mas codificam internamente uma estrutura de importância funcional nos tokens. Tokens críticos para o cálculo são distintos de tokens de "enfeite" linguístico.
• Ferramenta de Interpretabilidade: A poda gananciosa serve como um "probe" (sonda) para entender como o modelo organiza seu raciocínio, oferecendo uma alternativa mais controlada e explicativa do que métodos baseados apenas em amostragem estocástica.
• Eficiência e Distilação: O método oferece uma via principled (fundamentada) para comprimir dados de treinamento, permitindo que modelos menores aprendam raciocínio eficiente sem perder a precisão, superando métodos que dependem de modelos externos caros para compressão.
• Futuro: Sugere que sinais de atenção podem ser usados para prever a importância de tokens em tempo real, abrindo caminho para técnicas de poda on-the-fly durante a inferência.

Conclusão

O artigo estabelece que os LLMs possuem uma estrutura de importância funcional codificada em seus tokens de raciocínio. Através da Greedy Pruning, é possível expor e explorar essa estrutura para comprimir cadeias de raciocínio de forma eficaz, melhorando a eficiência computacional e a capacidade de distilação de modelos, enquanto revela que a "essência" do raciocínio reside principalmente nos componentes simbólicos e matemáticos, e não na narrativa linguística circundante.