TAB-PO: Preference Optimization with a Token-Level Adaptive Barrier for Token-Critical Structured Generation

Resumo Técnico: TAB-PO

1. O Problema

O artigo aborda os desafios de alinhar Grandes Modelos de Linguagem (LLMs) para tarefas de geração estruturada crítica, especificamente na anotação de comunicações médicas entre pacientes e provedores (tarefa chamada PV-Miner).

Embora a Otimização Direta de Preferência (DPO) tenha obtido sucesso em tarefas de geração livre (como diálogo e resumo), ela falha em cenários estruturados devido a duas características principais:

1. Baixa Separação de Preferências: Os pares de "escolhido" (ganho) e "rejeitado" (perdido) diferem apenas por um número muito pequeno de tokens (1 a 3 tokens), geralmente nas etiquetas de código ou limites de trechos (Spans). O DPO padrão, que opera no nível da sequência inteira, dilui o sinal de aprendizado entre os milhares de tokens estruturais idênticos (como a sintaxe JSON), falhando em corrigir esses erros finos.
2. Viés de Importância de Tokens: Em anotações estruturadas, tokens semânticos raros (rótulos hierárquicos, evidências) carregam a maior parte da importância da tarefa, enquanto tokens estruturais frequentes (scaffolding JSON) são menos críticos. O DPO padrão trata todos os tokens igualmente, levando a:
   • Colapso de Margem: Insuficiente separação de probabilidade entre preferências quase idênticas.
   • Esmagamento de Probabilidade (Likelihood Squeezing): O modelo aumenta a margem relativa, mas diminui a probabilidade absoluta da resposta correta, comprometendo a validade do esquema.
   • Diluição de Gradiente: O sinal de aprendizado é espalhado uniformemente, enfraquecendo a atualização dos tokens críticos.

2. Metodologia: TAB-PO

Os autores propõem o TAB-PO (Token-Adaptive Barrier Preference Optimization), um método de pós-ajuste (post-SFT) que modifica o DPO para focar em tokens críticos e estabilizar a geração.

Componentes Principais:

• Construção de Preferências Conscientes de Confusão:

  • Em vez de gerar negativos aleatórios, o TAB-PO cria pares de preferência baseados em erros reais do modelo SFT (Supervised Fine-Tuning).
  • Utiliza uma mistura de preferências curadas por especialistas (40%) e sintéticas baseadas em confusões empíricas do modelo (60%).
  • Perturbações focadas: Confusão de Código/Sub-código, anotações faltantes (deleção) e anotações extras (inserção espúria), mantendo a estrutura JSON intacta.

• Termo DPO Ponderado por Token (Token-Weighted DPO):

  • Substitui a vantagem de log-verossimilhança no nível da sequência por uma vantagem ponderada por token.
  • Atribui pesos maiores ($w_{Code}, w_{Sub-code}, w_{Span}$) aos tokens semânticos críticos e peso zero aos tokens estruturais (JSON).
  • Isso concentra o sinal de gradiente exatamente onde as preferências diferem, evitando a diluição.

• Barreira Adaptativa de Nível de Token (Adaptive Token-Level Barrier):

  • Introduz um mecanismo de regularização condicional para prevenir o "esmagamento de probabilidade".
  • Atua como um "âncora" baseada no modelo SFT. Se o modelo atual estiver pouco confiante em um token preferido (probabilidade abaixo de um limiar $\tau$), aplica-se uma penalidade estilo SFT para aumentar a probabilidade desse token específico.
  • Isso garante que a geração estruturada (JSON) permaneça válida e estável, enquanto permite a separação de preferências nos tokens semânticos.

• Normalização e Balanceamento:

  • Normalização por Comprimento: Para evitar que exemplos longos dominem o gradiente do lote.
  • Balanceamento de Classe: Pondera exemplos baseados na raridade dos códigos presentes, mitigando o viés de classes frequentes.

3. Contribuições Principais

1. Tarefa e Benchmark PV-Miner: Introdução de um novo benchmark para anotação estruturada de mensagens seguras de pacientes, envolvendo previsão conjunta de rótulos hierárquicos (Código/Sub-código) e trechos de evidência (Spans) com restrições estritas de esquema.
2. Método TAB-PO: Uma nova abordagem de otimização de preferência que:
   • Emphasiza tokens semânticos críticos via ponderação de valor.
   • Estabiliza a geração estruturada através de uma barreira adaptativa ancorada no SFT.
3. Superioridade em Cenários de Baixa Separação: Demonstra que a otimização de preferência tradicional falha em tarefas onde a diferença entre respostas boas e ruins é mínima, enquanto o TAB-PO resolve essas ambiguidades finas.

4. Resultados Experimentais

Os experimentos foram conduzidos em modelos de diferentes tamanhos (de 1.5B a 70B parâmetros, incluindo Llama-3 e Qwen2.5) na tarefa PV-Miner.

• Desempenho Geral: O TAB-PO superou consistentemente o SFT e todas as variantes recentes de DPO (DPO padrão, IPO, SimPO, Cal-DPO, DPO-Positive).
• Métricas:
  • Houve uma melhoria relativa média de ~4% no micro-F1 em comparação ao SFT.
  • Melhorias significativas foram observadas na previsão de Sub-códigos (rótulos finos), onde o TAB-PO aumentou o F1 em até 8.3 pontos percentuais (ex: Qwen2.5-1.5B).
  • Redução da variância entre execuções, indicando otimização mais estável.
• Comparação com Baselines: O TAB-PO superou todas as técnicas de otimização de preferência em F1 para Código, Sub-código e Span, demonstrando melhor resolução de rótulos hierárquicos confusos e preservação da estrutura JSON.

5. Significância

Este trabalho é significativo porque:

• Preenche uma Lacuna Crítica: Mostra que métodos de alinhamento padrão (DPO) não são suficientes para tarefas de alto risco e alta precisão (como saúde), onde erros de um único token podem invalidar todo o output.
• Inovação Técnica: Propõe uma mudança de paradigma de "otimização de sequência" para "otimização sensível a tokens", reconhecendo que nem todos os tokens em uma geração estruturada têm o mesmo valor para a tarefa.
• Aplicabilidade Prática: Oferece uma metodologia robusta para treinar modelos em cenários de anotação médica e extração de informações estruturadas, garantindo que os modelos sejam não apenas precisos semanticamente, mas também estruturalmente válidos.

Em resumo, o TAB-PO oferece uma solução eficaz para o problema de "diluição de gradiente" e "instabilidade de probabilidade" em tarefas de geração estruturada, permitindo que LLMs aprendam a distinguir nuances críticas em dados médicos complexos.