From Reasoning Traces to Reusable Modules: Understanding Compositional Generalization in Language Model Reasoning
For researchers and practitioners building reasoning LLMs, this work provides a theoretical framework and practical protocol for combining SFT and RL to achieve compositional generalization.
The paper proposes that the success of SFT+RL post-training for LLM reasoning is due to compositional generalization, formalized via a hierarchical latent selection model. Theoretically and experimentally, SFT supplies atomic modules in traces while RL decomposes and recombines them for generalization, with RL on compound traces outperforming training on isolated modules.
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.