ReLaX: Reasoning with Latent Exploration for Large Reasoning Models

Resumo Técnico: ReLaX

1. O Problema: Colapso de Entropia e Exploração Ineficiente

O artigo aborda um gargalo fundamental no treinamento de Grandes Modelos de Raciocínio (LRMs) utilizando Aprendizado por Reforço com Recompensas Verificáveis (RLVR).

• Colapso de Entropia: Embora o RLVR tenha demonstrado potencial para melhorar o raciocínio, ele tende a levar a política do modelo para um estado de sobre-determinismo. A distribuição de probabilidade dos tokens converge rapidamente, reduzindo a entropia e confinando o gradiente da política em um subespaço estreito.
• Exploração vs. Exploração: A recompensa esparsa no RLVR exacerba esse problema, fazendo com que a política "explorar" prematuramente (over-exploit), inibindo a busca por soluções inovadoras e levando a um desempenho subótimo.
• Limitações das Abordagens Atuais: Métodos existentes focam em aumentar a entropia ao nível de tokens (ex: regularização de entropia, recompensas baseadas em tokens). No entanto, o artigo argumenta que isso é um sintoma superficial. O verdadeiro problema reside na rigidez das dinâmicas latentes internas (estados ocultos) que governam a geração de tokens. Além disso, em Modelos Multimodais de Linguagem (MLLMs), há um desalinhamento entre os cálculos internos multimodais e as saídas textuais unimodais, tornando o feedback ao nível de tokens ineficiente para capturar a complexidade do processamento multimodal.

2. Metodologia: ReLaX e Dinâmica Latente

A proposta central do trabalho é a ReLaX (Reasoning with Latent eXploration), um framework que regula a exploração e exploração monitorando e intervindo nas dinâmicas latentes do modelo, em vez de apenas nos tokens de saída.

Componentes Principais:

• Teoria do Operador de Koopman:

  • Para tornar as dinâmicas não lineares e de alta dimensão dos estados ocultos do LRM analisáveis, os autores utilizam a Teoria do Operador de Koopman.
  • Eles mapeiam a evolução não linear dos estados ocultos ($x_t$) para um espaço de observáveis infinito-dimensional onde a dinâmica se torna linear.
  • Utilizam uma Dicionário Neural Koopman (baseado em ResKoopNet) para aprender essa representação linear de forma estável e precisa.

• Dispersão Espectral Dinâmica (DSD - Dynamic Spectral Dispersion):

  • Introduzem uma nova métrica, o DSD, para quantificar a heterogeneidade das dinâmicas latentes.
  • O DSD é definido como a variância das magnitudes dos autovalores do operador de Koopman aproximado.
  • Interpretação: Um DSD alto indica um espectro dinâmico rico e diversificado (flexibilidade computacional), enquanto um DSD baixo sugere um comportamento rígido e repetitivo (colapso de modo).

• Regularização Baseada em DSD no Objetivo de Otimização:

  • O DSD é incorporado diretamente na função de perda do GRPO (Group Relative Policy Optimization).
  • Termo de Regularização ($\tilde{L}_{xp}$): O objetivo penaliza a convergência para modos latentes homogêneos. A regularização é ponderada pelos advantages (vantagens) positivos, garantindo que a exploração ocorra apenas em trajetórias que geram recompensas úteis.
  • Regularização KL Adaptativa: Para evitar instabilidade causada por uma dispersão excessiva, aplica-se uma penalidade KL apenas nas trajetórias onde o DSD excede um certo limiar, permitindo que trajetórias com potencial de exploração prossigam livremente.

3. Contribuições Chave

1. Métrica DSD: Proposição de uma métrica fundamental que captura a heterogeneidade das dinâmicas de raciocínio latente, oferecendo uma caracterização mais profunda da exploração da política do que as estatísticas de superfície de tokens.
2. Framework ReLaX: Desenvolvimento de um novo paradigma de treinamento que utiliza a regularização baseada em DSD para mitigar o colapso de modos e facilitar um equilíbrio eficaz entre exploração e exploração.
3. Superação de Métodos Baseados em Tokens: Demonstração de que intervir no espaço latente contínuo é mais eficaz e generalizável do que tentar forçar a diversidade apenas no espaço discreto de tokens, especialmente em modelos multimodais.

4. Resultados Experimentais

Os autores realizaram extensos experimentos em modelos de 3B e 7B parâmetros, tanto puramente textuais (LLMs) quanto multimodais (VLMs/MLLMs).

• Raciocínio Multimodal (VLMs):

  • O ReLaX-VL-7B alcançou um novo estado da arte (SOTA) com média de 53.2 em 7 benchmarks de raciocínio multimodal (incluindo MMMU, MathVista, etc.), superando o anterior líder (VL-Rethinker-7B com 52.5).
  • O modelo de 3B superou vários modelos de 7B anteriores.
  • A análise de dinâmica de treinamento mostrou que o ReLaX mantém uma entropia e diversidade latente (DSD) mais altas e estáveis, evitando o colapso rápido observado no GRPO padrão.

• Raciocínio Matemático Textual (LLMs):

  • Em benchmarks como MATH500, Minerva, AMC e AIME, o ReLaX superou consistentemente o GRPO padrão e variantes avançadas como SimpleRL, DAPO, KL-Cov e FR3E.
  • No modelo Qwen2.5-7B-Base, superou o SOTA anterior (FR3E) em 4.3 pontos de média.

• Análise de Caso:

  • Estudos qualitativos mostraram que modelos treinados com ReLaX realizam verificações de raciocínio mais robustas e baseadas em princípios (ex: leis físicas/matemáticas), enquanto métodos que apenas aumentam a entropia de tokens tendem a gerar "alucinações" ou estratégias ineficazes (ex: gerar código Python sem ambiente de execução).
  • Em tarefas multimodais, o ReLaX demonstrou maior robustez a variações visuais (ex: mudanças numéricas em diagramas) comparado a métodos que focam apenas em tokens.

5. Significado e Impacto

O trabalho ReLaX representa uma mudança de paradigma na otimização de modelos de raciocínio:

• Do Token para o Latente: Demonstra que a chave para melhorar o raciocínio não está apenas em diversificar a saída de tokens, mas em manter a flexibilidade dos processos computacionais internos (estados ocultos).
• Solução para MLLMs: Oferece uma solução principial para o desalinhamento em modelos multimodais, onde o feedback baseado em texto é insuficiente para guiar o processamento visual complexo.
• Escalabilidade: Ao utilizar a teoria do Operador de Koopman, o método torna as dinâmicas complexas tratáveis e otimizáveis, abrindo caminho para o desenvolvimento de modelos de raciocínio mais robustos e generalizáveis, superando os limites de saturação impostos pelos métodos atuais de RLVR.

Em suma, o ReLaX prova que guiar a exploração através da dinâmica latente é uma solução mais robusta e escalável para o dilema exploração-exploração em Grandes Modelos de Raciocínio.