OLLM: Options-based Large Language Models

1. O Problema

Os Modelos de Linguagem (LLMs) tradicionais operam sob a premissa de prever uma única distribuição para o próximo token, dependendo de heurísticas de amostragem (como temperatura, top-k ou top-p) para introduzir diversidade. O artigo identifica duas limitações principais nessa abordagem:

• Ineficiência na Exploração: A busca por respostas corretas ocorre sobre todo o vocabulário, o que é computacionalmente caro e ineficiente para otimização por reforço (RL).
• Ambiguidade vs. Determinismo: A análise empírica mostra que, embora muitas posições na geração de texto sejam altamente ambíguas (com múltiplas continuações plausíveis), outras são quase determinísticas (ex: sequências matemáticas fixas). Os modelos atuais tratam tudo como uma única distribuição, fazendo com que tokens alternativos competam pela mesma massa de probabilidade dentro de um único softmax, o que pode levar a overfitting em um modo dominante ou a escolhas subótimas.
• Desalinhamento em RL: Técnicas de ajuste fino por reforço (RLHF) que operam diretamente no espaço de vocabulário completo podem levar a comportamentos indesejados, como troca de idioma ou raciocínio degenerado, devido à vasta e complexa espaço de ação.

2. Metodologia: Options LLM (OLLM)

Os autores propõem o OLLM, um método que substitui a previsão de um único próximo token por um conjunto de "opções" aprendidas, indexadas por uma variável latente discreta.

• Arquitetura (Plug-in Leve): O OLLM insere dois módulos leves (um Encoder e um Decoder) antes da camada de saída (lm_head) de qualquer LLM pré-treinado.
  • Encoder ($E_\psi$): Durante o treinamento supervisionado (SFT), recebe o estado oculto do contexto ($h_t$) e a representação do token alvo ($e(y_{t+1})$). Ele mapeia essa informação para uma distribuição sobre um espaço latente discreto e pequeno $Z$ (ex: $K=10$ opções). Isso permite que o modelo aprenda contextos sensíveis onde o mesmo token alvo pode pertencer a diferentes "opções" dependendo do trajeto de raciocínio.
  • Decoder ($D_\omega$): Recebe o estado oculto e a opção latente amostrada ($z$) para gerar um vetor de viés ($\Delta h_t$) que é adicionado ao estado oculto antes da projeção no vocabulário.
• Decomposição da Distribuição: A distribuição do próximo token é decomposta em uma mistura sobre as opções latentes. Em vez de tokens concorrentes lutarem por probabilidade em um único espaço, diferentes continuação plausíveis são otimizadas sob latentes diferentes, desacoplando-as.
• Aprendizado de Política no Espaço Latente:
  • Na inferência, o Encoder é substituído por uma Política Leve ($\pi_\phi$) que prevê a opção latente $z$ baseada apenas no estado oculto $h_t$.
  • A política é treinada para imitar o Encoder (usando Behavioral Cloning via SFT).
  • Isso reduz o espaço de ação para a otimização de RL de todo o vocabulário para o pequeno conjunto de opções $Z$, tornando a busca por recompensas muito mais eficiente e segura.

3. Contribuições Chave

1. Modelagem Explícita de Opções: Introduz uma representação latente discreta que captura explicitamente a multimodalidade nas previsões de tokens, permitindo que o modelo mantenha múltiplos caminhos de raciocínio plausíveis sem overfitting.
2. Eficiência Paramétrica: O método é um "plug-in" que requer o ajuste de apenas uma fração mínima dos parâmetros (no experimento, apenas 1,56% de um modelo de 1,7B parâmetros), tornando-o aplicável a modelos grandes existentes.
3. Alinhamento Estrutural: O alinhamento e a robustez são alcançados através da estrutura do modelo (restrição ao espaço de opções aprendidas no SFT) em vez de depender de penalidades KL adicionais ou funções de perda de alinhamento manuais.
4. Eficiência em RL: Ao restringir a exploração a um espaço de ação de baixa dimensão (as opções), o método mitiga problemas comuns de RLHF, como desalinhamento e raciocínio degenerado.

4. Resultados

O modelo foi aplicado em um backbone de 1,7B parâmetros (derivado de Qwen2.5/Deepseek R1), treinado no conjunto de dados OpenMathReasoning e avaliado no OmniMath.

• Desempenho:
  • Modelos de base com adaptadores LoRA de última geração atingiram um pico de 51% de precisão na resposta final.
  • O OLLM, sob seleção ótima de latentes (teórica), atingiu ~70% de precisão na resposta final e ~83% de precisão por token.
• Análise: O ganho de desempenho ocorre porque as alternativas de tokens não competem dentro de um único simplex de probabilidade. A separação em opções latentes permite preservar múltiplas continuações plausíveis.
• Comportamento: O modelo demonstra capacidade de aprender opções genuinamente diversas quando há ambiguidade (ex: início de frases ou raciocínio lógico), mas colapsa para uma escolha consistente em casos determinísticos (ex: sequências de símbolos matemáticos fixos).

5. Significado e Futuro

O trabalho demonstra que a modelagem de "próximo token" baseada em opções melhora significativamente a controlabilidade, robustez e eficiência no raciocínio matemático.

• Implicações para RL: A abordagem sugere que o aprendizado de políticas em espaços latentes compactos é uma direção promissora para Reinforcement Learning em LLMs, oferecendo um espaço de ação bem fundamentado e de baixa dimensão.
• Interpretabilidade: A estrutura de opções oferece interpretabilidade, pois diferentes latentes podem corresponder a trajetórias de raciocínio distintas, permitindo a exploração controlada de caminhos de solução alternativos.
• Aplicações Futuras: Os autores planejam treinar políticas neste espaço latente usando RL real (além do SFT) e estender o método para domínios como geração de código, QA científica e planejamento multi-etapa.

Em resumo, o OLLM oferece uma mudança de paradigma: em vez de forçar o modelo a escolher uma única distribuição complexa sobre o vocabulário, ele aprende a estruturar suas escolhas em um conjunto gerenciável de "opções", facilitando tanto a geração precisa quanto a otimização futura.