Training Language Models to Cooperate with Inference-Time Controllers
CALM trains language models via multi-task RL over structured controllers, enhancing generalization across diverse inference workflows.
Key Findings
Methodology
This work models inference controllers as modular, reusable components forming structured workflows. It formulates multi-controller training as a multi-task reinforcement learning problem, employing turn-level GRPO to optimize a shared policy. Modules are assigned roles (e.g., chain-of-thought, critic, verifier), with control flow dictating invocation order. The training incorporates module-level loss decomposition, gradient normalization, and adaptive weighting to balance contributions from different modules. Multiple variants, including uniform, single-role, and gradient-norm-based weighting, are explored. Experiments on mathematical reasoning tasks demonstrate that this approach enables the model to generalize to unseen controller compositions and shifts, outperforming single-controller fine-tuning baselines.
Key Results
- On GSM8K, CALM achieves a 5.2% accuracy improvement over single-controller fine-tuning on unseen controller combinations, and 3.8% higher accuracy in controller transfer tasks, indicating strong generalization.
- Module-level gradient balancing strategies, especially gradient normalization, significantly improve robustness across diverse controller structures, reducing training instability.
- The approach enhances multi-turn reasoning stability, with models maintaining high performance despite structural changes in controllers, validating the effectiveness of modular multi-task RL.
Significance
This research addresses the critical challenge of enabling language models to adapt to a broad spectrum of structured inference workflows. By training a single policy over a distribution of controllers, it overcomes the limitations of fixed-pattern fine-tuning, paving the way for versatile, robust AI reasoning systems. The modular, multi-task RL framework provides a scalable solution for future multi-controller environments, with implications for complex decision-making, multi-agent reasoning, and tool integration in AI systems.
Technical Contribution
The paper introduces a novel multi-controller training paradigm based on modular RL, employing turn-level GRPO with loss decomposition and gradient normalization. It conceptualizes controllers as compositions of reusable modules, enabling the model to learn generalized reasoning behaviors. The framework incorporates multiple training variants to balance module contributions, addressing gradient imbalance and stability issues. This approach extends multi-task RL to structured, heterogeneous control environments, offering a new methodology for training adaptable, multi-purpose language models.
Novelty
This is the first systematic formulation of multi-controller inference workflows as a multi-task reinforcement learning problem, emphasizing modularity and generalization. Unlike prior work focusing on single-controller fine-tuning, it introduces a module-level loss decomposition and adaptive gradient balancing, enabling models to generalize to unseen controller compositions. This innovation bridges the gap between structured inference pipelines and scalable training strategies, marking a significant advance in model robustness and flexibility.
Limitations
- Training stability remains challenging, especially with increasing interaction length and complex controller structures, requiring further algorithmic refinement.
- Model performance on highly novel or extreme controller configurations is limited, indicating room for improvement in robustness.
- Computational costs are high due to multi-turn RL sampling and gradient balancing, constraining large-scale deployment and real-time applications.
Future Work
Future directions include developing more stable and efficient training algorithms, integrating meta-learning for rapid adaptation, and expanding controller diversity. Exploring automatic controller search and optimization, as well as reducing computational overhead, will be key to deploying these models in real-world, resource-constrained environments. Additionally, extending the framework to multi-agent and multi-modal reasoning tasks offers promising avenues for research.
AI Executive Summary
Recent advances in large language models (LLMs) have demonstrated remarkable capabilities in reasoning, planning, and complex interaction tasks. However, their performance heavily depends on the inference workflows, which are often structured as multi-step, multi-module pipelines such as Chain-of-Thought, debate, or tool-augmented reasoning. Traditionally, models are fine-tuned on a single fixed interaction pattern, limiting their ability to generalize across diverse workflows. This mismatch hampers deployment flexibility and robustness in real-world applications.
To address this, the authors propose CALM (Controller-Aware Language Models), a novel training framework that explicitly incorporates multiple inference controllers into the training process. They model controllers as compositions of reusable modules—each representing a reasoning or verification step—and formulate the training as a multi-task reinforcement learning problem. The core algorithm employs turn-level Group Relative Policy Optimization (GRPO), which samples multiple trajectories under different controllers, optimizing a shared policy that can adapt to unseen controller configurations.
A key innovation is the module-level loss decomposition, combined with gradient normalization and adaptive weighting strategies, which balance contributions from different modules and mitigate training instability. The framework supports various training variants, including uniform, single-role, and gradient-norm-based approaches, allowing systematic analysis of their impact on generalization.
Experimental results on the GSM8K dataset show that CALM significantly outperforms traditional fine-tuning, achieving a 5.2% accuracy boost on unseen controller combinations and demonstrating robust transfer in controller shift scenarios. These findings validate the effectiveness of multi-controller training in enhancing model flexibility and robustness.
Overall, CALM represents a substantial step toward versatile, multi-purpose reasoning systems capable of adapting to a broad spectrum of inference workflows. Its modular, reinforcement learning-based approach opens new avenues for scalable, generalizable AI reasoning, with potential applications in automated reasoning, multi-agent systems, and interactive AI tools. Future work will focus on improving training stability, reducing computational costs, and expanding controller diversity to further advance the field.
Deep Dive
Abstract
Large language model (LLM) performance increasingly depends not only on the base model, but also on the inference-time controller used to organize reasoning. Existing post-training methods, however, typically optimize for a single fixed interaction pattern, despite real deployments relying on diverse controllers such as Chain-of-Thought, self-consistency, debate, planning, and verification pipelines. This creates a training--deployment mismatch and limits transfer to new workflows. We introduce CALM (Controller-Aware Language Models), a post-training framework that explicitly places controllers in the training loop. We formulate controller-aware post-training as multi-task reinforcement learning over controller-induced interaction protocols, where controllers are compositions of reusable local reasoning modules. This structure also induces a module-level decomposition of mixed-controller training under a turn-level GRPO objective, enabling a systematic study of controller and module-aware training strategies. We evaluate CALM on held-out controller compositions and broader controller shifts, showing that controller-aware post-training improves generalization across inference-time workflows beyond single-controller optimization.