RewardMap: Tackling Sparse Rewards in Fine-grained Visual Reasoning via Multi-Stage Reinforcement Learning

O artigo apresenta o RewardMap, um framework de aprendizado por reforço em múltiplos estágios que supera o desafio de recompensas esparsas no raciocínio visual detalhado de modelos de linguagem multimodal, utilizando o conjunto de dados ReasonMap-Plus e um design de recompensa sensível à dificuldade para alcançar melhorias consistentes em tarefas de raciocínio espacial e visual.

O essencial

Imagine que você está tentando ensinar um robô superinteligente a ler um mapa de metrô complexo e responder perguntas como "Qual é o melhor caminho da Estação A para a Estação B?" ou "Quantas estações existem entre essas duas?".

O problema é que, mesmo os robôs mais modernos (chamados de Modelos de Linguagem Multimodais) costumam se perder nesses mapas. Eles veem as linhas e os nomes, mas não conseguem "raciocinar" o caminho corretamente. E quando tentamos treiná-los com métodos comuns, eles desistem porque não recebem feedback suficiente: o robô só ganha um "ponto" se acertar a resposta final, mas se errar no meio do caminho, ele não sabe onde errou. É como tentar aprender a tocar piano apenas recebendo um "parabéns" ou "tente de novo" no final de uma música inteira, sem saber qual nota você tocou errado.

Aqui está a explicação do trabalho REWARDMAP, usando analogias simples:

1. O Problema: O "Deserto de Recompensas"

Pense no treinamento do robô como uma jornada em um deserto. O objetivo (a resposta certa) está muito longe. Se o robô só recebe um prêmio (recompensa) quando chega ao final, ele fica perdido no meio do caminho. Ele não sabe se está indo na direção certa ou se está andando em círculos. Isso é o que os cientistas chamam de recompensa esparsa (muito rara).

2. A Solução 1: O "Mapa de Recompensas" (REASONMAP-PLUS)

Os autores criaram um novo conjunto de dados chamado REASONMAP-PLUS.

• A Analogia: Em vez de jogar o robô direto no deserto difícil, eles criaram um "treinamento de obstáculos" que vai do fácil ao difícil.
• Como funciona: Eles começaram com perguntas simples, como "Quantas linhas tem neste mapa?" ou "Onde fica esta estação?". Essas perguntas são fáceis e dão prêmios frequentes (recompensas densas). Isso ajuda o robô a aprender o básico (percepção visual) antes de tentar resolver problemas complexos de lógica. É como aprender a andar de bicicleta com rodinhas antes de tentar descer uma montanha.

3. A Solução 2: O Treinador Inteligente (REWARDMAP)

Aqui entra a parte principal do paper: o sistema REWARDMAP. Ele é um método de treinamento em duas etapas principais:

• A. O Sistema de "Pontos Parciais" (Recompensa Detalhada):

  • A Analogia: Imagine um professor corrigindo uma prova de matemática. Se o aluno errar a conta final, mas fez a primeira etapa certa, o professor dá alguns pontos parciais.
  • Na prática: O REWARDMAP não espera apenas a resposta final. Ele verifica se o robô identificou corretamente a estação de partida, a de chegada, o nome da linha e as transferências. Se acertou a partida, ganha pontos. Se acertou a linha, ganha mais pontos. Isso transforma o "deserto" em uma trilha com várias placas de "Bom trabalho!", mantendo o robô motivado e sabendo exatamente onde corrigir.

• B. O Currículo em Etapas (Aprendizado Multi-Estágio):

  • A Analogia: É como subir uma escada. Você não pula do primeiro degrau para o último. Você sobe degrau por degrau.
  • Na prática: O robô primeiro treina apenas com perguntas fáceis (contar coisas, dizer sim/não). Depois, quando ele já é bom nisso, o sistema introduz perguntas médias e, por fim, as perguntas difíceis de planejar rotas. Isso evita que o robô fique frustrado ou confuso no início.

4. O Resultado: O Robô que "Enxerga" e "Pensa"

Depois de passar por esse treinamento especial, o robô não só ficou muito melhor em ler mapas de metrô, mas também melhorou em outras tarefas visuais.

• A Analogia: É como se, ao aprender a navegar em um labirinto de metrô com um guia excelente, o robô desenvolvesse um "instinto" melhor para entender qualquer desenho complexo, gráficos ou imagens detalhadas.
• O Sucesso: Os testes mostraram que esse método reduziu as "alucinações" (quando o robô inventa fatos) e a confusão visual. O robô passou a ver o mapa com mais clareza e a raciocinar o caminho de forma lógica.

Resumo em uma frase:

O REWARDMAP ensina robôs a entenderem mapas complexos não jogando-os direto no problema difícil, mas criando um "caminho de prêmios" que vai do simples ao complexo, dando feedback constante e detalhado a cada pequeno acerto, transformando um aprendizado difícil e confuso em uma jornada de sucesso passo a passo.

Resumo técnico

1. Problema e Motivação

O raciocínio visual de alta granularidade (fine-grained visual reasoning) em entradas estruturadas, como mapas de transporte público, representa um desafio significativo para os Modelos de Linguagem Multimodal (MLLMs). O benchmark REASONMAP demonstrou que mesmo modelos avançados falham em tarefas que combinam compreensão visual e raciocínio espacial (ex: planejamento de rotas).

A aplicação direta de Aprendizado por Reforço (RL) padrão para resolver essas tarefas enfrenta dois obstáculos principais:

1. Recompensas Esparsas: Em tarefas de raciocínio de longo prazo (long-chain reasoning), o sinal de recompensa (sucesso/fracasso) só é fornecido no final da cadeia de raciocínio. Se o modelo falhar em um passo intermediário, não recebe feedback imediato, o que desestabiliza a otimização e dificulta a exploração.
2. Ineficiência do Cold-Start: Estratégias tradicionais de Supervised Fine-Tuning (SFT) fornecem supervisão densa, mas muitas vezes falham em ensinar o modelo a realizar tomada de decisão de longo prazo e raciocínio complexo, levando a sobreajuste e rigidez cognitiva.

2. Metodologia: REWARDMAP

Para superar essas limitações, os autores propõem o REWARDMAP, um framework de RL multi-estágio, e o conjunto de dados REASONMAP-PLUS.

A. REASONMAP-PLUS (Construção do Dataset)

Os autores estenderam o dataset REASONMAP original criando o REASONMAP-PLUS, que introduz sinais de recompensa densos para facilitar o treinamento inicial (cold-start).

• Estrutura: Organizado em um continuum de dificuldade (fácil → médio → difícil).
• Tipos de Tarefas: Inclui 5 categorias de perguntas além do planejamento de rotas:
  1. Contagem Global: Número total de linhas no mapa.
  2. Contagem Local 1: Número de paradas intermediárias entre duas estações.
  3. Contagem Local 2: Número de linhas que passam por uma estação específica.
  4. Verdadeiro ou Falso 1: Relação espacial entre duas estações.
  5. Verdadeiro ou Falso 2: Se uma estação pertence a uma linha específica.
• Anotação: As perguntas são geradas automaticamente a partir de dados estruturados de metrô (Metro Data) e revisadas manualmente para garantir diversidade e equilíbrio de dificuldade.

B. Arquitetura do REWARDMAP

O framework utiliza o algoritmo GRPO (Group Relative Policy Optimization) e integra dois componentes principais:

1. Design de Recompensa Consciente da Dificuldade (Difficulty-Aware Reward Design):
   Para mitigar a esparsidade, a função de recompensa ($R$) é composta por três termos ponderados:
   $$R = W_{difficulty} \times (R_{format} + R_{correctness} + \alpha \times R_{detail})$$

   • Recompensa de Formato ($R_{format}$): Garante que a saída siga o formato exigido (ex: \boxed{}).
   • Recompensa de Corretude ($R_{correctness}$): Pontuação exata para a resposta final.
   • Recompensa de Detalhes ($R_{detail}$): Inovação Chave. Concede crédito parcial por itens corretos na resposta (ex: nome da linha, estações de partida/chegada, estações de transferência), mesmo que a resposta final não esteja 100% correta. Isso fornece feedback denso durante o treinamento.
   • Ponderação por Dificuldade ($W_{difficulty}$): Ajusta a magnitude da recompensa baseada na dificuldade do mapa (fácil/médio/difícil) e da pergunta (número de transferências necessárias).

2. Currículo de RL Multi-Estágio (Multi-Stage RL Scheme):
   Em vez de iniciar com SFT, o treinamento começa diretamente com RL, seguindo um currículo estruturado:

   • Princípio Global: As tarefas são agendadas de simples para complexas (ex: de tarefas de percepção binária/contagem para raciocínio de planejamento complexo).
   • Princípio Local Estocástico: Dentro de cada estágio, as amostras são embaralhadas para evitar sobreajuste a um trajeto fixo de dificuldade.
   • Objetivo: Usar as tarefas densas do REASONMAP-PLUS para estabelecer uma política robusta antes de enfrentar as tarefas esparsas de planejamento do REASONMAP original.

3. Contribuições Principais

1. REASONMAP-PLUS: Um dataset estendido com 4.018 perguntas, organizado por dificuldade e cobrindo múltiplos tipos de tarefas de raciocínio visual fino, fornecendo supervisão densa para o cold-start do RL.
2. Framework REWARDMAP: Uma nova abordagem de RL multi-estágio que combina:
   • Um esquema de recompensa com detalhes parciais para aliviar a esparsidade.
   • Um currículo de aprendizado que transita da percepção visual simples para o raciocínio complexo sem depender de SFT inicial.
3. Validação Empírica: Demonstração de que a combinação desses componentes supera significativamente as abordagens de SFT e RL padrão, melhorando não apenas em mapas, mas também em benchmarks gerais de raciocínio visual.

4. Resultados Experimentais

Os experimentos foram conduzidos em modelos baseados na família Qwen2.5-VL (3B, 7B, 32B, 72B) e Kimi-VL.

• Desempenho em REASONMAP e REASONMAP-PLUS:

  • O modelo REWARDMAP superou consistentemente todas as linhas de base (SFT, RL padrão, SFT+RL).
  • No REASONMAP, alcançou 31.51% de precisão ponderada (versus 29.51% do RL de base e 26.22% do RL sem recompensas de detalhe), aproximando-se do desempenho de modelos proprietários fechados como o Seed1.5-VL.
  • No REASONMAP-PLUS, superou modelos de código aberto de grande porte e até o Seed1.5-VL.

• Generalização em Outros Benchmarks:

  • O modelo treinado com REWARDMAP mostrou melhorias consistentes em 6 benchmarks diversos (SEED-Bench-2-Plus, SpatialEval, V*Bench, HRBench, ChartQA, MMStar).
  • Houve um ganho médio de 3.47% em todos os benchmarks.
  • Destaque no SpatialEval, com um aumento massivo de 13.51% na precisão, indicando uma melhoria substancial na inteligência espacial.

• Análise Qualitativa:

  • O REWARDMAP reduziu significativamente alucinações e confusão visual (ex: identificar linhas ou estações erradas) em comparação com modelos de base e RL padrão.
  • As recompensas de detalhe permitiram que o modelo aprendesse a localizar informações corretas no mapa mesmo antes de dominar o planejamento completo da rota.

• Estudos de Ablação:

  • Tanto o design de recompensa (detalhes) quanto o design multi-estágio contribuíram individualmente para o ganho de desempenho.
  • A abordagem é robusta em diferentes escalas de modelo (3B a 72B) e arquiteturas (Qwen, Kimi).

5. Significância e Impacto

O trabalho REWARDMAP oferece uma solução principiante para o problema fundamental de recompensas esparsas em raciocínio visual estruturado. Ao introduzir recompensas de detalhe e um currículo de aprendizado escalonado, os autores demonstram que é possível treinar MLLMs para tarefas complexas de raciocínio espacial sem depender exclusivamente de SFT massivo ou dados sintéticos de baixa qualidade.

A metodologia não se limita a mapas de transporte; os resultados em benchmarks de gráficos (ChartQA) sugerem que a abordagem é generalizável para qualquer domínio visual estruturado que exija raciocínio de múltiplos passos e compreensão de alta fidelidade. Isso avança o estado da arte na capacidade de MLLMs de realizar tarefas práticas de navegação e análise espacial no mundo real.