UpSkill: Mutual Information Skill Learning for Structured Response Diversity in LLMs

O artigo apresenta o UpSkill, um método de treinamento que adapta o aprendizado de habilidades baseado em informação mútua ao otimizador GRPO para aumentar a diversidade de respostas e melhorar as métricas pass@k em modelos de linguagem grandes, sem comprometer a precisão de tentativa única.

O essencial

Imagine que você tem um gênio da lâmpada (uma Inteligência Artificial) que é muito bom em resolver problemas de matemática. Se você pedir a ele uma solução, ele geralmente acerta. Mas, e se você pedir para ele tentar resolver o mesmo problema cinco vezes seguidas?

O problema é que, se você pedir isso para a maioria das IAs hoje, elas vão te dar cinco respostas quase idênticas. É como se você tivesse cinco amigos tentando adivinhar a senha do Wi-Fi, mas todos eles estivessem usando exatamente a mesma lógica e, se um errar, os outros cinco também vão errar da mesma maneira. Isso é desperdício de tempo e esforço.

O artigo que você enviou apresenta uma nova técnica chamada UpSkill. Vamos explicar como ela funciona usando uma analogia simples: O Restaurante de Estratégias.

O Problema: O Restaurante "Sem Variedade"

Atualmente, quando treinamos IAs para acertar respostas (como em provas de matemática), nós as recompensamos apenas quando elas acertam. Com o tempo, a IA aprende que a única coisa que importa é o "prato perfeito". Então, ela para de experimentar sabores novos e começa a fazer o mesmo prato, da mesma maneira, todas as vezes.

• Resultado: Se o prato estiver levemente estragado (uma pequena falha no raciocínio), a IA vai servir esse mesmo prato estragado 10 vezes seguidas. Você não ganha nada tentando de novo.

A Solução: O Chef com "Cartas de Estratégia" (UpSkill)

Os autores do UpSkill tiveram uma ideia brilhante: e se, em vez de pedir apenas "uma resposta", nós pedíssemos à IA que usasse diferentes "modos" ou "estratégias" para pensar?

Imagine que a IA agora tem um menu com 5 cartas mágicas (chamadas de variáveis latentes ou "z").

• Carta 1: "Resolva isso como um matemático rigoroso, passo a passo."
• Carta 2: "Resolva isso desenhando um gráfico mental."
• Carta 3: "Resolva isso tentando chutar e corrigindo."
• Carta 4: "Resolva isso usando uma analogia com culinária."
• Carta 5: "Resolva isso de trás para frente."

O segredo do UpSkill é ensinar a IA a gostar de usar essas cartas diferentes. Eles criaram uma regra de recompensa (chamada de "Recompensa de Informação Mútua") que diz:

"Ei, IA! Se você usar a Carta 1 e der uma resposta, e depois usar a Carta 2 e der uma resposta totalmente diferente, você ganha pontos extras! Mas se você usar a Carta 1 e a Carta 2 e der a mesma resposta, você não ganha nada."

Como isso funciona na prática?

1. Treinamento: A IA aprende que para ganhar pontos, ela precisa ser criativa e variar sua abordagem. Ela descobre que a Carta 1 funciona bem para um tipo de problema, e a Carta 3 funciona melhor para outro.
2. O Teste: Quando você quer resolver um problema difícil, você pede para a IA tentar 5 vezes, cada uma com uma carta diferente.
3. O Resultado: Em vez de 5 respostas iguais, você recebe 5 caminhos de raciocínio diferentes.
   • Se a IA errar na abordagem "Matemática Rigorosa", ela pode acertar na abordagem "Desenho Mental".
   • Isso aumenta drasticamente a chance de que pelo menos uma das 5 tentativas seja correta.

A Analogia da Caça ao Tesouro

Pense em procurar um tesouro enterrado em uma floresta:

• IA Antiga: Você envia 5 exploradores, mas todos seguem o mesmo mapa. Se o mapa estiver errado, todos se perdem juntos.
• IA com UpSkill: Você envia 5 exploradores, mas dá a cada um um mapa diferente (um vai pelo norte, outro pelo sul, outro pelo rio, etc.). Mesmo que 4 se percam, é muito provável que o quinto encontre o tesouro porque ele estava explorando um lugar que os outros não estavam.

O Que os Resultados Mostram?

Os autores testaram isso em modelos famosos (como Llama e Qwen) com problemas de matemática do ensino fundamental.

• Sem UpSkill: A IA melhora um pouco em acertar na primeira tentativa, mas se você pedir 5 tentativas, a chance de sucesso não aumenta muito porque as respostas são repetitivas.
• Com UpSkill: A chance de acertar em 5 tentativas (pass@5) aumentou significativamente (cerca de 3% a 10% a mais), sem piorar a chance de acertar na primeira tentativa.

Resumo em uma frase

O UpSkill ensina a Inteligência Artificial a não ser "teimosa" e a tentar resolver problemas de várias maneiras diferentes ao mesmo tempo, garantindo que, se uma estratégia falhar, outra provavelmente vai funcionar, aumentando muito a chance de sucesso em tarefas complexas.

É como transformar um time de jogadores que todos jogam no mesmo estilo em um time diversificado, onde cada um tem uma habilidade única, tornando o time muito mais difícil de ser derrotado.

Resumo técnico

1. O Problema

Os Grandes Modelos de Linguagem (LLMs) têm demonstrado excelente desempenho em tarefas de raciocínio verificável (como matemática e programação) quando otimizados para precisão em uma única tentativa (pass@1) usando Aprendizado por Reforço com Recompensas Verificáveis (RLVR). No entanto, existe um problema fundamental:

• Colapso de Diversidade: Quando esses modelos são amostrados múltiplas vezes para o mesmo prompt, tendem a gerar respostas altamente similares ou idênticas.
• Impacto no Pass@k: Em cenários onde a solução é avaliada pela probabilidade de pelo menos uma entre $k$ tentativas ser correta (métrica pass@k), a falta de diversidade reduz o número efetivo de tentativas independentes. Se todas as $k$ amostras forem redundantes, a probabilidade de sucesso não aumenta significativamente em relação a uma única tentativa.
• Limitação das Abordagens Atuais: Métodos de inferência (como aumento de temperatura ou perturbação de prompts) são frágeis e exigem ajuste manual. Métodos de treinamento existentes muitas vezes não expõem uma base controlável de estratégias ou sacrificam a precisão individual para ganhar diversidade.

O objetivo do trabalho é desenvolver um mecanismo de treinamento que aumente a diversidade de respostas de forma controlada, gerando modos de raciocínio semanticamente distintos e reproduzíveis, sem degradar a precisão da tentativa única (pass@1).

2. Metodologia: UpSkill

Os autores propõem o UpSkill, um método de treinamento que adapta o Aprendizado de Habilidades por Informação Mútua (MISL) para LLMs. A ideia central é introduzir uma variável latente discreta $z$ que estrutura a resposta, de modo que diferentes valores de $z$ correspondam a diferentes estratégias de resolução.

Componentes Principais:

1. Variável Latente Discreta ($z$):

   • Um token de prefixo leve (ex: "Estratégia {z} |") é injetado no prompt.
   • Durante o treinamento, $z$ é amostrado uniformemente de um conjunto $\{1, ..., N\}$.
   • Durante a inferência, selecionam-se $k$ valores distintos de $z$ para gerar $k$ tentativas diversas.

2. Recompensa de Informação Mútua (MI) em Nível de Token:

   • O objetivo é maximizar a informação mútua condicional $I(\tau; z | x)$, onde $\tau$ é a trajetória (resposta) e $x$ é o prompt.
   • Isso incentiva duas coisas:
     • Alta entropia marginal das trajetórias (cobertura ampla do espaço de soluções).
     • Baixa entropia condicional dado $z$ (cada estratégia $z$ deve ser reproduzível e distinta).
   • A recompensa é calculada token a token como a diferença entre o log-probabilidade da resposta dada a estratégia específica e a log-probabilidade média sobre todas as estratégias:
     $$r_{TMI}(\tau; x, z) = \sum_{t} \left[ \log p(y_t | x, z, y_{<t}) - \log p(y_t | x, y_{<t}) \right]$$
     Onde $p(y_t | x, y_{<t})$ é uma mistura uniforme de todas as estratégias.

3. Otimização com GRPO:

   • O método utiliza o Group Relative Policy Optimization (GRPO).
   • Para cada prompt, o modelo gera um grupo de completions ($C$) condicionadas a uma estratégia $z$.
   • A recompensa total combina:
     • Recompensa de correção verificável ($r_{corr}$).
     • Penalidade KL para manter o modelo próximo da base.
     • Recompensa de Informação Mútua ($r_{TMI}$): O termo inovador que promove a diversidade estruturada.

3. Contribuições Chave

• Método de Treinamento Controlado: UpSkill induz diversidade estruturada sem engenharia de prompts complexa, aprendendo modos de raciocínio distintos indexados por $z$.
• Conexão Teórica: Os autores provam teoricamente que a melhoria no pass@k está intimamente ligada ao objetivo de maximizar a informação mútua $I(\tau; z | x)$. Eles demonstram que a melhoria no pass@k é limitada inferiormente por uma função da informação mútua, garantindo que maximizar a MI leva a ganhos no pass@k.
• Desempenho sem Rótulos de Verdade: Demonstram empiricamente que o UpSkill pode melhorar o pass@k mesmo sem usar recompensas de correção (apenas com a recompensa de MI), embora a combinação seja mais eficaz.
• Preservação do Pass@1: Ao contrário de métodos que sacrificam a precisão para ganhar diversidade, o UpSkill mantém ou melhora a precisão de tentativa única em modelos base mais fortes.

4. Resultados Experimentais

Os experimentos foram conduzidos no conjunto de dados GSM8K (problemas de matemática do ensino fundamental) e em um ambiente de aritmética controlado, utilizando três modelos de pesos abertos: Llama 3.1-8B, Qwen 2.5-7B e R1-Distilled-Qwen2.5-Math-1.5B.

• GSM8K (Modelos Fortes):

  • No Qwen 2.5-7B, o UpSkill resultou em um ganho médio de ~3.4% no pass@k e ~9.1% no plurality@k, mantendo a precisão do pass@1.
  • No Llama 3.1-8B, também foram observados ganhos significativos no pass@k.
  • O método demonstrou que diferentes latentes $z$ correspondem a estratégias distintas (ex: abordagens algébricas vs. aritméticas diretas).

• Ambiente Aritmético Controlado:

  • Em um ambiente onde o modelo escolhe operadores e dígitos, o UpSkill evitou o colapso de entropia observado no treinamento padrão (GRPO apenas).
  • Enquanto o GRPO padrão convergia para uma estratégia determinística (pass@1 ≈ pass@5), o UpSkill manteve trajetórias diversas, elevando o pass@5 de 0.793 para 0.897, mesmo com um pass@1 ligeiramente menor (0.390), demonstrando a eficácia da diversidade.

• Análise de Modelos Menores (R1):

  • No modelo menor (R1-Distilled-1.5B), o método puro de MI causou degradação, sugerindo que modelos com capacidade limitada ou otimização prévia pesada podem ser sensíveis à introdução de múltiplas estratégias sem regularização KL adequada.

5. Significado e Conclusão

O trabalho UpSkill oferece uma solução principial para o dilema "exploração vs. exploração" em LLMs de raciocínio.

• Impacto Prático: Permite que sistemas de IA resolvam problemas complexos com maior confiabilidade ao gerar múltiplas tentativas que são verdadeiramente independentes e complementares, em vez de apenas variações superficiais da mesma resposta.
• Avanço Teórico: Estabelece um vínculo formal entre a teoria da informação (Informação Mútua) e métricas de sucesso prático (pass@k), validando que a diversidade estruturada é um motor direto para a melhoria de desempenho em tarefas de múltiplas tentativas.
• Reprodutibilidade: O código e os dados foram abertos, permitindo que a comunidade estude e refine a diversidade de respostas em modelos de linguagem.

Em resumo, o UpSkill demonstra que ensinar LLMs a "pensar de maneiras diferentes" (através de variáveis latentes e recompensas de informação mútua) é uma estratégia eficaz para superar as limitações de redundância e melhorar a robustez na resolução de problemas.