Resumo Técnico: Rumo à Internalização Robusta na Destilação On-Policy

1. Definição do Problema

Avanços recentes na destilação on-policy (OPD) permitiram que modelos estudantes internalizem contextos privilegiados (ex: prompts de sistema, dicas de tarefas ou descrições de estado de jogo) em seus parâmetros, permitindo que desempenhem bem sem o contexto no momento da inferência. No entanto, este artigo identifica um fenômeno anteriormente não estudado, denominado degradação induzida pelo contexto.

Ao contrário da intuição de que o conhecimento internalizado deveria tornar o contexto redundante (uma operação idempotente), os autores observam que reintroduzir o contexto privilegiado a um estudante destilado frequentemente degrada seu desempenho e infla o comprimento das respostas, mesmo em instâncias que o modelo resolve corretamente sem o contexto. Isso ocorre porque a OPD padrão otimiza exclusivamente a fidelidade privilegiada (corresponder à saída do professor condicionada ao contexto), mas não impõe restrições ao comportamento do estudante quando o contexto está presente. Consequentemente, o estudante falha em alcançar a removibilidade de contexto: a propriedade onde a saída do modelo permanece estável, independentemente de o contexto privilegiado ser fornecido ou não.

Os autores categorizam as interações pós-destilação de contexto em três regimes:

• Regime A: O estudante apresenta bom desempenho sem o contexto, mas a reintrodução do contexto causa quedas significativas de precisão e altas "taxas de dano" (inversão de previsões corretas para incorretas). Este é o regime mais prevalente em seus experimentos.
• Regime B: O contexto ainda atua como um andaime útil, indicando uma internalização incompleta.
• Regime C: O estudante internalizou com sucesso o contexto, e o desempenho é estável com ou sem ele.

2. Metodologia: Ancoragem Sem Contexto (NCA)

Para abordar a degradação induzida pelo contexto, os autores propõem a Ancoragem Sem Contexto (NCA), um regularizador de consistência leve. A ideia central é forçar que a saída do estudante condicionada ao contexto ($q_c$) permaneça consistente com sua saída sem contexto ($q_x$), tratando esta última como uma âncora.

Função Objetivo

A OPD padrão minimiza a divergência KL entre a saída sem contexto do estudante e a distribuição condicionada ao contexto do professor:
$$\mathcal{L}_{OPD} = \mathbb{E}_{x \sim D} [\text{KL}(q_x \parallel f(y \mid x, c))]$$

A NCA aumenta isso com um termo de consistência:
$$\mathcal{L} = \mathcal{L}_{OPD} + \beta \cdot \mathcal{L}_{NCA}$$
onde $\mathcal{L}_{NCA}$ é definida como:
$$\mathcal{L}_{NCA} \triangleq \text{KL}(\text{sg}[g_\theta(\cdot \mid x)] \parallel g_\theta(\cdot \mid x, c))$$

Principais Escolhas de Design

1. Ancoragem com Stop-Gradient: A visão sem contexto $q_x$ passa por um operador de stop-gradient (sg). Isso garante que $q_x$ receba gradientes apenas da perda primária de OPD, enquanto $q_c$ recebe gradientes apenas da perda de NCA. Isso elimina a competição de gradientes entre os dois objetivos.
2. Divergência KL Direta (Forward KL): Os autores escolhem $\text{KL}(q_x \parallel q_c)$ (forward KL) em vez de reverse KL. Esta escolha é "cobertora de modos" (mode-covering), evitando que $q_c$ colapse os modos de $q_x$. Além disso, a esperança recai sobre $q_x$, permitindo o reuso de amostras on-policy geradas para a OPD sem exigir um novo rollout de $q_c$.
3. Eficiência Computacional: A NCA requer apenas uma passagem adicional de forward por etapa de treinamento (para computar $q_c$). Os logits para $q_x$ são desconectados do grafo de computação, e as sequências on-policy são reutilizadas, tornando o overhead negligenciável.

3. Configuração Experimental

Os autores avaliaram a NCA em 12 configurações abrangendo:

• Domínios: QA de prompt de sistema (MedMCQA, datasets de Segurança) e Jogos de Texto (Sokoban, FrozenLake).
• Modelos: Vários modelos Llama-3 e Qwen (variando de 1.7B a 8B parâmetros), incluindo cenários de autodestilação e destilação entre modelos distintos.
• Métricas: Precisão/Taxa de Vitória com e sem contexto, Taxa de Dano (probabilidade de o contexto inverter uma resposta correta), taxas de "Ambos Corretos/Ambos Venceram" e mudanças relativas no comprimento da resposta.

4. Principais Resultados

Em todas as 12 configurações, a NCA demonstrou os seguintes efeitos:

• Mitigação da Degradação Induzida pelo Contexto: A NCA reduziu o dano induzido pelo contexto em 11 de 12 configurações. Em configurações de Regime A (onde a degradação era severa), a NCA reduziu significativamente a queda de precisão após a reintrodução do contexto.
• Melhoria na Estabilidade: A NCA consistentemente aumentou as taxas de "AmbosCorretos" ou "AmbosVenceram", indicando que o modelo mantém o comportamento correto sob ambas as visões.
• Eliminação da Inflação do Comprimento da Resposta: A OPD padrão frequentemente causava que o modelo gerasse respostas significativamente mais longas quando o contexto era reintroduzido (até +39.5% de aumento no comprimento). A NCA eliminou efetivamente essa inflação na maioria das configurações, muitas vezes reduzindo a mudança de comprimento para quase zero.
• Melhoria no Desempenho Sem Contexto: Em 7 de 12 configurações, a NCA melhorou o objetivo primário de treinamento (precisão sem contexto) em 0,6 a 3,6 pontos percentuais. Os autores atribuem isso ao fato de a restrição de consistência atuar como um regularizador que impede que as atualizações de parâmetros acomodem comportamentos conflitantes entre as visões.
• Casos Limítrofes: Em um cenário de Regime B (FrozenLake, cross-model), a NCA aumentou ligeiramente o dano, sugerindo que quando o contexto carrega informações únicas e não internalizadas, forçar uma removibilidade estrita pode trocar parte do desempenho condicionado ao contexto pela estabilidade.

5. Análise Mecanística

O artigo fornece um estudo de caso mecanístico na tarefa Sokoban para explicar como a NCA alcança a removibilidade de contexto:

• Nível de Representação: A análise de estados ocultos revelou que, embora a OPD padrão reduzisse a divergência entre as visões de contexto e sem contexto, uma lacuna perceptível permanecia nas camadas médias e tardias. A NCA, contudo, manteve uma similaridade de cosseno quase perfeita ($\ge 0.997$) em todas as camadas, confirmando que a removibilidade de contexto é alcançada no nível da representação.
• Alinhamento de Gradiente: Na OPD padrão, os gradientes para o objetivo de destilação e a visão condicionada ao contexto eram quase ortogonais. A NCA introduz uma nova direção de gradiente que equilibra ativamente as forças de aproveitar o contexto para destilação e permanecer invariante a ele.
• Atualizações de Parâmetros: Tanto a OPD quanto a NCA concentram atualizações de parâmetros em camadas rasas (particularmente $k\_proj$ e $q\_proj$). No entanto, as atualizações específicas impulsionadas pela NCA são distintas das da OPD, sugerindo que a removibilidade de contexto é alcançada principalmente através de modificações nas representações de conhecimento de camadas rasas que a OPD sozinha não produz.

6. Significância e Alegações

Os autores afirmam que seu trabalho altera o paradigma de avaliação para a destilação de contexto privilegiado. Eles argumentam que a internalização robusta requer não apenas o correspondência ao comportamento do professor (fidelidade privilegiada), mas também garantir a removibilidade de contexto.

• Novidade: Este é o primeiro estudo a identificar e abordar o fenômeno onde a reintrodução de contexto internalizado degrada o desempenho do modelo.
• Praticidade: O método proposto (NCA) é simples, exigindo apenas uma passagem de forward e um regularizador de consistência padrão, mas resolve efetivamente o problema da degradação sem sacrificar os objetivos primários de destilação.
• Generalização: O método melhora o desempenho através de diversas famílias de modelos, tarefas e configurações de destilação (auto e cross-model).

O artigo conclui instando a comunidade a avaliar modelos destilados não apenas pelo seu desempenho sem contexto, mas também pela sua robustez à reintrodução de contexto, garantindo que o conhecimento internalizado não se torne um passivo quando o contexto é acidental ou intencionalmente fornecido.