Resumo Técnico: Manifold Bandits – Aprendizado de Currículo Bayesiano sobre a Geometria Latente de Grandes Modelos de Linguagem

1. Declaração do Problema

O Aprendizado por Reforço (RL) tornou-se um método primário para aprimorar as capacidades de raciocínio de Grandes Modelos de Linguagem (LLMs), particularmente através de algoritmos de otimização de grupo-relativa como GRPO, DAPO e GSPO. Esses métodos dependem da variância da recompensa dentro de um grupo de rollouts para gerar gradientes de política. Se um modelo resolve consistentemente ou falha consistentemente em um prompt, a variância da recompensa é zero, resultando em nenhum sinal de aprendizado e desperdiçando recursos computacionais.

Métodos existentes de aprendizado de currículo adaptativo tentam abordar isso priorizando prompts de dificuldade intermediária (aqueles com probabilidade de gerar variância não nula). No entanto, essas abordagens tipicamente enquadram a seleção de problemas como um problema de bandido multi-braços padrão com braços independentes. Esse enquadramento ignora duas dinâmicas críticas do treinamento de LLMs:

1. Não-Estacionariedade Endógena: A distribuição de recompensa sobre os prompts evolui não apenas devido a fatores externos, mas como uma consequência direta das próprias atualizações e decisões de amostragem do modelo.
2. Heterogeneidade Estruturada: Problemas não são independentes; prompts de tipos semelhantes compartilam padrões de raciocínio e características latentes. Ignorar esses relacionamentos leva a uma exploração ineficiente.

Além disso, as estratégias atuais focam apenas na produtividade (maximizar o sinal de aprendizado imediato) enquanto negligenciam a diversidade (cobertura do espaço de tarefas) e a utilidade (relevância para benchmarks de avaliação de downstream). Isso pode levar a uma cobertura estreita e generalização subótima.

2. Metodologia

Os autores propõem o Bayesian Manifold Curriculum (BMC), um framework que reestrutura a amostragem de problemas como um problema de bandido estruturado em variedade (manifold) com não-estacionariedade endógena. A metodologia consiste em três componentes principais:

A. Árvores de Tarefas Latentes (Latent Task Trees)

Em vez de tratar prompts como braços independentes, o BMC constrói uma Árvore de Tarefas Latente hierárquica diretamente das representações internas do modelo de política.

• Construção: A árvore é construída recursivamente usando embeddings de uma camada intermediária do modelo de política. Em cada nó, os embeddings são reduzidos via PCA, projetados via UMAP e agrupados usando HDBSCAN.
• Critério de Parada: A recursão para quando uma região passa por um "Teste de Gráfico" (Chart Test), verificando se a variedade local é aproximadamente suave e coerente (baixa dimensionalidade intrínseca e conectividade).
• Propósito: Esta hierarquia aproxima a variedade de tarefas latente do modelo, organizando prompts em regiões de múltiplas escalas que refletem a percepção interna do modelo sobre as relações entre tarefas, em vez de depender de anotações externas ou taxonomias fixas.

B. Bayesian Manifold Curriculum (BMC)

O BMC realiza a tomada de decisão estruturada sobre a Árvore de Tarefas Latente usando um ciclo de três etapas:

1. Seleção Top-Down (Thompson Sampling Hierárquico): Para formar um lote de treinamento, agentes percorrem a árvore da raiz até as folhas. Em cada nó, eles amostram estimativas de recompensa para os nós filhos a partir de uma distribuição de crença Gaussiana ($\mathcal{N}(\mu, \sigma^2)$) e selecionam o filho com a maior amostra. Isso equilibra exploração (diversidade) e explotação (produtividade) através de diferentes níveis de abstração.
2. Modelagem de Crença Não-Estacionária (Filtragem Bayesiana): Os nós folha (prompts individuais) mantêm crenças sobre seu sinal de aprendizado esperado (variância da recompensa). Essas crenças são atualizadas online usando um filtro Bayesiano leve inspirado no modelo Normal–Inverse-Gamma. O mecanismo de atualização incorpora surpresa (erro de previsão padronizado) e obsolescência (tempo desde a última amostragem) para se adaptar às capacidades evolutivas do modelo.
3. Propagação de Crença Bottom-Up (Empirical Bayes): Após as atualizações, a informação se propaga para cima. As crenças dos nós pais são estimadas como médias ponderadas pela precisão de seus filhos, ajustadas para a heterogeneidade do sub-árvore. Isso permite que observações de um prompt informem as crenças sobre prompts relacionados na mesma região, explorando a estrutura latente.

C. Extensão Orientada à Utilidade (BMC-T)

Para abordar o desalinhamento entre a produtividade de treinamento e o desempenho de avaliação downstream, os autores introduzem o BMC-T. Esta variante enviesa a amostragem para regiões do treinamento que sobrepõem uma distribuição alvo (ex: um benchmark de avaliação específico) no espaço latente. Ele trata a Árvore de Tarefas Latente como um modelo de intervenção aproximado, assumindo que treinar em tarefas próximas produz efeitos de transferência mais fortes.

3. Principais Contribuições

1. Árvores de Tarefas Latentes: Um método hierárquico para aproximar a estrutura da variedade de tarefas diretamente dos embeddings do modelo, fornecendo uma representação de múltiplas escalas das relações entre prompts sem rótulos externos.
2. Bayesian Manifold Curriculum (BMC): Um framework de currículo adaptativo consciente da estrutura que realiza tomada de decisão Bayesiana sobre essas árvores, permitindo uma exploração eficiente enquanto considera a não-estacionariedade endógena induzida pelo modelo.
3. Trade-offs Produtividade–Diversidade–Utilidade: Uma análise diagnóstica demonstrando que currículos adaptativos não podem ser caracterizados apenas pelo desempenho downstream. Os autores distinguem três eixos:
   • Produtividade: Sinal de aprendizado (variância da recompensa).
   • Diversidade: Cobertura da variedade de tarefas.
   • Utilidade: Relevância para benchmarks de avaliação específicos.
     Eles mostram que esses eixos não são intercambiáveis; otimizar um não garante a otimização dos outros.

4. Resultados Experimentais

Os autores avaliaram o BMC em tarefas de raciocínio matemático usando o dataset DAPO-Math-17K e modelos Qwen3 (4B e 8B) com GSPO e GRPO.

• Eficiência de Treinamento (Produtividade): O BMC e sua ablação "Difficulty Only" alcançaram velocidades de aprendizado comparáveis ao Dynamic Sampling (que regenera lotes para garantir variância), mas com tempo de execução significativamente menor, permanecendo próximo à eficiência do Uniform Sampling.
• Cobertura e Compartilhamento de Informação (Diversidade): O BMC equilibrou com sucesso a diversidade do Tree Only (que enfatiza a cobertura) com a produtividade do Difficulty Only. Métricas mostraram que o BMC proporcionou melhor cobertura de regiões sub-representadas (exposição ponderada pela raridade) e que a Árvore de Tarefas Latente explicou uma variação significativa nos sinais de aprendizado (ganho de estrutura), confirmando que os sinais de aprendizado estão correlacionados com a estrutura latente.
• Desempenho de Avaliação (Utilidade):
  • O Dynamic Sampling teve um desempenho ruim em benchmarks menos representados (ex: matemática em chinês, ciência OOD) devido ao filtragem implícita de tipos de prompts raros.
  • O BMC alcançou um forte desempenho no benchmark out-of-domain GPQA-Diamond, sugerindo que a otimização conjunta de produtividade e diversidade promove capacidades não capturadas por avaliações padrão de matemática in-distribution.
  • O BMC-T demonstrou que a utilidade pode ser direcionada: visar benchmarks específicos (ex: AIME2024) melhorou o desempenho nesses benchmarks específicos enquanto preservava a produtividade geral, confirmando que diferentes tipos de problemas carregam diferentes relevâncias de avaliação.

5. Significância e Alegações

O artigo argumenta que a dificuldade por si só é insuficiente para um forte desempenho downstream. Maximizar o sinal de aprendizado (produtividade) não garante melhores resultados se o modelo não for exposto aos tipos de problemas que impulsionam a generalização para tarefas alvo.

Os autores alegam que a seleção de problemas deve ser enquadrada como um problema de bandido estruturado em variedade sobre a geometria latente da política. Ao derivar a estrutura do currículo das próprias representações aprendidas pela política (aderindo à "Lição Amarga" de usar dados e computação para aprender a estrutura em vez de impor definições humanas), o BMC oferece uma forma de:

• Navegar pelos trade-offs entre produtividade, diversidade e utilidade.
• Explorar a não-estacionariedade endógena onde as decisões de amostragem moldam futuros cenários de recompensa.
• Generalizar através de domínios, línguas e modalidades sem exigir taxonomias definidas por humanos ou modelos de embedding externos.

O trabalho destaca um fenômeno de "zona morta de avaliação": um modelo pode ser treinado produtivamente em um conjunto específico de problemas, mas se esses problemas não se alinharem com a distribuição de avaliação, o desempenho downstream pode não melhorar. O BMC-T fornece um mecanismo para endereçar isso ao direcionar o esforço de treinamento para regiões relevantes ao alvo na variedade latente.