🤖 machine learning

From Reasoning Traces to Reusable Modules: Understanding Compositional Generalization in Language Model Reasoning

This paper proposes that the success of combining supervised fine-tuning (SFT) and reinforcement learning (RL) in enhancing language model reasoning stems from compositional generalization, where SFT provides raw atomic modules within traces and RL decomposes and recombines them to solve novel configurations, a theory validated by experiments showing that training on compound traces yields superior generalization compared to isolated modules.

Original authors: Lingjing Kong, Xin Liu, Guangyi Chen, Martin Q. Ma, Xiangchen Song, Yuekai Sun, Mikhail Yurochkin, Taylor W. Killian, Ruslan Salakhutdinov, Kun Zhang, Eric P. Xing, Zhengzhong Liu

Published 2026-06-17

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Lingjing Kong, Xin Liu, Guangyi Chen, Martin Q. Ma, Xiangchen Song, Yuekai Sun, Mikhail Yurochkin, Taylor W. Killian, Ruslan Salakhutdinov, Kun Zhang, Eric P. Xing, Zhengzhong Liu

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Picture: How AI Learns to "Think"

Imagine you are teaching a robot to solve a complex puzzle. You have two main ways to teach it:

Show-and-Tell (SFT): You show the robot a perfect, step-by-step solution to a specific puzzle.
Trial-and-Error (RL): You let the robot try to solve the puzzle on its own. If it gets the right answer at the end, you give it a "gold star" (reward). If it fails, it gets nothing.

For a long time, researchers noticed that combining these two methods works best. The robot learns the steps from "Show-and-Tell" but becomes much better at solving new puzzles it has never seen before after "Trial-and-Error."

The Question: Why does this combination work so well? What is actually happening inside the robot's brain?

The Core Idea: The "Lego" Theory

The authors of this paper propose a new way to understand this process. They argue that reasoning isn't just memorizing a long story; it's like building with Lego bricks.

They break down a reasoning trace (the robot's thought process) into two types of reusable parts:

Skills (The Bricks): These are small, local actions. Examples: "Add these two numbers," "Remove this word," or "Check if this is true."
Routing (The Instructions): These are the rules for how to connect the bricks. Examples: "Take the result of step 1 and feed it into step 2," or "If the number is big, skip to step 4."

The Problem with "Show-and-Tell" (SFT):
When you show the robot a perfect solution, the "bricks" and the "instructions" are glued together in a specific order. The robot sees a long, solid block of Lego. It learns to copy the whole block, but it doesn't realize that the block is made of separate, reusable pieces. If you give it a puzzle that requires the bricks in a different order, the robot gets stuck because it only knows how to copy the original block.

The Magic of "Trial-and-Error" (RL):
When the robot starts trying to solve problems on its own and gets gold stars for correct answers, something amazing happens. It starts to realize: "Wait, I used that same 'Add' brick in three different places, and it worked every time!"

The RL process acts like a deconstructor. It takes those glued-together blocks from the "Show-and-Tell" phase and breaks them back down into individual, reusable bricks and instructions. Once the robot has a box of loose, reusable bricks, it can snap them together in new ways to solve puzzles it has never seen before.

The Key Findings (The Recipe for Success)

The paper ran experiments to prove this theory. Here is what they found, explained simply:

1. You need the raw materials first.
You can't teach a robot to build with Lego if you never show it what a Lego brick looks like. The "Show-and-Tell" phase is crucial because it provides the raw materials (the atomic skills and routing mechanisms). Without this, the robot has nothing to break down.

2. Breaking things down is better than just showing them.
If you only show the robot isolated bricks (e.g., "Here is how to add numbers"), it learns to add, but it doesn't learn how to combine adding with other steps.

The Finding: It is much better to show the robot a full, complex solution (a compound trace) and let the "Trial-and-Error" phase figure out how to take it apart. The robot learns to recombine the pieces much faster and more effectively than if you just gave it a pile of isolated bricks.

3. The "Glue" must be broken in the right place.
The authors discovered a specific recipe for the best results:

Phase 1 (Show-and-Tell): Show the robot many different complex solutions that cover every possible type of brick and instruction. Make sure it sees all the tools in the toolbox.
Phase 2 (Trial-and-Error): Let the robot practice on puzzles that mix these tools in new, weird combinations that it hasn't seen before.
Why? If you let the robot practice on the same combinations it saw in Phase 1, it just memorizes them. But if you let it practice on new combinations, it is forced to figure out how to snap the loose bricks together on the fly. This is where the "generalization" (the ability to solve new things) happens.

A Simple Analogy: The Chef

Imagine you are training a chef.

SFT (Show-and-Tell): You show the chef a recipe for "Spaghetti Carbonara." You show them exactly how to crack the eggs, boil the pasta, and mix the sauce. The chef learns to copy this one dish perfectly.
RL (Trial-and-Error): Now, you tell the chef, "Make me a delicious dinner," but you don't give a recipe. You let them try. If they make a great meal, you say "Good job!"

The Paper's Insight:
If you only show the chef the Carbonara recipe, they might think "Cracking eggs" and "Boiling pasta" are one single, unbreakable action called "Making Carbonara."
But when you let them experiment (RL), they realize: "Oh! I can crack eggs for an omelet, or I can boil pasta for a salad." They realize the skills (cracking, boiling) are separate from the routing (what to do with the pasta next).

Once they realize this, they can invent a brand new dish (like "Pasta Omelet") that they've never seen before, because they have learned to treat the ingredients as separate, reusable tools rather than a single fixed script.

Summary

Reasoning is modular: It's made of small skills and rules for connecting them.
SFT provides the parts: It gives the model the raw materials (the skills).
RL does the assembly: It forces the model to break the materials apart and learn how to snap them together in new ways.
The Best Strategy: Show the model a wide variety of complex examples first (to get all the parts), then let it practice on new combinations of those parts (to learn how to build).

This paper explains why the current "Show-and-Tell then Trial-and-Error" method is so powerful: it turns a rigid memorization machine into a flexible builder.

Technical Summary: From Reasoning Traces to Reusable Modules

Problem Statement

Post-training pipelines combining Supervised Fine-Tuning (SFT) and Reinforcement Learning (RL) have recently driven significant improvements in the reasoning capabilities of Large Language Models (LLMs). Empirically, SFT alone often leads to models that memorize canonical "golden" reasoning traces but fail when required to recombine familiar steps in unfamiliar ways (Out-of-Distribution or OOD compositions). Conversely, SFT followed by RL handles such OOD compositions more reliably. However, the underlying mechanism driving this combined success remains unclear: specifically, what distinct contributions do SFT and RL make, and what data configurations are optimal for each stage?

Methodology

Theoretical Framework: Hierarchical Latent Selection Model

The authors formalize reasoning as a Hierarchical Latent Selection Model. In this framework:

Problem Descriptor ( $P$ ): The input prompt induces a hierarchy of discrete latent selection variables.
Latent Variables ( $S$ ): These variables select from a library of reusable atomic modules:
- Skills: Local operations (e.g., arithmetic steps, rewrite rules).
- Routing Mechanisms: Decisions on how intermediate information is selected, reused, and composed (e.g., forwarding previous outputs, branching).
Observable Trace ( $D$ ): The final reasoning trace is generated by the cascade of these selections.

The core hypothesis is that reasoning traces are generated by composing these atomic modules. SFT provides the raw material (traces containing these modules), while RL is responsible for decomposing these compound traces to identify the latent modules and enable their flexible recombination.

Theoretical Contributions

The paper establishes two key theoretical results:

Identifiability (Theorem 3.1): Under mild conditions (faithfulness, rank faithfulness, restricted choice sets, and neighborhood coverage), the latent selection hierarchy (skills and routing) is identifiable from the observable distribution of traces.
Compositional Generalization (Theorem 3.4): Once modules are identified, they can be recombined on novel compositions if every local configuration in the new prompt has a "local witness" in the training data (i.e., the interface between modules has been observed during training).

The authors further characterize the roles of SFT and RL through support analysis:

SFT Limitations (Proposition 3.2): SFT on canonical traces often leaves atomic modules "entangled" with their specific contexts because it reveals only a subset of valid traces. Critical local events that would disentangle modules may lie outside the SFT support.
RL Mechanism (Proposition 3.3): RL enriches the distribution by re-sampling rollouts based on reward. If a "hidden" local event (e.g., a specific routing choice) leads to higher rewards, RL increases its probability, effectively decomposing the entangled trace and making the atomic module identifiable.

Empirical Setup

To validate these theories, the authors conducted controlled experiments on synthetic string transformation tasks (based on Yuan et al., 2025).

Atomic Skills: 24 deterministic string transformation functions.
Routing Mechanisms: 10 distinct input-structure definitions (determining which intermediate outputs feed into the next step).
Protocol: Models were trained on traces of varying depths ( $L$ $L$ ). The study systematically intervened on the data structure to test:
- Decomposition of compound traces into atoms.
- The necessity of compositional exposure vs. isolated atomic exposure.
- The optimal distributional relationship between SFT and RL data (overlap vs. disjoint).
Supplementary Evidence: Analysis of open-source models (Qwen3-4B) on real-world math benchmarks (MATH, AIME) to observe skill tokenization and n-gram diversity.

Key Results

RL Decomposes and Recombines: Models trained with RL on compound traces ( $L=3$ ) significantly outperformed SFT-only models on unseen compositions ( $L=4$ to $L=8$ ). RL enabled the model to recover atomic skills and routing mechanisms from compound traces and recombine them for generalization.
Compound Traces are Superior for RL: When RL was trained on isolated atomic modules ( $L=1$ ) versus compound traces ( $L=3$ ), the latter yielded significantly higher OOD accuracy (+37.8% gain vs. +10.0% gain). This confirms that exposure to module interfaces (compositions) is critical for learning how to recombine them.
Re-injection Requires Composition: When specific atomic skills or routers were withheld from RL training, performance dropped. Simply re-injecting isolated atomic skills ( $L=1$ ) did not restore performance. Performance was only recovered when compositional traces containing the held-out atoms were re-injected, proving that atoms must be learned within a compositional context to be reusable.
Optimal SFT-RL Data Pairing: The strongest OOD generalization occurred when SFT and RL data distributions were disjoint ( $C_{SFT} \cap C_{RL} = \emptyset$ $C_{S F T} \cap C_{R L} = \emptyset$ ).
- SFT Role: Should cover the full inventory of atomic modules through compositional traces.
- RL Role: Should focus on novel compositions outside the SFT support to drive exploration and identify new interfaces.
- Overlapping data ( $C_{SFT} \supset C_{RL}$ ) resulted in the weakest generalization.
Real-Model Evidence: Analysis of Qwen3 models showed that RL-tuned models exhibited higher diversity in short skill n-grams (n=2, 3), indicating more flexible local recombination of abstracted steps, consistent with the theoretical model of modular reuse.

Significance and Claims

The paper claims to provide a latent-variable explanation for the success of SFT+RL pipelines in compositional reasoning. Its primary significance lies in:

Formalizing the Mechanism: Moving beyond empirical observation to a theoretical framework where SFT supplies "materials" (entangled traces) and RL performs "decomposition" (identifying latent modules via reward-driven support expansion).
Curriculum Design: Offering a concrete protocol for post-training: SFT should ensure coverage of all atomic modules via compositions, while RL should target novel compositions to explore new interfaces.
Identifiability Conditions: Providing sufficient conditions under which latent reasoning structures can be recovered from observed data, linking the success of RL to the enrichment of "local witnesses" required for compositional generalization.

The authors position this work as a step toward understanding how LLMs learn to reason compositionally, suggesting that the "black box" of RL is effectively a mechanism for discovering and recombining reusable atomic reasoning units.