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

2026-06-16

AI Summary

This work investigates compositional generalization in large language models on reasoning tasks and proposes a hierarchical latent variable selection model that represents reasoning trajectories as compositions of reusable atomic skills and a routing mechanism. It formally characterizes, for the first time, the generative process underlying compositional generalization and reveals an asymmetric complementary relationship between supervised fine-tuning (SFT) and reinforcement learning (RL): SFT provides composite reasoning trajectories, while RL disentangles atomic modules from them to enable cross-task recomposition and generalization. Through controlled synthetic experiments and module disentanglement techniques, the study demonstrates that training on composite trajectories outperforms training on isolated modules. Furthermore, it introduces a synergistic SFT-RL training strategy based on coverage and exploration, which substantially enhances the model’s compositional generalization capabilities.

Abstract

Post-training pipelines that combine supervised fine-tuning (SFT) with reinforcement learning (RL) have emerged as the key recipe for transforming large language models (LLMs) into robust reasoners. We argue that this combined success is driven by compositional generalization, which we formalize through a hierarchical latent selection model. In this framework, reasoning traces are generated by a cascade of discrete latent selection variables corresponding to reusable atomic modules, including both skills (local operations) and routing mechanisms (how intermediate information is selected, reused, and composed). Within this model, we theoretically show that SFT and RL play asymmetric, complementary roles: SFT supplies the raw module materials in compositional traces, and RL decomposes those traces to identify the latent atomic modules and enable compositional generalization. We design controlled experiments to validate this theory. Our results demonstrate that RL can extract atomic modules from compound traces supplied by SFT and recombine them to solve new configurations. Moreover, we find that training on compound traces yields stronger generalization than training on isolated atomic modules. Finally, we investigate the relationship between SFT and RL data and identify an effective protocol in which SFT ensures coverage of all atomic modules through compositional traces, while RL focuses on novel compositions outside the SFT support to drive exploration.

Problem

• compositional generalization
• reasoning traces
• reusable modules
• supervised fine-tuning
• reinforcement learning

Innovation

• compositional generalization
• reasoning traces
• atomic modules
• supervised fine-tuning
• reinforcement learning