Language models learn complex reasoning better when trained on whole tasks
Compositional Reasoning in Language Models under Reinforcement Learning Post-Training
Artificial IntelligenceMachine Learning
Summary
Solving big problems often means breaking them into smaller skills and then combining those skills. This paper looks at how language models improve their problem-solving after extra training with rewards, especially for combining skills. The authors found it’s easier for the models to learn big combined tasks first and then do smaller parts, rather than the other way around. They explain why this happens and check if it holds true in different tricky settings and even some real-world examples.
What this means in practice
- •For machine learning engineers: Improve training strategies for language models to solve complex tasks by starting with composed skill training to enhance overall reasoning ability.
- •For software developers building ai assistants: Design AI tools that better generalize problem-solving by focusing training on full composed tasks rather than disjointed smaller skills.
Authors
Yu He, Yingxi Li, Yifei Wang, Ellen Vitercik
Abstract
Compositional reasoning is critical for real-world problem solving: since training data is necessarily limited, models must generalize by composing learned skills in new ways. While post-training methods such as reinforcement learning (RL) have substantially improved the reasoning abilities of language models (LMs), their effects on compositional reasoning remain less well understood. We propose a dependency-graph framework to formalize compositional reasoning, yielding three levels of compositionality with increasing complexity. Empirically, we instantiate this framework with data-structure tasks, which provide deterministic reward computation and clear compositional structure. We find a consistent decomposed-to-composed asymmetry: decomposed-skill training does not reliably transfer to composed tasks, whereas composed-task training transfers more readily back to decomposed tasks. We provide theoretical explanation for this asymmetry, and further evaluate compositional generalization under length extrapolation, structural distribution shift, and transfer to tasks requiring unseen skills. Finally, we present a pilot study on real-world tool-calling benchmarks, showing preliminary evidence that the decomposed-to-composed asymmetry can extend to practical settings.