Training Language Models to Cooperate with Inference-Time Controllers

2026-07-26

AI Summary

Current post-training approaches for large language models typically optimize for a single, fixed reasoning controller, limiting their ability to generalize across the diverse inference protocols encountered in real-world deployment—such as chain-of-thought, self-consistency, or debate-based reasoning. To address this limitation, this work proposes CALM, a framework that explicitly incorporates reasoning controllers into the post-training process. CALM employs multi-task reinforcement learning over interaction protocols defined by these controllers, modeling each controller as a composition of reusable, local reasoning modules. The approach introduces a turn-level Generalized Reward Policy Optimization (GRPO) objective to enable modular, controller-aware policy learning. Experiments demonstrate that CALM significantly outperforms conventional single-controller optimization methods in both unseen controller combinations and cross-controller transfer settings, substantially enhancing the model’s adaptability to diverse reasoning workflows.

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.

Problem

• inference-time controllers
• training-deployment mismatch
• controller generalization
• post-training
• reasoning workflows

Innovation

• Controller-Aware Training
• Inference-Time Controllers
• Multi-Task Reinforcement Learning
• Modular Reasoning
• GRPO