Token-Level Off-Policy Learning for Faithful Generation Under Distribution Shift

Resumo Técnico: Aprendizado Off-Policy ao Nível de Token para Geração Fiel

Definição do Problema
Grandes Modelos de Linguagem (LLMs) frequentemente lutam contra alucinações e inconsistências fatuais, particularmente quando geram conteúdo sob mudanças de distribuição (cenários Out-of-Distribution ou OOD). Embora o Aprendizado por Reforço com Feedback Humano (RLHF) e métodos off-policy como o Direct Preference Optimization (DPO) tenham melhorado o alinhamento, eles frequentemente dependem de supervisão ao nível de sequência. Essa granularidade grosseira limita a capacidade de identificar erros fatuais específicos. Além disso, métodos on-policy (ex: GRPO) sofrem com alta complexidade computacional devido à amostragem online, enquanto abordas off-policy existentes ao nível de token frequentemente tratam a modelagem de recompensa como um componente separado, em vez de um objetivo de treinamento unificado. O artigo identifica a necessidade de um método que mantenha a simplicidade do treinamento off-policy, ao mesmo tempo que fornece supervisão fina ao nível de token para melhorar a factualidade e a generalização OOD.

Metodologia: Rotulagem Off-Policy ao Nível de Token (TOPL)
Os autores propõem o Token-Level Off-Policy Labeling (TOPL), um paradigma de treinamento que reformula o pós-treinamento como uma tarefa de classificação binária ao nível de token, em vez de uma tarefa de previsão do próximo token.

• Construção de Dados: O método utiliza um conjunto de dados (baseado em FAVA) contendo documentos, resumos de verdade fundamental (ground-truth) e resumos perturbados ("corrompidos"). As corrupções são introduzidas via inserções, deleções ou substituições de tokens para simular alucinações. Cada token na resposta recebe um rótulo binário: $1$ para fiel (correto) e $0$ para não fiel (corrompido).
• Objetivo de Treinamento: Em vez de otimizar um cabeçalho de modelagem de linguagem para a previsão do próximo token, o TOPL anexa um cabeçalho de recompensa leve a uma camada intermediária do modelo pré-treinado. O modelo é treinado para prever a probabilidade de um token específico ser factualmente consistente com o contexto de entrada.
• Arquitetura e LoRA: O treinamento é realizado utilizando Low-Rank Adaptation (LoRA). O modelo aprende a distinguir tokens bons e ruins via uma perda de entropia cruzada binária. Crucialmente, o cabeçalho de recompensa é descartado após o treinamento; os pesos LoRA aprendidos são fundidos (merged) ao modelo base para criar um modelo generativo que utiliza o cabeçalho de decodificação original.
• Interpretação como Direcionamento Condicional: Os autores interpretam o TOPL através da lente do direcionamento condicional (conditional steering). Eles hipotetizam que a matriz LoRA-A atua como um mecanismo de projeção condicional (classificando a correção do token), enquanto a matriz LoRA-B define vetores de direcionamento que deslocam os estados ocultos para longe ou para perto de um comportamento fiel com base no sinal de classificação.

Principais Contribuições

1. Novo Paradigma de Treinamento: Introdução do TOPL, que unifica o aprendizado de recompensa ao nível de token e a otimização de política em um único framework off-policy, evitando a complexidade da amostragem online.
2. Insight Mecanístico: Uma análise demonstrando que o TOPL induz atualizações interpretáveis no modelo, onde os componentes LoRA funcionam como cabeçalhos de classificação linear (LoRA-A) e vetores de direcionamento (LoRA-B).
3. Validação Empírica: Avaliação abrangente mostrando que o TOPL supera tanto os baselines de nível de sequência (SFT, DPO, SOPL) quanto os baselines de nível de token existentes (TDPO, TLDR, Unlikelihood Training) em cenários OOD.

Resultos Experimentais

• Sumarização de Documentos: Avaliado em 11 datasets OOD do AggreFact usando Qwen3-8B, Llama-3.1-8B e Gemma-3-4B. O TOPL alcançou as maiores pontuações médias de factualidade em todos os modelos de base.
  • Ablação de Granularidade: Análogos de nível de sequência (SOPL) falharam em igualar o desempenho do TOPL, confirmando que sinais finos ao nível de token são críticos.
  • Seleção de Camada: Aplicar LoRA em camadas intermediárias yielded os melhores resultados; atualizar as camadas finais proporcionou ganhos limitados ou degradou o desempenho, provavelmente devido à interrupção dos comportamentos generativos pré-treinados.
  • Sensibilidade de Rótulo: O método depende de supervisão contrastiva (presença de ambos os tokens bons e ruins). Embora a proporção de tokens bons/ruins não tenha sido o fator dominante, o treinamento com apenas um tipo de token levou à degradação do desempenho.
• Tradução Automática: O TOPL foi transferido com sucesso para tarefas de tradução automática (MLQE-PE e FLORES-Plus), alcançando desempenho estado da arte em OOD, sugerindo que os benefícios se generalizam para tarefas de geração fiel.
• Análise de Representação:
  • Separabilidade: Modelos treinados com TOPL exibiram significativamente maior separabilidade (AUROC) entre tokens factuais e não factuais no subespaço LoRA-A em comparação ao SFT.
  • Direcionamento (Steering): Intervenções post-hoc usando vetores LoRA-B mostraram que os modelos TOPL respondem coerentemente à força de direcionamento, enquanto os modelos SFT mostraram baixa sensibilidade.
  • Correlação: Uma maior separabilidade de LoRA-A correlacionou-se fortemente com melhores pontuações de generalização OOD.

Significância e Alegações
O artigo afirma que o TOPL oferece uma solução robusta para melhorar a factualidade do modelo sob mudanças de distribuição, mudando o objetivo de treinamento de "qual token vem a seguir" para "este token é correto". Os autores argumentam que esta abordagem evita as armadilhas de treinar diretamente modelos para gerar tokens off-policy. Em vez disso, ao treinar o modelo para distinguir a correção, o método naturalmente guia o processo de geração em direção a saídas fiéis.

A significância do trabalho reside na demonstração de que:

1. A supervisão ao nível de token é superior à supervisão ao nível de sequência para generalização OOD em geração fiel.
2. O objetivo específico importa: Ter apenas rótulos ao nível de token é insuficiente; o objetivo de classificação binária usado no TOPL cria uma estrutura de representação interna específica (LoRA-A separável e LoRA-B direcionável) que aumenta a robustez.
3. Interpretabilidade: O método fornece uma explicação mecanística clara para seu sucesso, ligando as atualizações LoRA aprendidas a mecanismos de direcionamento condicional que podem ser analisados e potencialmente controlados.

Os autores permanecem modestos quanto à causalidade, observando que, embora observem uma forte correlação entre a separabilidade de LoRA-A e o desempenho OOD, a relação causal estrita entre a estrutura de representação e a generalização permanece uma questão aberta para trabalhos futuros.