Como o Raciocínio Flui? Rastreando o Fluxo de Informação Induzido pela Atenção para RL Direcionado em LLMs

Resumo

A atribuição de crédito em nível de token continua sendo um obstáculo fundamental para o aprendizado por reforço (RL) em modelos de linguagem de grande porte (LLMs), onde as receitas de RL normalmente tratam todos os tokens igualmente, falhando em distinguir etapas de raciocínio decisivas de formatação rotineira ou preenchimento fluente. Tentativas recentes aproveitam sinais internos do modelo para atribuir crédito mais granular, mas essas são frequentemente heurísticas pontuais que ignoram a estrutura global da propagação de informações. Propomos o FlowTracer, uma estrutura de RL que rastreia o fluxo de raciocínio direcionado à resposta em um grafo acíclico direcionado induzido por atenção, no qual nós correspondem a tokens e as capacidades das arestas vêm de pesos de atenção agregados, e deriva crédito de token dessa estrutura global. As capacidades das arestas são reponderadas para reter apenas a influência que pode alcançar a região da resposta, enquanto impõem conservação de fluxo local para que tokens intermediários não percam nem ganhem massa efetiva devido ao comprimento do caminho ou ramificações irrelevantes. Nesse grafo, o FlowTracer extrai uma espinha dorsal de fluxo de informações conectando a pergunta à resposta e pontua tokens por vazão de fluxo, revelando hubs de alto impacto e pontos de verificação de agregação que medeiam dependências de longo alcance. Essas importâncias derivadas são usadas para moldar recompensas em nível de token, permitindo que os sinais de aprendizado se concentrem precisamente nos tokens que roteiam informações em direção a (ou afastando de) respostas corretas, e proporcionando ganhos de desempenho consistentes em uma variedade de tarefas de raciocínio.

English

Token-level credit assignment remains a key obstacle for reinforcement learning (RL) in large language models (LLMs), where RL recipes typically treat all tokens equally, failing to distinguish decisive reasoning steps from routine formatting or fluent filler. Recent attempts leverage model-internal signals to assign finer-grained credit, but these are often point-wise heuristics that ignore the global structure of information propagation. We propose FlowTracer, an RL framework that traces answer-targeted reasoning flow on an attention-induced directed acyclic graph in which nodes correspond to tokens and edge capacities come from aggregated attention weights and derives token credit from this global structure. The edge capacities are reweighted to retain only the influence that can reach the answer region, while enforcing local flow conservation so intermediate tokens neither lose nor gain effective mass due to path length or irrelevant branches. On this graph, FlowTracer extracts an information-flow backbone connecting the question to the answer and scores tokens by flow throughput, revealing high-impact hubs and aggregation checkpoints that mediate long-range dependencies. These derived importances are used to shape token-level rewards, enabling learning signals to focus precisely on the tokens that route information toward (or away from) correct answers and delivering consistent performance gains across a range of reasoning tasks.