TLDR: CoSMo-RL is a novel reinforcement learning framework designed to train Large Multimodal Reasoning Models (LMRMs) that are both highly capable and inherently safe. It integrates safety and general reasoning development into a single, stable pipeline, utilizing a two-stage training process, a robust optimization algorithm (CPGD), a multiobjective reward function, and multimodal jailbreak data augmentation. This approach leads to LMRMs that demonstrate improved safety awareness, reduced over-refusal, enhanced moral reasoning, and strong general reasoning performance, proving that safety and capability can co-evolve.
Large Multimodal Reasoning Models (LMRMs) are becoming increasingly common in real-world applications, but ensuring their safety and reliability remains a significant challenge. These advanced AI models, which can process both images and text, often struggle with safety issues, such as being tricked by combined visual and textual inputs to bypass safety measures, or exhibiting inconsistent behavior like over-refusing benign requests or complying with risky ones.
A new framework called CoSMo-RL aims to address these critical issues by training reasoning-oriented LMRMs under a comprehensive set of multimodal, multitask, and multiobjective signals. The core idea behind CoSMo-RL is to allow safety and capability to develop together in a stable training process, rather than having them compete against each other during the alignment phase.
Key Principles of CoSMo-RL
The CoSMo-RL framework is built upon four fundamental principles:
-
Strong General Reasoning Leads to Safe Behavior: Models with broad capabilities are better equipped to understand and navigate complex multimodal environments, enabling them to follow instructions accurately and anticipate risky situations.
-
Staged Safety Alignment: Safety measures need to be introduced and balanced over time. Early safety training can be overwritten by later, more complex task training, so a careful balance is crucial for long-lasting alignment.
-
Policy Stability is Essential: Uncontrolled updates during reinforcement learning can lead to unpredictable or undesirable behaviors, undermining both performance and safety. CoSMo-RL incorporates mechanisms to ensure stable policy updates.
-
Robustness Through Exposure: To resist real-world attacks, models must be trained on adversarial multimodal scenarios, not just evaluated on them.
To implement these principles, CoSMo-RL combines supervised pretraining with a two-stage reinforcement learning schedule. In the first stage, the model focuses on acquiring broad reasoning skills. The second stage then jointly optimizes for safety, helpfulness, and overall capability. Stability is maintained using an algorithm called Clipped Policy Gradient with Policy Drift (CPGD), and robustness is enhanced by training with multimodal ‘jailbreak’ data and preference-driven objectives.
A Unified Reward System
CoSMo-RL employs a unique multiobjective reward function that guides the learning process across various tasks. This reward is composed of four parts:
-
Visual-Focus: Encourages the model to pay attention to and correctly ground key visual elements in its reasoning.
-
Helpful: Promotes safe, accurate, and informative answers while discouraging risky or misleading outputs.
-
Format: Ensures that the model’s outputs follow a structured format, making its reasoning steps transparent.
-
Task-Aware: Covers safety, value, knowledge, and general reasoning dimensions, penalizing unsafe content, promoting factual reasoning, and ensuring completeness.
This unified design simplifies reward assignment, stabilizes training, and improves the model’s ability to generalize across different tasks.
Also Read:
- MaskGRPO: A Unified Reinforcement Learning Approach for Multimodal Discrete Diffusion Models
- AGENTRL: A Scalable Framework for Training Generalist LLM Agents
Experimental Success
Experiments with CoSMo-R1, a model trained using this framework, demonstrate significant improvements. CoSMo-R1 showed enhanced safety awareness, consistently performing well across various safety benchmarks and even surpassing some leading proprietary models. Crucially, it also reduced ‘over-refusal’—the tendency of models to unnecessarily reject benign inputs in the name of safety. The framework also improved the model’s moral reasoning and value alignment, achieving competitive performance without relying on massive model scale or proprietary data.
Furthermore, CoSMo-R1 maintained and often improved multimodal reasoning and instruction following across a diverse set of benchmarks, proving that safety enhancements do not have to come at the cost of general capability. The framework’s generalizability was also confirmed, showing consistent gains when applied to different model architectures.
The development of CoSMo-RL represents a significant step towards creating more trustworthy and capable LMRMs, where safety and general intelligence evolve together. For more details, you can refer to the original research paper.
