Comment le raisonnement circule-t-il ? Traçage du flux d'information induit par l'attention pour un RL ciblé dans les LLMs

Résumé

L'attribution de crédit au niveau des jetons demeure un obstacle majeur pour l'apprentissage par renforcement (RL) dans les grands modèles de langage (LLMs), où les recettes de RL traitent généralement tous les jetons de manière égale, échouant à distinguer les étapes de raisonnement décisives du formatage de routine ou des remplissages fluides. Des tentatives récentes exploitent les signaux internes des modèles pour attribuer un crédit plus fin, mais il s'agit souvent d'heuristiques ponctuelles qui ignorent la structure globale de la propagation de l'information. Nous proposons FlowTracer, un cadre de RL qui trace le flux de raisonnement ciblant la réponse sur un graphe orienté acyclique induit par l'attention, dans lequel les nœuds correspondent aux jetons et les capacités des arêtes proviennent des poids d'attention agrégés, et qui dérive le crédit des jetons à partir de cette structure globale. Les capacités des arêtes sont repondérées pour ne conserver que l'influence pouvant atteindre la région de la réponse, tout en imposant une conservation locale du flux afin que les jetons intermédiaires ne perdent ni ne gagnent de masse effective en raison de la longueur du chemin ou de branches non pertinentes. Sur ce graphe, FlowTracer extrait une colonne vertébrale de flux d'information reliant la question à la réponse et évalue les jetons en fonction du débit de flux, révélant des hubs à fort impact et des points de contrôle d'agrégation qui médiatisent les dépendances à longue portée. Ces importances dérivées sont utilisées pour façonner des récompenses au niveau des jetons, permettant aux signaux d'apprentissage de se concentrer précisément sur les jetons qui acheminent l'information vers (ou loin des) réponses correctes, et générant des gains de performance constants dans un éventail de tâches de raisonnement.

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.