Adaptive Decoding via Test-Time Policy Learning for Self-Improving Generation

1. Problema

As estratégias de decodificação (como amostragem gulosa, top-k ou top-p fixo) determinam a qualidade da saída dos Grandes Modelos de Linguagem (LLMs). No entanto, essas heurísticas são estáticas e agnósticas à tarefa, o que leva a:

• Qualidade de geração subótima ou inconsistente em domínios que exigem flexibilidade estilística ou estrutural.
• Dificuldade em equilibrar trade-offs complexos (fluência, diversidade, precisão factual e controle estilístico).
• A necessidade de retreinar ou fazer fine-tuning do modelo inteiro para melhorar o comportamento de decodificação, o que é computacionalmente caro e impraticável em escala.

O artigo propõe que a decodificação deve ser vista como um ponto de intervenção para auto-aperfeiçoamento recursivo, permitindo que o modelo monitore sua própria geração e ajuste suas ações sem modificar seus parâmetros internos.

2. Metodologia

Os autores propõem um amostrador de decodificação baseado em Aprendizado por Reforço (RL) que trata a decodificação como um problema de tomada de decisão sequencial.

• Formulação como MDP (Processo de Decisão de Markov):
  • Estado ($s_t$): Inclui o prompt, o prefixo gerado, os logits do modelo, representações ocultas (hidden states) e a entropia da distribuição do próximo token.
  • Ação ($a_t$): O agente (uma política leve) seleciona dinamicamente os hiperparâmetros de decodificação, especificamente temperatura e top-p, em tempo real. A política é uma Rede Neural Leve (MLP de 2 camadas) que mapeia o estado para uma distribuição Gaussiana sobre os parâmetros.
  • Recompensa ($R$): Sinais derivados de métricas específicas da tarefa (resumo) que guiam o agente.
• Algoritmo de Treinamento: Utiliza PPO (Proximal Policy Optimization) para treinar a política de forma estável e eficiente em termos de amostragem.
• Condição de Congelamento: Os pesos do LLM base (Granite-3.3-2B e Qwen-2.5-0.5B) permanecem congelados durante todo o processo. Apenas a política leve de decodificação é treinada.
• Função de Recompensa Proposta: Uma combinação ponderada e normalizada de:
  • ROUGE-L F1 (sobreposição com referência).
  • Penalidade de comprimento (desvio do comprimento ideal).
  • Bônus de cobertura (sobreposição com tokens importantes da fonte).
  • Penalidade de repetição (se >30% das palavras se repetem).
  • Penalidade de completude (falta de pontuação final).

3. Principais Contribuições

1. Decodificação Adaptativa sem Retreinamento: Introdução de um amostrador baseado em RL que formaliza a decodificação como um processo de decisão em tempo de teste, permitindo adaptação sem alterar os pesos do LLM.
2. Validação Empírica: Avaliação inicial em dois modelos de tamanhos diferentes (2B e 0.5B) e três conjuntos de dados distintos (BookSum, arXiv, WikiHow), demonstrando ganhos consistentes sobre estratégias estáticas.
3. Análise de Recompensas: Estudos de ablação que revelam como diferentes componentes da função de recompensa influenciam o comportamento de aprendizado e a estabilidade da política.

4. Resultados

Os experimentos foram conduzidos nos conjuntos de dados BookSum (narrativa), arXiv (científico) e WikiHow (instrucional).

• Desempenho Geral: A política aprendida por RL superou consistentemente as bases de comparação (decodificação gulosa e amostragem estática com temperatura fixa).
  • Ganhos Relativos: Até +88% no dataset BookSum (modelo Granite) e +79% no WikiHow (modelo Qwen).
  • Robustez: Em tarefas mais estruturadas como arXiv, o RL manteve desempenho superior ou equivalente às bases estáticas, sem degradação.
• Aprendizado Efetivo: A curva de recompensa ao longo do tempo (Figura 2) mostrou uma tendência de alta, confirmando que o agente aprende uma política não trivial e não colapsa para uma heurística estática.
• Ablação de Recompensas:
  • Recompensas baseadas apenas em ROUGE (sobreposição) foram insuficientes, resultando em ganhos negligenciáveis ou negativos.
  • Termos de "shaping" estruturado (comprimento, repetição, completude) foram críticos para a estabilidade e melhoria sustentada.
  • A remoção do termo de "cobertura" resultou em ganhos grandes e estáveis, sugerindo que essa métrica específica pode introduzir ruído em certos contextos.
• Eficiência Computacional: O overhead de computação é insignificante, pois a política é um MLP pequeno e os recursos de estado (logits, hidden states) já são gerados durante a decodificação padrão.

5. Significado e Implicações

• Auto-Aperfeiçoamento Prático: O trabalho demonstra que o RL pode ser uma ferramenta prática para a adaptação em tempo de teste, permitindo que sistemas de IA melhorem a qualidade da geração dentro de um único episódio de inferência, alinhando-se com visões de sistemas auto-aperfeiçoadores baseados em feedback e controle de política.
• Desacoplamento de Controle e Modelo: Ao separar o controle de amostragem do treinamento do modelo principal, o método evita o custo de retreinamento e permite uma reversibilidade fácil (basta desativar a política aprendida).
• Aplicabilidade em Modelos Menores: A técnica é eficaz mesmo em modelos pequenos (0.5B), tornando-a viável para cenários onde o fine-tuning é impossível, mas a adaptação em tempo de execução é desejável.
• Importância do Design de Recompensa: O estudo enfatiza que a definição da função de recompensa é tão crucial quanto o algoritmo de RL em si; recompensas mal projetadas podem levar a instabilidade ou falha em capturar a qualidade desejada.

Em resumo, o artigo propõe uma mudança de paradigma: em vez de treinar modelos maiores para serem melhores em tudo, podemos usar agentes leves e adaptativos para controlar como esses modelos geram texto, otimizando a saída para tarefas específicas sem alterar o núcleo do modelo.