Resumo Técnico: Desvendando Estratégias de Raciocínio Latentes em Modelos de Linguagem

Declaração do Problema

Modelos de linguagem (LMs) pré-treinados em tarefas de raciocínio frequentemente resolvem problemas usando múltiplas estratégias distintas e de alto nível (por exemplo, prova por contradição vs. indução, ou diferentes algoritmos de ordenação). No entanto, essas estratégias são implícitas e emaranhadas dentro da distribuição de resposta do modelo $p_\theta(y|x)$. Não existe uma variável explícita para identificar ou controlar qual estratégia está sendo usada durante a geração.

O artigo formula o problema de decompor essa distribuição de resposta em uma representação estruturada e condicionada à estratégia. O objetivo é aprender uma fatoração de roteador-gerador $(r_\phi(z|x), g_\phi(y|x, z))$ tal que:

1. Fidelidade Marginal: A distribuição integrada $\int r_\phi(z|x)g_\phi(y|x, z) dz$ aproxima a distribuição do modelo base original $p_\theta(y|x)$.
2. Alinhamento de Estratégia: A variável latente $z$ é informativa e alinhada à estratégia, o que significa que variar $z$ altera a estratégia de alto nível utilizada, e o mesmo valor latente carrega um significado de estratégia consistente através de diferentes entradas.

O Desafio Central: Diferente da modelagem de variáveis latentes padrão, onde um modelo é treinado do zero, o modelo base $p_\theta$ já captura a distribuição alvo sem qualquer variável latente. Consequentemente, uma fatoração que ignora $z$ (onde $g_\phi(y|x, z) = p_\theta(y|x)$ para todo $z$) está disponível na inicialização e é globalmente ótima para objetivos que medem apenas o ajuste marginal. A inferência variacional padrão (ex: ELBO condicional) sofre de colapso de posterior neste cenário, falhando em incentivar $z$ a codificar a variação ao nível da estratégia.

Metodologia

Arquitetura

Os autores adaptam um Transformer pré-treinado em uma estrutura de roteador-gerador:

• Roteador ($r_\phi$): Lê a entrada $x$ e produz uma distribuição sobre um $z$ latente contínuo (modelado como uma Gaussiana diagonal).
• Gerador ($g_\phi$): Um modelo autorregressivo que condiciona em $x$ e um embedding projetado $E(z)$ do latente amostrado.
• Posterior ($q_\xi$): Um modelo separado usado apenas durante o treinamento para inferir $z$ a partir do par completo entrada-resposta $(x, y)$.
• Parametrização: O roteador e o gerador compartilham parâmetros inicializados a partir do modelo base $\theta$, atualizados via adaptação de baixo rank do tipo LoRA.

Objetivo de Reconstrução Dirigido pelo Modelo

Para superar a subespecificação do ELBO padrão, os autores propõem um objetivo de reconstrução dirigido pelo modelo. Este objetivo utiliza o modelo base fixo $p_\theta$ como referência para direcionar onde a variável latente $z$ deve explicar a variação.

A função objetivo $J_\theta$ consiste em um termo de reconstrução $R_\theta$ e uma penalidade KL:
$$J_\theta(x, y; \phi, \xi) = R_\theta(x, y; \phi, \xi) + \beta \text{KL}(q_\xi(\cdot|x, y) \parallel r_\phi(\cdot|x))$$

Componentes principais do termo de reconstrução:

1. Normalização do Modelo Base: A perda de reconstrução é normalizada pela perda esperada de resposta do modelo base ($c_\theta$), expressando a perda em "unidades de referência do modelo base".
2. Ponderação Dirigida pela Surpresa (Surprisal): Em vez de uma ponderação uniforme de tokens, a perda é ponderada pela surpresa de nível de token do modelo base ($b_{\theta,t}$).
   • $b_{\theta,t}(x, y) = -\log p_\theta(y_t | x, y_{<t})$
   • Tokens com alta surpresa (pontos de ramificação onde o modelo base é incerto) recebem maior pressão de reconstrução.
   • Isso encoraja $z$ a codificar a variação em posições onde o modelo base deixa múltiplas continuações plausíveis, o que tipicamente corresponde a pontos de escolha de estratégia.

O termo de reconstrução é definido como uma fração ponderada e normalizada da perda do modelo base:
$$R_\theta = \frac{1}{c_\theta} \mathbb{E}_{z \sim q_\xi} \left[ \sum_{t} a_{\theta,t}(x, y) [-\log g_\phi(y_t | x, z, y_{<t})] \right]$$
onde os pesos $a_{\theta,t}$ interpolam entre pressão uniforme e pressão dirigida pela surpresa.

Configuração Experimental e Benchmark

Os autores introduzem um benchmark controlado de tarefas algorítmicas multi-estratégia para tornar o alinhamento de estratégia mensurável.

• Tarefas: Somação de listas, algoritmos de ordenação, busca de caminho em grade (grid pathfinding), resolução de equações lineares, conversão de base e adição de múltiplos dígitos.
• Design: Cada tarefa permite múltiplos traços de solução válidos (estratégias) para a mesma instância do problema. Rótulos de estratégia de verdade fundamental (ground-truth) são gerados, mas ocultados durante o treinamento, sendo usados apenas para avaliação.
• Métricas de Avaliação:
  1. Fidelidade Distribucional: A fração de traços gerados que permanecem soluções válidas para a tarefa.
  2. Alinhamento de Estratégia (Consistência Analógica): A probabilidade de que amostrar um $z$ de uma entrada fonte $X$ e aplicá-lo a uma entrada alvo relacionada $X'$ resulte em um traço de solução $Y'$ usando a mesma estratégia que $Y$ gerado a partir de $X$.

Principais Resultados

1. Informação Implícita Existe: Sondas lineares nos estados ocultos do modelo base podem decodificar rótulos de estratégia, confirmando que a estrutura alvo existe implicitamente antes da fatoração.
2. Falha do ELBO Padrão: Objetivos variacionais padrão preservam a fidelidade da distribuição, mas resultam em colapso latente, onde $z$ permanece inerte e falha em se alinhar com as estratégias.
3. Sucesso da Reconstrução Dirigida pelo Modelo: O objetivo proposto recupera códigos latentes que:
   • Mantêm alta fidelidade de distribuição (preservando as saídas válidas do modelo base).
   • Alcançam alta Consistência Analógica, demonstrando que $z$ carrega semântica de estratégia reutilizável entre diferentes entradas.
   • Separam-se em regiões distintas de tarefa e estratégia no espaço latente ao longo do treinamento.
4. Estudos de Ablação:
   • Remover a escala global (normalização) ou a ponderação de token (direção da surpresa) degrada o alinhamento da estratégia.
   • Usar ponderação de surpresa inversa (enfatizando tokens de baixa surpresa) enfraquece o alinhamento, confirmando que a pressão deve ser aplicada a tokens incertos/de ramificação para recuperar a variação da estratégia.

Análise Teórica

O artigo inicia uma análise teórica do problema de recuperação. Ele enquadra o objetivo como a maximização de um score perfilado envolvendo a informação mútua $I(Y; Z|X)$ menos uma penalidade de complexidade de aproximação. Os autores fornecem condições suficientes sob as quais um objetivo variacional penalizado por complexidade seleciona uma fatoração latente alinhada com uma estrutura de estratégia alvo $S^*$, desde que a "discrepância de recuperação" entre a solução ótima e as subótimas seja positiva.

Significância e Alegações

O artigo afirma abordar uma forma patológica específica de colapso de posterior, única para a adaptação de modelos pré-treinados em vez de treiná-los do zero. Ao introduzir um objetivo de reconstrução dirigido pelo modelo, os autores demonstram um método para:

• Desemaranhar estratégias de raciocínio implícitas em uma variável latente explícita e controlável.
• Preservar a distribuição de resposta do modelo base enquanto tornam a variável latente informativa.
• Habilitar Intervenção: A estrutura latente recuperada permite direcionar a geração para estratégias de alto nível específicas (ex: forçar um método de prova específico) sem o retreinamento do modelo base.

Os autores observam que a metodologia é validada em tarefas algorítmicas sintéticas com estratégias de verdade fundamental. Eles declaram explicitamente que estender isso para tarefas de raciocínio de código aberto e caracterizar totalmente a conexão entre o critério de seleção de oráculo teórico e a dinâmica real de treinamento permanecem questões abertas para trabalhos futuros. As aplicações potenciais mencionadas incluem interpretabilidade de modelos, controle e exploração dirigida em aprendizado por reforço pós-treinamento.