Resumo Técnico: LARK (Seleção de Trajetória Baseada em Aprendibilidade)

Declaração do Problema

O artigo aborda o desafio da destilação de raciocínio, um paradigma onde um modelo de linguagem grande (LLM) estudante menor é ajustado (fine-tuned) para imitar trajetórias de raciocínio geradas por modelos professores maiores. Embora os métodos de seleção de dados existentes frequentemente operem ao nível da questão ou dependam de heurísticas como qualidade da trajetória, correção ou confiança do modelo, eles frequentemente negligenciam um fator crítico: a aprendibilidade.

Especificamente, para uma única questão, podem existir múltiplas trajetórias de raciocínio candidatas (geradas por diferentes professores, amostragens ou estilos). Mesmo que essas trajetórias levem à mesma resposta correta final, elas fornecem sinais de treinamento vastamente diferentes. Uma trajetória que parece ser de "alta qualidade" ou "natural" não é necessariamente aquela da qual o modelo estudante pode aprender de forma mais eficiente. A questão central posta é: Podemos projetar um critério de seleção de trajetória baseado em princípios que garanta que as trajetórias de raciocínio destiladas sejam aprendíveis pelo modelo estudante?

Metodologia: LARK

Os autores propõem o LARK (Rank de Tempo de Âncora baseado em Aprendibilidade), um método que seleciona trajetórias com base em uma definição de aprendibilidade fundamentada na teoria da otimização, em vez de métricas de qualidade estáticas.

1. Aprendibilidade como um Objetivo de Otimização

O LARK enquadra a seleção de trajetórias como um problema de otimização de política. Ele define um fator de aprendibilidade $\rho(\theta; q)$, que caracteriza a taxa na qual a perda de treinamento do estudante diminui. Este fator é definido como a razão entre a norma do gradiente ao quadrado e a perda:
$$\rho(\theta; q) \triangleq \frac{\|\nabla_\theta \mathcal{L}(\theta; q)\|^2}{\mathcal{L}(\theta; q)}$$
onde $q$ é o vetor de pesos sobre $K$ trajetórias candidatas. Sob a condição Polyak–Łojasiewicz (PL), um $\rho$ mais alto implica uma convergência mais rápida.

2. Aproximação de Tempo de Âncora

Otimizar diretamente o objetivo de aprendibilidade integrado no tempo ao longo de toda a trajetória de ajuste fino é intratável. O LARK introduz a condição relativa à âncora, que postula que, para modelos superparametrizados no regime de "treinamento preguiçoso" (lazy training), a taxa de aprendibilidade permanece proporcional ao seu valor nos parâmetros de âncora iniciais $\theta_{ref}$. Isso permite que o objetivo seja aproximado pela maximização de $\rho(\theta_{ref}; q)$.

3. Linearização Local e Proxy de Passagem Direta (Forward-Pass)

Maximizar $\rho(\theta_{ref}; q)$ diretamente é problemático porque é uma função quase-convexa que colapsa para uma solução "one-hot" degenerada (selecionando apenas uma única trajetória). Para evitar isso, o LARK:

• Linearização Local: Expande $\rho$ em torno da distribuição uniforme $p$ (SFT padrão sem seleção) usando uma expansão de Taylor de primeira ordem.
• Surrogate de Passagem Direta ($\hat{g}_k$): Em vez de calcular gradientes exatos (que exigem retropropagação/backward passes para cada candidato), o LARK deriva um proxy computável $\hat{g}_k$ usando apenas passagens diretas (forward passes). Este proxy estima o gradiente local da taxa de aprendibilidade com base nos resíduos de predição de token do estudante e na perda de entropia cruzada.
  • Especificamente, o $\hat{g}_k$ favorece trajetórias onde o estudante está "confiantemente errado" de forma estruturada: a perda $\ell_k$ é alta (indicando uma má predição), mas a massa residual está concentrada em alguns poucos tokens de "quase acerto" (fornecendo um sinal de correção claro).

4. Política de Seleção com Regularização $\chi^2$

Para evitar que o seletor "hackeie" o critério de aprendibilidade ao se desviar demais da distribuição original dos dados (o que prejudica a generalização), o LARK impõe um termo de regularização $\chi^2$. Isso leva a uma regra de seleção de forma fechada e parametrizada pelo orçamento:

• As trajetórias são classificadas por seu escore $\hat{g}_k$.
• As $B$ melhores trajetórias são selecionadas.
• Ao contrário da truncagem rígida de top-$B$, o LARK atribui pesos suaves proporcionais à margem do escore acima da $(B+1)$-ésima trajetória. Esta regra de ponderação é derivada para maximizar o limite inferior da melhoria da aprendibilidade enquanto controla a divergência $\chi^2$ em relação ao prior uniforme.

Principais Contribuições

1. Formulação: Os autores formulam a seleção de trajetória de raciocínio como um problema de otimização de política baseado em aprendibilidade, introduzindo a taxa de aprendibilidade de tempo de âncora $\rho$ como um objetivo fundamentado.
2. Algoritmo: Eles derivam uma regra de seleção de forma fechada e parametrizada pelo orçamento que utiliza um surrogate de passagem direta $\hat{g}_k$ e regularização $\chi^2$. Isso evita o custo computacional de retropropagações por trajetória, enquanto garante teoricamente limites sobre o erro de estimativa e generalização.
3. Validação Empírica: O artigo demonstra que o LARK supera consistentemente as linhas de base existentes (incluindo métodos baseados em heurística, qualidade e pontuação do lado do estudante) em múltiplos modelos base (Qwen-2.5-7B/1.5B, Llama-3.2-3B) e benchmarks de raciocínio (AIME, AMC, GPQA, MATH-500).

Resultados Experimentais

• Desempenho: O LARK alcançou o maior Acc@5 médio em todas as configurações de orçamento de modelo testadas. Por exemplo, no Qwen-2.5-7B com um orçamento de $B=3$, o LARK melhorou o desempenho em 7,07 pontos percentuais sobre a linha de base mais forte (RSR).
• Eficiência: O LARK é um método de apenas passagem direta (forward-pass-only), igualando o custo computacional das linhas de base de menor custo do lado do estudante, enquanto métodos baseados em gradientes exatos exigiriam muito mais computação.
• Análise Diagnóstica:
  • Decaimento da Perda: As trajetórias selecionadas pelo LARK induziram a redução mais rápida da perda de SFT durante o treinamento.
  • Validade Preditiva: O escore LARK $\hat{g}_k$ mostrou uma forte correlação positiva (Pearson $r=0,69$) com a precisão downstream quando agregada por modelo professor.
  • Comportamento ao Nível de Token: As trajetórias selecionadas exibiram probabilidades de tokens corretos menores, mas maior "concentração de tokens errados", confirmando que o LARK seleciona exemplos onde o estudante está errado, mas o erro é estruturado e corrigível.

Significância e Alegações

O artigo afirma que a aprendibilidade é um princípio distinto e eficaz para a destilação de raciocínio eficiente, separado da qualidade ou dificuldade da trajetória. Ao fundamentar a seleção na dinâmica de otimização do modelo estudante, o LARK identifica trajetórias que fornecem o sinal de aprendizado mais eficiente.

Os autores enfatizam que seu método:

• Preserva a Generalização: Através da regularização $\chi^2$, evita que a política de seleção sofra overfitting para um conjunto estreito de trajetórias "fáceis" ou "hackeadas".
• É Teoricamente Fundamentado: A regra de seleção é derivada de primeiros princípios da teoria da otimização (condições PL, fluxo de gradiente) com limites de erro explícitos.
• É Prático: Não requer ajuste de hiperparâmetros para a própria regra de seleção (a força da regularização é determinada pelo orçamento $B$) e evita o overhead computacional de retropropagações.

O trabalho sugere que extensões futuras poderiam aplicar essa perspectiva baseada em aprendibilidade para pools de candidatos mais ruidosos e domínios de raciocínio mais amplos, mas o escopo atual é estritamente limitado à destilação de raciocínio na resolução de problemas matemáticos.