Resumo Técnico: Além do Escalonamento de Temperatura Post-Hoc: Otimização Bilevel para Calibração de LLMs

Declaração do Problema

Grandes Modelos de Linguagem (LLMs), particularmente aqueles alinhados via Aprendizado por Reforço com Feedback Humano (RLHF) ou Otimização de Preferência Direta (DPO), frequentemente exibem uma calibração deficiente, manifestando-se como excesso de confiança (overconfidence). Embora a confiança seja frequentemente usada como um proxy para confiabilidade, modelos mal calibrados em domínios de alto risco (ex: saúde, jurídico) podem levar a uma confiança indevida e a decisões ruins.

As soluções existentes baseiam-se primariamente no Escalonamento de Temperatura (TS) post-hoc, onde um escalar de temperatura único é ajustado em um conjunto de validação para reescalonar os logits após o treinamento. No entanto, o artigo identifica uma limitação crítica: o TS é inerentemente dependente do conjunto de dados. Uma temperatura otimizada para um domínio frequentemente falha ao generalizar para outros, tornando difícil encontrar uma temperatura "universal" que calibre bem em diversas tarefas. Além disso, métodos post-hoc deixam os parâmetros do modelo fixos, perdendo a oportunidade de aprender um modelo que seja intrinsecamente mais propenso à calibração.

Metodologia: CALM (Calibração para Grandes Modelos via Otimização Bilevel)

Os autores propõem o CALM, um framework de tempo de treinamento que reformula a calibração como um problema de otimização bilevel. Em vez de ajustar um parâmetro post-hoc, o CALM otimiza conjuntamente os parâmetros do modelo e hiperparâmetros de perda "semelhantes à temperatura" para maximizar a entropia preditiva enquanto preserva o desempenho da tarefa.

1. Formulação Bilevel

O problema é estruturado da seguinte forma:

• Nível Inferior (Otimização da Tarefa): Treina os parâmetros do modelo $\theta$ para minimizar uma perda paramétrica padrão (entropia cruzada) no conjunto de treinamento $S_T$, dado um conjunto de hiperparâmetros $\alpha$.
  $$\theta(\alpha) \in \arg\min_{\theta'} \mathcal{L}_{tr}(\theta'; \alpha)$$
• Nível Superior (Otimização da Calibração): Otimiza os hiperparâmetros $\alpha$ para maximizar a entropia da distribuição preditiva em um conjunto de calibração mantido fora (held-out) $S_V$, sujeito à restrição de que $\theta$ permaneça ótimo para a tarefa de nível inferior.
  $$\min_{\alpha} \mathcal{L}_{ent}(\alpha, \theta(\alpha))$$

Os hiperparâmetros $\alpha$ generalizam o conceito de escalonamento de temperatura. Em vez de um escalar único $\tau$, o CALM utiliza um parâmetro de maior dimensão $\alpha = (l, \tau)$, onde $l \in \mathbb{R}^K$ introduz deslocamentos de logit aditivos e $\tau \in \mathbb{R}^K_{>0}$ fornece temperaturas dependentes da classe. Isso permite ajustes por vocabulário em vez de um reescalonamento global.

2. Funções Objetivo

• Objetivo de Nível Superior: Maximiza a entropia de Shannon $H(f_{\alpha, \theta}(\cdot | x))$ para desencorajar predições excessivamente concentradas (overconfident). Isso fornece uma direção controlável para mover o modelo do excesso de confiança em direção ao sub-confiança, potencialmente encontrando um ponto onde a confiança corresponde à acurácia.
• Restrição de Nível Inferior: Garante que o modelo retenha a capacidade discriminativa ao minimizar a entropia cruzada nos dados de treinamento.

3. Aproximação de Primeira Ordem Eficiente

A otimização bilevel exata requer o cálculo de hipergradientes envolvendo derivadas de segunda ordem (inversas de Hessian), o que é computacionalmente proibitivo para LLMs com bilhões de parâmetros. Para abordar isso, o CALM adota uma aproximação de primeira ordem inspirada no BOME (Bilevel Optimization Made Easy):

• Ele reformula o problema como uma otimização restrita sobre a variável conjunta $u = (\alpha, \theta)$.
• Ele aproxima a função de valor de nível inferior executando $T$ passos de gradiente interno a partir do atual $\theta$ para gerar um $\tilde{\theta}$ substituto.
• Ele constrói uma restrição de "resíduo de plug-in" e calcula uma direção de atualização conjunta usando apenas derivadas de primeira ordem, evitando a inversão explícita de Hessian.

Principais Contribuições

1. Perspectiva de Tempo de Treinamento: O artigo desloca o paradigma de calibração de ajustes post-hoc específicos de conjuntos de dados para a otimização em tempo de treinamento, visando aprender modelos que admitam ajustes de calibração generalizáveis entre conjuntos de dados.
2. Design de Perda Bilevel: Formula o design de perda orientado à calibração como um problema bilevel onde o nível superior otimiza hiperparâmetros semelhantes à temperatura usando um objetivo de entropia, e o nível inferior treina o modelo.
3. Maximização de Entropia: Os autores propõem um objetivo de maximização de entropia no nível superior para mitigar o excesso de confiança em LLMs alinhados sem exigir rótulos rígidos adicionais ou cabeças auxiliares.
4. Framework Escalável: Eles introduzem um esquema de otimização bilevel de primeira ordem eficiente (CALM) que evita computações dispendiosas de segunda ordem, tornando sua abordagem prática para o ajuste fino de LLMs em larga escala.

Resultados Experimentais

O método foi avaliado em quatro LLMs alinhados (Llama-3.1-Tulu-8B, Vicuna-7B, OLMo-2-7B, Mistral-7B) em questões de múltipla escolha (MCQA) e QA generativo aberto.

Generalização Fora do Domínio (OOD)

• Calibração Superior: O CALM alcançou o melhor erro de calibração esperado de confiança (conf-ECE) em três dos quatro modelos no cenário OOD, reduzindo significativamente o erro em comparação com os baselines. Por exemplo, no Llama-3.1, reduziu o conf-ECE de 0.1784 para 0.1050.
• Preservação de Acurácia: Ao contrário de algumas ablações (ex: "Iterate" ou "Regularization" sem a restrição bilevel) que sofreram colapso catastrófico de acurácia, o CALM manteve alta acurácia de tarefa enquanto melhorava a calibração.
• Robustez: Em QA generativo aberto (transferência de domínio entre PopQA e TriviaQA), o CALM alcançou o menor Semantic-ECE médio (0.0671), superando tanto o baseline não calibrado quanto fortes baselines de tempo de treinamento como o Calibration-Aware Fine-Tuning (CFT).

Desempenho Dentro do Domínio (ID)

• No cenário ID, onde as distribuições de treinamento e avaliação coincidem, o Escalonamento de Temperatura teve o melhor desempenho em ECE (como esperado para métodos post-hoc). No entanto, o CALM foi o único método baseado em treinamento que consistentemente melhorou tanto a acurácia quanto a calibração em todos os quatro modelos comparado ao baseline.

Estudos de Ablação

• Estrutura Bilevel: Remover a restrição bilevel (Regularization) ou o termo de acoplamento (Iterate) resultou em pior calibração ou perda significativa de acurácia, demonstrando que o acoplamento entre desempenho da tarefa e calibração é essencial.
• Parametrização: Usar parâmetros vetoriais por vocabulário ($d_y, l_y$) foi muito superior ao uso de um parâmetro escalar colapsado (TS aprendido), destacando a necessidade de ajustes finos para lidar com a má calibração assimétrica entre as opções de resposta.

Significância e Alegações

O artigo afirma que o CALM fornece uma solução robusta de tempo de treinamento para o problema do excesso de confiança em LLMs alinhados. Sua principal significância reside em:

• Generalização: Diferente do TS post-hoc, o CALM aprende um mecanismo de calibração que generaliza entre domínios, abordando a "dependência de conjunto de dados" dos métodos atuais.
• Gestão de Trade-off: Ele alcança um trade-off superior entre acurácia e calibração, evitando o colapso de utilidade frequentemente observado ao tentar melhorar a calibração via regularização simples.
• Praticidade: Ao utilizar uma aproximação de primeira ordem, torna a otimização bilevel viável para LLMs modernos, oferecendo um caminho para modelos "calibrados por design" em vez de "calibrados por pós-processamento".

Os autores concluem que seu método é particularmente vantajoso em cenários fora do domínio, onde métodos alternativos frequentemente se tornam instáveis ou sofrem perda substancial de utilidade, tornando-o uma abordagem viável para implantar LLMs confiáveis em diversos cenários do mundo real.