Boosting Multi-Domain Reasoning of LLMs via Curvature-Guided Policy Optimization
Authors: Xize Liang
Lin Yang
Jie Wang
Rui Liu
Yang Lu
Jinliang Zeng
Hanzhu Chen
Dong Li
Jianye HAO
Venue: ICLR 2026 Poster (2026)
Keywords: large language models
policy optimization
multi-domain reasoning
Primary area: foundation or frontier models, including LLMs
TL;DR: We propose CGPO, a scalable, curvature-guided RL method that leverages cross-domain gradient interactions to enhance multi-domain reasoning in LLMs, achieving faster and more consistent improvements across diverse tasks.

Abstract: Multi-domain reinforcement learning (RL) for large language models (LLMs) involves highly intricate reward surfaces, posing significant challenges in finding parameters that excel across all domains. Recent empirical studies have further highlighted conflicts among domains, where gains in one capability often come at the expense of another. However, approaches to mitigate such conflicts and enhance multi-domain reasoning remain largely underexplored. To address this challenge, we propose **C**urvature-**G**uided **P**olicy **O**ptimization (**CGPO**), a principled and scalable training framework to advance the multi-domain reasoning of LLMs. Inspired by Newton's method, CGPO exploits the geometric structure in the reward surface, while sidestepping the prohibitive cost of Hessian computation. At each update, CGPO processes domains in random order, preconditioning their gradients with curvature information from other domains to foster richer cross-domain interactions. This mechanism further promotes implicit gradient alignment by maximizing inter-domain inner products in expectation, steering the parameters toward regions that jointly enhance multi-domain performance. Extensive experiments on a mixed dataset covering math, coding, science, and creative writing, evaluated across seven widely-used benchmarks, show that CGPO significantly outperforms all baselines in terms of faster reward improvement and stronger multi-domain capability.

Published as a conference paper at ICLR 2026
BOOSTING MULTI-DOMAIN REASONING OF LLMs VIA CURVATURE-GUIDED POLICY OPTIMIZATION
Xize Liang!",, Lin Yang" Jie Wang Rui Liu', Yang Lu', Jinliang Zeng', Hanzhu Chen', Dong Li?, Jianye Hao
'MoE Key Laboratory of grain-insspired Intelligent Perception and Cognition,
University of Science and Technology of China
HHuawei Technologies Co., Ltd.
'College of Intelligence and Computing, Tianjin University xi zeel angêma 1. ustc. edu.
ABSTRACT
Multi-domain reinforcement learning (RL) for large language models (LLMs) involves highly intricate reward surfaces, posing significant challenges finding parameters that excel across all domains. Recent empirical studies have further highlighted conflicts among domains, where gains in one capability often come at the expense of another. However, approaches mitigate such conflicts and enhance multi-domain reasoning remain largely underexplored. To address this challenge, we propose Curvature- survature-Guided Guided Policy Optimization (CGPO), a principled and scal- able training framework to advance the multi-domain reasoning of LLMs. Inspired by Newton's method, CGPO exploits the geometric structure in the reward surface, while sidestepping the prohibitive cost of Hessian computation. At each update, CGPO processes domains in random order, preconditioning their gradients with curvature information from other domains to foster richer cross-domain interac- tions. This mechanism further promotes implicit gradient alignment by maximizing inter-domain inner products in expectation, steering the parameters toward regions that jointly enhance multi-domain performance. Extensive experiments on a mixed dataset covering math, coding, science, and creative writing, evaluated across seven widely- dused benchmarks, show that CGPO significantly outperforms all baselines in terms of faster reward improvement and stronger multi-domain capability.
INTRODUCTION
Large language models (LLMs) have recently achieved remarkable progress in complex reasoning, including mathematical problem solving (Yang et al. 2024 Yu et al. 2025a), code generation (Ye
al. 2025 Zeng et al., 2025). circuit generation (Bai et (0266) optimization modeling (Liu
et al. 2026), and creative writing (Fein et al. 2025 Carrera et al. 2025). A key driver behind these advances reinforcement learning (RL), particularly policy optimization methods such as PPO
(Schulman et al. 2017) and GRPO (Shao et al. 2024). While earlier work primarily focused on applying RI. within single domains (Hu et al. 2025) Yu et al. (025a)), more recent studies have moved toward multi-domain reasoning, constructing diverse datasets (Cheng et al.) (025), training general reward models (Ma et al. 2025), and empirically examining interactions among different reasoning capabilities (Li et al. 2025b Cheng et al. 2025).
Despite these advances, multi-domain RL for LLMs still confronts significant challenges. The coexistence of diverse data distributions and reward signals produces highly complex reward surfaces, making it difficult to find parameters that excel across all domains simultaneously (Vithayathil Vargh & Mahmoud, 2020) Crawshaw 2020). Recent studies further show that, although multi-domain RI. can yield overall benefits, is often hindered by cross-domain conflicts, where gains in one capa- bility are accompanied by losses in another (Cheng al. 2025 Li al. 2025b). These difficulties
"Work conducted during the research internship of Xize Liang (xizeliang www.l..stc.ed..c.o at Huawei.
1 Equal contribution.
Corresponding author. Email: www.annnxxuttcctucccc.m

---

Published as a conference paper at ICLR 2026
further compounded by the nature of RL training: on one hand, online sampling (i.e., rollouts) introduces unpredictable interactions among domain-specific samples; on the other hand, generating rollouts is computationally expensive, and much of this effort is wasted when cross-domain conflicts cancel out the contributions. These considerations make it crucial to develop RI. frameworks that fully exploit mixed datasets to enhance L.MMs' reasoning across diverse domains.
Cross-domain conflicts often manifest as gradient conflicts (Chen et al. 2025b Liang et al. 2025). yet widely-used approaches for mitigating them face notable limitations in the context of RL for LLMs. Most existing methods intervene during gradient aggregation once conflicts occur, aiming to balance updates across domains. On the one hand, they do not leverage the underlying geometry of the reward surface or loss landscape (Liu et al.] [2023] [eener & Koltun (0188). On noisy, rollout- based gradients, such purely reactive strategies tend to amplify update variance and degrade both stability and performance. On the other hand, many techniques require storing and manipulating all domain gradients simultaneously on the device (Yu et al., 2020 [iu et al. 2024, 202T). This incurs substantial memory overhead that grows rapidly with the number of domains and can even result in out-of-memory failures, severely limiting the scalability of multi-domain RL for LLMs Alternatively, recent work suggests that second-order methods such as Newton's 's method and its approximation SOAP Vaas et al., (2255) can mitigate gradient conflicts in PINNs (Wang et al., (025). but their reliance on Hessian computations renders them infeasible for the high-dimensional, rollout-heavy setting of RL for LLMss These limitations compellingly motivate the following question: How to mitigate cross-domain confticts in a manner that is both consistent with the nature of RL and efficient at scale, thereby enhancing the multi-domain reasoning capabilities of LLMs?
In this paper, we propose CGPO, a principled and scalable policy optimization framework, to enhance multi-domain reasoning for LLM CGPO draws inspiration from Newton's 's method, while incorporating a design specifically adapted to the distinct challenges of multi-domain RL for LLMs. Newton's 's method exploits the geometric structure of the loss landscape (i.e., the Hessian matrix) to precondition gradients, correcting directional deviations induced by anisotropy and facilitating efficient convergence. To retain these benefits while circumventing the computational burden of full Hessian computation, we adapt the preconditioning step into a lightweight mechanism tailored for efficient RL training of LLMs. Specifically, at each parameter update, domains are processed in random order, with each domain's s gradient modulated by curvature information from others, thereby inducing rich cross-domain interactions. Another appealing feature of this mechanism is that it implicitly aligns domain gradients by maximizing their inner products in expectation, guiding the parameters toward regions of high cross-domain consistency. We validate CGPO on a diverse dataset of 20k samples spanning mathematical reasoning, code generation, scientific --..5----Innnte and Dowen2.5-7B-Instreee evaluated across
QA, and creative writing using seven benchmarks. Our results demonstrate that CGPO consistently outperforms a broad spectrum of maselines-includingg curriculum learning strategies, gradient balancing techniques, and joint learning- achieving faster reward gains and markedly stronger multi-domain reasoning capabilities.
2 PRELIMINARIES
2.1 MULTI-DOMAIN LANGUAGE MODELING AS REINFORCEMENT LEARNING
An LLM no (with parameters 0) defines conditional probability distribution over output responses
y = "pt....,yr). given a query x D. represented as T x) = II ( x, y1t--))). To align LLMs with desired behaviors, recent work formulates language generation as a reinforcement learning (RI.) problem. The model acts as a policy that interacts with an environment by generating responses y to queries and each response receives a reward R(x, y) E R that reflects its quality.
In many real-world applications, LLMs are expected to perform well across multiple domains, each corresponding to distinct type of query or task. Formally, let there be K domains with query distributions {K Each domain & defines its own reward function R& (:::,, reflecting task- specific quality criteria. Assuming equal importance for all domains, the multi-domain training objective is to maximize the average expected reward (we abbreviate y ~ Tt(- x) as y ~ π J(0)- 1 Ex1 [f) = K Σx Ex~DD yymm R& (x. y)]. We provide discussion on extending K this formulation to non-uniform domain importance in Appendix[D.2
Additional discussion on the applicability of CGPO LLM pre-training provided in Appendix D.1
2

---

Published as a conference paper at ICLR 2026
2.2 POLICY OPTIMIZATION ALGORIHTM
The multi-domain formulation in Section 2.1 reduces to the standard RL objective when expressed with generic query distribution 'D and reward function R. i.e., I(0) Ex-D,y- [[(x, y)]-
Directly optimizing I(0) is challenging due to the discrete, variable-length output space and the dependency of the distribution Th on the parameters 0. Instead, the policy gradient the- orem (Sutton et al. 1998) provides an unbiased estimator for the gradient, i... VaJ(O) Ex-DD Vo log To (y x)A(x.XX y)], where y(x.y) R(x,y) b(x) denotes the advantage of response y over a baseline b(x). In practice, the true advantage function is unknown and must be estimated from rollouts. This is typically done by training value function (x( to approximate
the expected reward, and then computing an estimated advantage Á(x,y) y) = R(x, y) - (x)).. By
combining this estimator with importance sampling using rollouts from old policy Taaa one can
define surrogate objective L00 loid, D) y.yeeed To.ta (y/x)
moly|x) A(x.y)
While the theoretical surrogate objective using the true advantage A has a gradient that coincides exactly with V (0) at 0 Oota. practical objectives using the estimated advantage À serve as first-order approximation. This approximation is reliable as long the updated policy my re- mains close to To.i Building on this, Proximal Policy Optimization (PPO) (Schulman et al. 2017) ensures stable policy updates by maximizing a clipped surrogate objective LPPo Bold, D)
Ex--Dyy--ooa
min
Teat [y]x]
(ylx)) A(x,y), clipp
clipp 1+
Koai (y|x)
xo(y)x))
A(x.y))) where is small hyper-
parameter and clip cli Yiow
clip(:; Now, Thigh) is the clipping function.
However, the reliance of PPO on separately trained critic model to estimate b(x) introduces substantial memory and computational overhead. To address this, recent critic-free methods rep- resented by GRPO (Shao et al., (024) have emerged. GRPO estimates the baseline directly from a group of sampled responses. Specifically, it samples G responses {
Ji=1 for each query x, Jiz1, and then computes a normalized advantage for each response: A [r(9) mean ((r())) [-=1)] /std (fr6))))) The overall GRPO surrogate objective is
obtains their rewards {
LGRPO (0; Oord, D) (ExDD((r()) (1()))
min (o( x)
A0) clipi Tata (y(i)
1tenihh
no (y((t)
x)
A
(6) -D (TO (TTo|| (reet)
(1)
Teaaa (y(1) x)
x)
where Elow, Fhigh, and B are hyperparameters, Tee is a reference policy (typically the initial model), and Di (Tell (Ter) is a sample-based KL divergence penalty. In this work, we adopt GRPO as our base policy gradient algorithm due to its efficiency and scalability.
Surrogate Objectives as Faithful Gradient Approximators. While the policy gradient theorem provides unbiased gradient for the true advantage A, practical algorithms rely on estimated advantages Å, which introduce variance. Surrogate objectives like PPO and GRPO are designed to stabilize these gradients: PPO uses clipping to enforce a trust region, making V ygLL (0; fold. D) a reliable approximation of ((). while GRPO's combination of clipping and KI. regularization similarly produces a stable gradient VoLGRPO (0; foid, D) that approximates the KL-rregularized objective Ve (J) - 3'L (Tel [ret)).
2.3 NEWTON'S METHOD FOR GRADIENT PRECONDITIONING
Newton's 's method is a classical second-order optimization algorithm that leverages the curvature of the objective to accelerate convergence. Given a twice-differentiabll loss T.))). the Newton update
is 0.+1 0, H(),) (((0)))-((8)) (0.), where g(O,) = (eL(0,) and H(O,) (0)( is the Hessian. By preconditioning the gradient with local curvature, Newton's 's method corrects for anisotropy, producing more direct steps toward optimum. It is particularly effective in complex, conflicting landscapes;
e.g., Wang et al. (2025) shows that Newton's 's method and its approximate variant SOAP (Vyas et al. 2025 mitigate gradient conflicts in PINNs and accelerate convergence.
However, directly applying Newton's 's method to RL for LLMs is impractical: the Hessian is high- dimensional and costly compute invert, and rollout-based gradients noisy. Still, the principle
3

---

Published as a conference paper at ICLR 2026
Quir
Response
Rescadi
Advantagn
Interpalation
A
α a(OK 2
R
d
61
K
A
(o)
A
2
(2)
L8
A
Single-Domsin Gradient
622
onf
K
k-1
H. (o0g)
K
(
K
Ox
Croas-Domain Inseraction
Figure 1: Illustration of CGPO (one update step). After generating responses, computing rewards, and estimating advantages for each domain, CGPO randomly permutes the domain order and applies updates sequentially, followed by interpolation with the original model. The parameter change A can be approximately decomposed into a single-domain gradient term-capturing per-domain learning- and cross-domain interaction term that facilitates transfer across domains. Note that CGPO introduces only negligible additional computation overhead (see Section 4.3 for details).
of leveraging curvature to guide updates provides a valuable foundation for designing optimization strategies that handle conflicting gradients and complex surfaces, we explore Section 3
CURVATURE-GUIDER POLICY OPTIMIZATION
Building on the preliminaries, seek to leverage the insight that Newton's method couples gradients with curvature information- property that can be particularly valuable LLMs, where interactions between domains often complex and interdependent. Rather than directly approximating the Newton update, which would be computationally prohibitive in our setting, distill its essential idea into lightweight mechanism that induces cross-domain gradient-curvature interactions via sequential task updates. Our method unfolds in three parts: Section motivates motivvveee the design by analyzing the structure the Newton update, Section 3.2 presents a simple perturbation- based procedure to capture the desired interactions, and Section 3.3 integrates these components into a practical algorithm, i.e., our proposed CGPO. An overview of CGPO illustrated in Figure
RL for
3.1 MOTIVATION: WHY HESSIAN-GRADIEN INTERACTIONS MATTER
The starting point of CGPO is an informal observation about Newton's 's method. Although exact second-order updates are infeasible large-scale RL for LLMs, the Newton term Hg (omitting 0,) couples gradient and curvature, suggesting that such interactions may help reconcile conflicting gradients in multi-domain learning. To illustrate, consider heuristic expansion: H g (I- H))-'g g (I+ (I (I + OOIOOIIIIIII )))) 2g Hg + + H)?g), ?g) where the approximations informal and serve to reveal the structure rather than provide rigorous formula. In the multi-domain setting, where g Ek=
1= gt and H = Σk Hę, the product -g then contains cross-domain terms H (i * j), which the curvature domain j modulates the gradient of domain
These interactions effectively transmit curvature signals across tasks, amplifying, dampening, or redirecting updates capabilitie absent in first-order methods. This motivates our key design principle: instead of computing Hessians explicitly, we seek tractable mechanisms that induce such cross-domain interactions better align multi-domain optimization.
3.2 APPROXIMATE -DOMAN INTERACTIONS VIA SEQUENTIAL UPDATES
Given the motivation above, the question is how to induce Hessian Hessian-gradient interactions without explicitly computing Hessians. Our key idea is to approximate them by observing how the gradient of one domain changes after parameter updates from another.
4

---

Φx. the gradient of Published as a conference paper at ICLR 2026
Algorithm 1 CGPO (one epoch illustration)
1: Input: min reward functions { {RR 1K 1 datasets {D
1K
2: Hyperparameter: number of steps T, M, learning rate 7, mixing coefficient a
3: Initialization: Tref + Romme Bowe &- 1iit
4: for t 1,..., do
5: Taaa <- T0ncm
6: Sample a batch D(o,. = (x i=1
from DE for k
(0.k)
7: Generate responses (1(0.k]
(1.)
(.1x for lis Doyah 1 ((3)) G (1),b j=1
Compute rewards {
(1.3) and advantages (0).k j=1
8:
for < (P 1 k K
9:
for m ..... do
10:
Sample mini-batch ( from D for 1 k < K
11:
Let (1)...).... (KK) denote a random permutation of 1...., K
12:
② Onew for k 1, K do
13:
Update parameters by maximizing Eq. ! with D ((.,). and associated responses:
ID (D(L.).
ØF = ōb-1 - 1) Ex= (
EGRPO (0x-1; fold D ((t.)),o(k))
15:
0 fne t- Φu + alØx 20)
Output: To Tonow
Consider two domains and j. Let domain updates the parameters from of1 to Denoting the Hessian of domain j at ffie by H2, (8p0e) the gradient of domain j then shifts as
o(0)
gi (oost)
gj
( pre
Hy
(0()) pre
(4) post
g(o) pre
nH
(glt) pro
gi
( pre
which corresponds the cross-domain product H,g; This approximation derived from a first- order Taylor expansion and policy gradient ascent (see Appendix [B. for the detailed derivation). Thus, sequential updates naturally generate the desired interaction term. Further, extend beyond two domains, we randomize the order of domains at each iteration. Over time, this exposes every pair of domains to such interactions, allowing curvature information to propagate across domains. Intuitively, each domain feels the curvature of others: one nudges the parameters, another responds, producing coordinated updates that help reconcile conflicting objectives.
3.3 FULL ALGORITHM: RANDOMIZED CROSS-TASK INTERACTIONS
Building on the insights above, we now introduce CGPO, a principled algorithm for multi-domain policy optimization, illustrated in Figure T with pseudocode Alg. T At each training step, sample batches from all domains and generate multiple candidate responses under the current policy (Lines 6-7). These responses are evaluated by domain-specific reward functions to obtain rewards and advantage estimates (Line 8). We then repeatedly draw mini-batches (1 (Lines 9-10) and perform randomized sequential update: domains are visited according to random permutation (Lines 11-13), and at each step the parameters are updated with respect to domain, conditioned on perturbations induced by previously visited domains (Line 14). Finally, the updated parameters are interpolated with the original ones using mixing coefficient (Line 15), stabilizing training by balancing curvature- curvatureiifformed exploration with retention of the base policy.
To understand how sequential updates induce cross dros-d Hessian gradient interactions. consider Lines 11-15. Let the domain order be (11)..... o (K). and denote the loss, gradient, and Hessian of domain parameter e by Ta (6), gi: and H (Ø).. With do Φ1

---

Published as a conference paper at ICLR 2026
domain or (k) at Φk-1 can be expanded (see Appendix B.2 as
7 Do
(o) (0x- (0)))
(oo)) O
(3)
E==1 (Da(s)
For simplicity, assume uniform batch sizes |D(())/
k
(o((s) 1/K, then
k-1
(x) = (x) - K
Σ
H. Ho (co)
())))))))))))))))))))) (o) + (11²).
(4)
1=1
Aggregating over k, the overall parameter change after one sequential pass is (see Appendix B.3
K k-1
a(ók - 0o)
am K
gh (6o)
az K2
ΣΣ
Ho( (00 (o + 0
The first term is the aggregated gradient; the second term contains cross-domain Hessienngradient products. Importantly, the expression above describes the update for fixed permutation Because our algorithm re-samples a independently at every iteration, the quantity relevant for understanding CGPO's behavior is the expectation over the random permutation a. Taking expectation over makes every ordered pair (d..j) appear with equal probability; symmetrizing their contributions then yields (0o))) ())))) + (oo) (do) = Dég (Øo) T (őo)) (please see Appendix B.4 for
0
details). This shows that the update encourages alignment of domain gradients. For an analysis of why joint learning does not induce the same cross doos--d effect, please see Appendix D.3
After illustrating how the parameter change encodes both aggregated gradients and cross-domain interactions, is helpful to clarify the role of the final interpolation step. The vector ΦK- -0 provides geometry-informed update direction enriched by these interactions. The mixing coefficient a then controls how far we move along this direction (for ablations, see Section 4.3 sufficiently large enables the method to benefit from cevvature-iffomed coordination across domains, whereas an excessively large value may push the update outside the locally smooth region where gradient-basee approximations remain reliable, potentially destabilizing chiiiigg-nallgee large learning rate in standard optimization. Conversely, setting a too small would under-utilize the information encoded in Φx do and collapse the update to near-identity update, losing the benefits introduced sequential interactions. The interpolation therefore functions as a principled mechanism that balances stability and effective of cross-domain geometric information.
to taking an overly
Crucially, this analysis is not restricted to surrogate losses Ls: as argued in Section 2.2 GRPO surrogates provide faithful approximations of the true policy gradients within their trust regions. Thus, the induced interactions improve alignment not only among surrogate gradients but also among the true policy gradients (o)x In effect, randomized sequential updates encourage cooperation across domains by introducing curvature-gradient couplings that steer optimization toward coordinated improvements on the full multi-domain objective Σx (0).
Discussion, We highlight two clarifications to better situate our approach.
Sequential updates is common technique across different learning paradigms. For example, in meta-learning, Reptile (Nichol et al., 2018) adopts sequential updates to learn an initial model for rapid adaptation to new tasks, while in federated learning, methods such FedAvg (McMahan et al., (017) aggregate sequential client updates to improve global optimization. However, these precedents do not diminish the novelty of our contributions. First, our sequential update originates from our observation of Newton's 's method and its capability to navigate complex landscapes, where interactions naturally emerge across domains. Second, we adapt this mechanism to the multi-domain RL for LLMs setting, where domain-specitee rewards and surrogate policy gradients pose unique challenges absent in meta-learning or federated learning. Finally, we integrate randomized ordering, surrogate faithfulness (via GRPO), and stabilization through interpolation into a unified algorithm tailored for large-scale scale RLHF These innovations collectively distinguish CGPO as novel and practical solution for multi-doman policy optimization.
inherent crrvature-gradicn
A natural concern is that multiple updates per step could inflate the effective learning rate. To avoid this, scale each gradient proportionally to its mini-batcl size and normalize by the total across domains. This ensures that the overall update magnitude is consistent with that of using a single aggregated batch, thereby preserving comparability with standard mini-batch optimization.

---

Published as a conference paper at ICLR 2026
4 EXPERIMENTS
4.1 EXPERIMENTAL SETTINGS
Tasks and Datasets. We focus on enhancing the LLMs' overall capabilities across four domains- mathematical reasoning, code generation, scientific QA, and creative writing. These domains not only represent core areas of current research interest but also span four distinct forms of reward feedback, thereby ensuring both comprehensiveness and diversity. For mathematics, code, and science, we construct subsets from the Guru dataset (Cheng et al.) (025) with attention to dataset size and sample difficulty (as Guru poses non-trivial challenges for 7B-scale models): the math subset contains 6,250 samples, consisting of the 5,000 casiest problems (ranked by the pass rate of Swwen..5-BBIIstrruee and 1,250 more challenging ones; the code subset totals 4,740 samples, comprising all 3,791 problems with a Qwwe2.5-7BB-Isttree additional 949 randomly sampled from the remainder, ensuring a roughly 4:1 ratio between easier and harder samples; and the scientific QA subset includes the entire STEM split of Guru, with 3,591 samples. For creative writing, we randomly sample 2,000 samples each from the three most popular datasets available on Huggingface (LitBench (Fein et al., 2025). ) Creative_Writing--
pass rate of at least 25% plus an
(Nitral-Al, 2024), and samples. For details of the datasets, please see Appendix C.1
(kevinshin (025 yyeldingg a dataset of 6,000
Baselines. We compare our CGPO with several representative baselines. For vanilla strategies, we include joint learning, which directly trains on multi-domati dataset without any special strategies. For curriculum learning (CL), following the taxonomy in (Soviany et al.] (022)) we include Omni- Thinker (Li et al. 2025a), a progressive CL method, and self-pacedd CL. which schedules training from easier to harder examples based on task difficulty (measured by pass rate). For gradient balancing, we include FAMO (Liu et al. 2023), categorized in (Chen et al.; 20256) representative approach for balancing gradient magnitudes across domains. We also attempted to implement gradient manipulation methods such as PCGrad (Yu et al. 2020), but these require simultancously storing and operating on multiple per-domain gradients on devices, which leads to in the RL for LLM setting. For more details of baselines, please refer to Appendix C.2
(OOM) issues
Training Details. We train Qwen2.5-3B Instruct and Owen2.5-7B-Iie dataset using the verl framework (Sheng et al. 2025). For the implementation of multi-domain training in terms of data processing and reward design, we follow the codebases of (Cheng et al. 2025) and (Ma et al. 2025). For math, we adopt rule-based rewards; for coding, we evaluate models' outputs using unit test cases based on SandboxFusion for scientific QA, we use a 1.5B Verifer (Ma et 2025)) to assess the consistency between model outputs and groundtruth answers; and for creative writing, we compare model responses with reference answers using -wen2.5-72B-Inss Instruct. Besides, we require the model to enclose its reasoning process within <think> tttiithiikkk nn> tags and penalize responses that violate this format requirement, along with domain-specific constraints. Details of the reward functions are provided in Appendix C.3 We use a learning rate of x 10 6 a prompt batch size of 128, a mini-batch size of 64. a group size of 8, rollout temperature of 1.0, Elow = 0.2, Ehigh = 0.28, and 3 = 0.001 for CGPO and all baselines. Our code is available at [nt ttttps:/ / /tt:://ittti heg...hubb.... com MIMMMMAAAAA... USTC/
on the multi-
Siredance-Seed-Foundee
et al. 2025);
For more details of hyperparameters, please see Appendix C.4
We evaluate our models on seven widely-used benchmarks: MATH500 (Hendrycks) et al. 2021), AMC 2023 (MAA, (023). HumanEval (Chen et al. 2021), MBPP (Austin et al. [0221) -diamond (Rein et al. 2023), SuperGPQA (Team et al. 2025). and Writing (Wu et al.
Evaluation.
2025). ). To ensure consistent scaling across benchmarks, the scores on WritingBench are multiplied by 10. We use vLLM (Kwon et al. 2023) for efficient inference, generating 4 responses per query with temperature of 0.6 and top-p sampling of 0.95. Further details can be found in Appendix C.5
4.2 MAIN RESULTS
CGPO boosts the multi-domain reasoning of LLMs. Table presents the results across different methods. From the table we make the following observations: CGPO achieves the highest average performance for both model scales (3B and 7B), ranking either first second in most individual domains. This demonstrates its effectiveness in enhancing multi-domain reasoning capabilities of LLMs. (2) For smaller models (3B), CGPO consistently outperforms other baselines generation and creative writing, while maintaining competitive performance on math and scientific QA. FAMO and Omni-Thinker also provide gains over joint learning, particularly in code generation
code
7

---

Published as a conference paper at ICLR 2026
Table 1: Performance of models (Owen2.5-3B-Instruv and Qwen2.5-7B-Instur trained on the multi-domain dataset with different methods, evaluated on multiple benchmarks. The bold font indicates the best result and an underline indicates the second-best result.
Math
Code Generation
Scientific QA
Creative Writing
Methods
AVG
MATH500
AMC
HvmanEval
MBPP GPQA-diamond SuperGPQA
WritingBench
Qwen2.5-3B--nssee
Joint Leaminn
64.50
39.38
72.39
59.40
24.87
24.12
58.61
49.04
Omni-Thinker
65.65
41.50
71.55
58.80
21.34
26.75
57.90
49.13
Self-paced CL
65.30
38.75
70.12
58.80
24.37
24.72
57.82
48.55
FAMO
63.80
39.12
72.48
59.20
23.47
26.51
58.46
49.01
CGPO
64.20
39.71
74.29
60.80
24.37
26.63
63.04
50.42
Qwen2.5-777 2--22 In-ttuutt
Joint Learning
76.00
56.25
79.88
68.60
19.70
32.75
63.15
56.62
Omni--Thinker
75.10
53.75
82,93
68.60
23.86
30.63
62.35
56.75
Self-paced CL
74.70
51.88
82,93
68.00
21.72
30.25
63.68
$6.17
FAMO
75.65
55,63
82,54
68.80
23.07
31.49
63.62
57.26
CGPO
75.55
59.38
84.15
72.00
26.77
32.75
66.52
59.59
Code Generation (3B)
Creative Writing (38)
Math (38)
Scientific QA (38)
0.2
jent
CCPO
0.1
0.6
Jent
CCPO
-0.3
jort
CCPO
CGPO
Ja
Reward
0.2
-1.2
0.675
-0.2
2.050
0.1
-0.2
-13
5.000
250
152
Step
Step
Step
Generation
Creative Writing (78)
Scientific
jeint COPO
CGPO
CGPO
CGPO
Reward
0.2
0.-
6.
0.1
-0.3
0.0
30
75
100
125
250
23
30
75
102
150
25
30
75
150
:25
25
20
75
120
155
Step
Step
Step
Step
Figure 2: Training reward curves for -wwen.5533--Inttii and D0wen2.5-7B-Insteeen on four domains (code, creative writing, math, and scientific QA), comparing CGPO and joint learning.
and scientific QA, but they lag behind CGPO in creative writing. Self-paced CL remains the weakest overall, likely due to imbalanced domain difficulty and insufficient coverage of informative responses different training stages. (3) For larger models (7B), CGPO achieves clear improvements across nearly all domains, with the largest gains on code generation and creative writing. highlighting that its benefits scale with model capacity. Notably, FAMO shows competitive results, especially in math and creative writing, confirming that gradient balancing can help. but it still falls short of CGPO in aggregating multi-domain knowledge effectively. These results collectively indicate that curriculum learning and gradient weighting methods can provide partial improvements, but their reliance on task difficulty, loss, or gradient magnitude alone is insufficient. In contrast, CGPO leverages geometric information via randomized sequential updates and interpolation, enabling coordinated multi-domaii optimization and consistent performance gains across mathematical reasoning, code generation, scientific QA, and open-ended creative tasks.
CGPO achieves faster reward improvement across all domains. Figure 2resents the training reward curves of Qwen2.5-3BB --3--uInsttt and Qwen2.5-7B. -B--truutttt on the four domains, with all curves smoothed using EMA to clearly reveal trends. For both model sizes, the curves of CGPO consistently remain above those of joint learning. The advantage is particularly pronounced in code generation and creative writing, while in math and scientific QA the improvement is evident but less striking. Notably, compared with the other three domains, creative writing more subjective, requiring the model to generate diverse and creative outputs rather than strictly structured or precise answers; this makes potential confficts with the other domains the largest. The substantial advantage of CGPO
8

---

Published as a conference paper at ICLR 2026
Table Ablation study on domain order randomization in CGPO with Qwen2.5-7B-Instruee The bold font indicates the better result.
Math
Code Generation
Scientific QA
Creative Writing AVG
Methods
MATH500
AMC
HumanEval
MBPP GPQA-diamond
SuperGPQA
WritingBench
CGPOns
77.20
56.88
83.54
69.60
23.08
31.75
67.30
58.48
CGPO
75.55
59.38
84.15
72.00
26.77
32.75
66.52
59.59
Table 4: Ablation study on the effect of the mixing coefficient a in CGPO with Qwen2.5-7B-
Instruct. The bold font indicates the best result and underline indicates the second-best result.
Math
Code Generation
Scientifie QA
Creative Writing
AVG
MATH500
AMC
HumanEval
MBPP
GPQA-diamond SuperGPQA
WritingBench
0.9
75.85
55.88
84.15
71.20
21.72
32.25
66.01
58.15 15
1.2
75.55
59.38
84.15
72.00
26.77
32.75
66.52
59.59
1.5
75.55
55.25
81.10
69.20
23.36
35.37
66.47
58.04
in the reward curve for creative writing compared to joint learning provides strong evidence that CGPO effectively mitigates cross-domain conflicts. We also observe considerable differences in initial reward levels across domains. Taking Qwen2.5-7B. Instrust as an example, creative writing and scientific QA start near -0.4 and O, respectively, reflecting largely incorrect outputs, whereas math and especially coding begin from higher baselines (coding around 0.1). This indicates that the models enter RI. training with uneven domain-specific capabilities. Importantly, CGPO delivers varing degrees of acceleration even for domains with comparable starting points, suggesting that factors such as dataset difficulty or reward function design may influence the speedup. Investigating the underlying causes of these differences is left for future work.
4.3 ANALYSIS AND ABLATIONS
CGPO introduces only negligible additional com- putation overhead. In multi-domain RL for LLMs, the dominant computational bottleneck typically lies in generating responses and computing rewards- particularly domains such as coding and creative writing rather than in the forward and backward passes of the model itself. Against this backdrop, the additional operations introduced by CGPO are minimal. The sequential updates across domains essentially equivalent splitting mini-batch into smaller chunks and processing them sequentially, which incurs almost the same computational cost standard mini-batch training. Furthermore, the final interpolation with the mixing coefficient a amounts to single vector operation, which is computationally negligible. Taken together, these factors ensure that the overall overhead of CGPO practically insignificant, and the total training cost remains nearly identical to that of joint learning. As shown in Table [ the per-step wall-clock time under CGPO is only slightly higher than joint learning, confirming that our method adds no meaningful overhead in practice. For timing experiments on 32B and 72B models, please see Appendix E.1
Table 2: Computation cost comparison be- tween joint learning and CGPO (1 epoch). Note that the units of total time and per-step time are different (hours vs. minutes).
Methods
Total (h)
Step (min)
# Qwen2.5-3B Inssttrrct
Joint Learning
14.8
5.58
CGPO
16.0
6.04
# Siven2.5-7B-Insee Joint Learning
17.8
6.72
CGPO
18.6
7.02
Randomizing domain order is necessary for effective cross-domain interactions. We conduct ablations to examine the necessity of randomizing domain order. Specifically, we compare the standard randomized variant with a fixed-order variant (CGPOnx). where the sequence of domains remains unchanged throughout training. As shown in Table 3 randomizing the order consistently leads to higher average performance across all benchmarks. This result highlights that randomization is essential: it ensures balanced sequential updates among domains, avoiding systematic bias in Hessian Messian-gradient interactions. In contrast, fixed ordering allows earlier domains to dominate updates, while later domains can only adapt passively, reducing overall multi-domain coordination.

---

Published as a conference paper at ICLR 2026
The mixing coefficient a plays critical role in balancing stability and curvature exploitation. To study its effect, we experiment with a E {0.9, ].5] and report the corresponding multi multi-domain performance in Table Among these choices, a 1.2 achieves the best overall average, reflecting favorable trade-off between retaining the base policy and incorporating curvature-informer updates. Notably, the average performance of all tested values exceeds that of the strongest baseline, FAMO (57.26), indicating that CGPO is robust to the choice of a. The fact that all a values are close to 1.0 suggests that the interpolation does not substantially change the effective learning rate; the observed gains therefore arise from the curvature- sequential updates rather than step size adjustments.
5 RELATED WORK
Multi-domain RL for LLMs. It is challenging to stably improve LLMs' multi-domain reasoning performance with RL. One difficulty lies in designing reward functions that generalize across diverse tasks. Some works propose broadly applicable reward computation, such simplifying binary rewards via ground-truth properties (Zhou et al. (025) or using correctness likelihood (Yu et al.] 2025b). Others adopt domain-specific reward designs, e.g., hybrid rule-based, sandbox, and LLM-as- a-judge systems (Li et al. 2025a). A second challenge is understanding cross-domain interactions: Cheng et al. (2025) analyze how single-domin training affects other skills, while Li et al. (2025b) extend this to math, coding. and puzzles. TAPPA (Yang et al. (026) reveals consistent properties of L.M attention patterns via query self-similarity across multiple tasks, but is primarily diagnostic and offers no optimization mechanism for mitigating cross-domain interference in multi-domain RL.. EVIC (Liang et al. 2025) effectively improves multi-domain SFT for L.MMs via curriculum learning guided by evolving gradient-based interactions; yet it does not readily extend to RL, since responses are generated online and the corresponding gradients cannot be precomputed in advance.
Mitigating Gradient Conflicts. Gradient interference is major obstacle in multi-task learning (Chen et al. 2025b). Approaches such as GradNorm (Chen et al. 2018). PCGrad (Yu et al. 2020), MGDA (Sener& & Koltun 20188), ConFIG (Liu et al. (2024). and CAGrad (Liu et al. (0221) resolve conflicts by balancing or projecting task gradients. While effective in standard MTI settings, they are difficult scale to RL for L.MMs: many require storing all task gradients on-device, causing memory bottlenecks, or operate reactively without leveraging reward-landscape geometry, leading to high variance under noisy rollout-based gradients. These limitations motivate scalable, memory-efficient mechanisms for mitigating cross-domain conflicts, as pursued by CGPO.
Second-Order Optimization Methods. The loss landscapes of deep neural networks are often highly complex, posing challenges for first-order optimization algorithms, such as gradient descent. Without insights into the geometric structure of the landscape, first-order methods can easily get trapped in saddle points or narrow valleys, making it difficult to reach better local optima. In contrast, second-order optimization methods, such as Newton's 's method, exploit geometric information like the Hessian matrix to precondition gradients according to the local curvature, offering stronger theoretical guarantees. To mitigate the computational cost of full Hessian computation, various approximate Newton methods have been proposed, including AdaGrad (Duchi al.] (011), K-FAC (Martens & Grosse 2015), GGT (Agarwal et al.) 2018). Shampoo (Gupta et al. 2018). and SOAP (Vyas al. 2025). Recent studies show that Newton's method and SOAP (Vyas et al. 2025) can alleviate gradient conflicts in PINNs (Wang et al. 2025), providing inspiration for our approach. However, due to the massive parameter scale of L.Mss, directly applying approximate variants of Newton's method in RL for LLMs is fundamentally infeasible (we provide detailed discussion in Appendix D.4b. Motivated by this, we distill the core idea of leveraging curvature information and develop CGPO, principled and scalable framework for multi- dolti-domain RL in LLMs.
6 CONCLUSION AND LIMITATIONS
We present CGPO. principled and scalable framework for multi-domain RL of LLMs. Inspired by Newton's 's method, CGPO leverages the geometric structure of the reward surfaces to precondition gradients, while avoiding the cost of full Hessian computation. Through randomized sequential updates, each domain's 's gradient is modulated by curvature information from other domains, fostering cross-domain interactions and implicitly aligning gradients. Experiments on diverse multi-domain dataset covering mathematical reasoning, code generation, scientific QA, and creative writing show that CGPO outperforms all baselines, achieving faster reward improvement and stronger multi-domain reasoning across all benchmarks. For the limitations of this paper, please refer to Appendix [F]
10

---

Published as a conference paper at ICLR 2026
7 ETHICS STATEMENT
This work studies multi-domain reinforcement learning for LLMs using publicly available or appro- priately licensed datasets across domains such as mathematies, coding, scientific QA, and creative writing. No human subjects were directly involved. While our methods improve cross-domain optimization, models trained with them could be misused to produce plausible but incorrect or unsafe outputs. We strongly discourage any deployment outside research contexts and emphasize that reward functions and training setups are designed to encourage safe and aligned outputs. All research was conducted in accordance with the ICLR Code of Ethics, with no conflicts of interest or external influence on methodology or results.
REPRODUCIBILITY STATEMENT
To facilitate reproducibility, we provide detailed descriptions our algorithm (CGPO) in Section 3.3 and Algorithm|] []]
including pseudo-code and key hyperparameters. Experimental setups, including data processing, reward functions, and evaluation benchmarks, are described in Section and Appendix C Where applicable, we provide references publicly available datasets. Our code is available at athttps: // httpttt.ttt.cog/ com Ctt.tb..cco//MRLLa........ Saiiuaabb-STC/CCCC USTC/CGPO All derivations, approximations, and additional analyses supporting the method are included in Appendix B Together, these materials provide sufficient information for replication of the reported results.
ACKNOWLEDGMENTS
This work was supported in part by National Key R&D Program of China under contract 2022ZD01 19801, National Nature Science Foundations of China grants U23A20388 and 62021001. We would like to thank all the anonymous reviewers for their insightful comments.
REFERENCES
Naman Agarwal, Brian Bullins, Xinyi Chen, Elad Hazan, Karan Singh, Cyril Zhang, and Yi Zhang. The case for full-matrix adaptive regularization. arXiv preprint arXiv: 1oKiv:1806.022 1806.02958, pp. 404-413, 2018.
Jacob Austin, Augustus Odena, Maxwell Nye, Maarten Bosma, Henryk Michalewski, David Dohan, Ellen Jiang, Carrie Cai, Michael Terry, Quoc Le, and Charles Sutton. Program synthesis with large language models, 2021. URL https:: http:///axvvvvoo org/ abs 2108 07732
Yinqi Bai, Jie Wang, Lei Chen, Zhihai Wang, Yumeng Li, Mingxuan Yuan, Jianye HAO, Defu Lian, and Enhong Chen. Evolving graph structured programs for circuit generation with large language models. In The Fourteenth International Conference on Learning Representations, 2026. URL https: /opps://openrvviii net/ htte:/formmiiddd00 (errummii-DDDSSSKHHHH
Sebastian Borgeaud, Arthur Mensch, Jordan Hoffmann, Trevor Cai, Eliza Rutherford, Katie Mil- lican, George Bm Van Den Driessche, -aptiste Lespiau, Bogdan Damoc, Aidan Clark, Diego De Las Casas, Aurelia Guy, Jacob Menick, Roman Ring, Tom Hennigan, Saffron Huang. Loren Maggiore, Chris Jones, Albin Cassirer, Andy Brock, Michela Paganini, Ge- offrey Irving, Oriol Vinyals, Simon Osindero, Karen Simonyan, Jack Rac, Erich Elsen, and Laurent Sifre. Improving language models by retrieving from trillions of tokens. In Kama- lika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Proceedings of the 39th International Conference Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 2206-2240, PMLR, 17-23 Jul 2022. URL https://proceeeinngs. ml www.mr..rrss/eee 62/ hwt.:///eoooobbeebbaeeaaaaaa .tml
:. Yao Cheng, Jianfeng Chen, Jie Chen, Li Chen, Liyu Chen, Wentao Chen, Zhengyu Chen, Shijie Geng, Aoyan Li, Bo Li, Bowen Li, Linyi Li, Boyi Jiaheng Liu, Kaibo Liu, Qi Liu, Shukai Liu, Siyao Liu, Tianyi Liu, Tingkai Liu, Yongfei Liu, Rui Long, Jing Mai, Guanghan Ning, Z. Y. Peng. Kai Shen, Jiahao Su, Jing Su, Tao Sun, Yifan Sun, Yunzhe Tao, Guoyin Wang. Siwei Wang, Xuwu Wang, Yite Wang, Zihan Wang, Jinxiang Xia, Liang Xiang, Xia Xiao, Yongsheng Xiao, Chenguang Xi, Shulin Xin, Jingjing Xu, Shikun Xu,
11

---

Published as a conference paper at ICLR 2026
Hongxia Yang. Jack Yang, Yingxiang Yang, Jianbo Yuan, Jun Zhang, Yufeng Zhang, Yuyu Zhang, Shen Zheng, He Zhu, and Ming Zhu. Fullstack bench: Evaluating ilms as full stack coders, 2025.
http::htta/aatttt org/ abs 2412. 00535
Dashiel Carrera, Zixin Zhao, Ashish Ajin Thomas, and Daniel Wigdor. Nabokov's cards: An ai assisted prewriting system to support bottom-up up creative writing. In Proceedings of the 2025 Conference on Creativity and Cognition, pp. 546-559, 2025.
Jie Chen, Zhipeng Chen, Jiapeng Wang, Kun Zhou, Yutao Zhu, Jinhao Jiang, Yingqian Min, Xin Zhao, Zhicheng Dou, Jiaxin Mao, Yankai Lin, Ruihua Song, Jun Xu, Xu Chen, Rui Yan, Zhewei Wei, Di Hu, Wenbing Huang, and Ji-Rong Wen. Towards effective and efficient continual pre- training of large language models. In Wanxiang Che, Joyce Nabende, Ekaterina Shutova, and Mohammad Taher Pilehvar (eds.), Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 5779- 5795, Vienna, Austria, July 2025a. Association for Computational Linguistics. ISBN 7777998889977-551-1177 doi: 10. 10.18633/11/022 acl-long. 289. URL wtt.ltttss//aaaattttloiooooo 2tt.aaaaaa.ooooooo..2.5...o....
Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde de Oliveira Pinto, Jared Ka- plan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, Alex Ray, Raul Puri, Gretchen Krueger, Michael Petrov, Heidy Khlaaf, Girish Sastry, Pamela Mishkin, Brooke Chan, Scott Gray, Nick Ryder, Mikhail Pavlov, Alethea Power, Lukasz Kaiser, Mohammad Bavarian, Clemens Winter, Philippe Tillet, Felipe Petroski Such, Dave Cummings, Matthias Plappert, Fotios Chantzis, Elizabeth Barnes, Ariel Herbert-Voss William Hebgen Guss, Alex Nichol, Alex Paino, Nikolas Tezak, Jie Tang. Igor Babuschkin, Suchir Balaji, Shantanu Jain, William Saunders, Christopher Hesse, Andrew N. Carr, Jan Leike, Josh Achiam. Vedant Misra, Evan Morikawa, Alec Radford, Matthew Knight, Miles Brundage, Mira Murati, Katie Mayer, Peter Welinder, Bob McGrew, Dario Amodei, Sam McCandlish, Ilya Sutskever, and Wojeiech Zaremba. Evaluating large language models trained on code, 2021. htRttttss://// //axxiv hbwgggaas////// 07. 03374
Weiyu Chen, Baijiong Lin, Xiaoyuan Zhang, Xi Lin, Han Zhao, Qingfu Zhang, and James T Kwok. Gradient-based multi-objective deep learning: Algorithms, theories, applications, and beyond. arXiv preprint arXiv:2501. 10945, 2025b.
Zhao Chen, Vijay Badrinarayanan, Chen-Yu Yu Lee, and Andrew Rabinovich. Gradnorm: Gradient normalization for adaptive loss balancing in deep multitask networks. In International conference on machine learning, pp. 794-803. PMLR, 2018
Zhoujun Cheng, Haoyu Dong, Zhiruo Wang, Ran Jia, Jiaqi Guo, Yan Gao, Shi Han, Jian-Guang Lou, and Dongmei Zhang. Hitab: A hierarchical table dataset for question answering and natural language generation. In Proceedings of the 60th annual meeting of the association for computational linguistics (volume I: long papers). pp. 1094-1110, 2022.
Zhoujun Cheng, Shibo Hao, Tianyang Liu, Fan Zhou, Yutao Xie, Feng Yao, Yuexin Bian, Yonghao Zhuang, Nilabjo Dey. Yuheng Zha, et al. Revisiting reinforcement learning for Ilm reasoning from a cross-domain perspective. arXiv preprint arXiv:2506. 144.1.55.....55. 2025.
Michael Crawshaw. Multi-task learning with deep neural networks: A survey. arXiv preprint arXiv:2009.09796, 2020.
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research, 12(7), 2011.
Daniel Fein, Sebastian Russo, Violet Xiang, Kabir Jolly, Rafael Rafailov, and Nick Haber. Litbench: A benchmark and dataset for reliable evaluation of creative writing, arXiv preprint arXiv:2507 .00669, 2025.
Jiwwei Gu, Zacc Yang, Chuanghao Ding, Rui Zhao, and Fei Tan. CMR scaling law: Predict- ing critical mixture ratios for continual pre-training of language models. In Yaser Al-Onaizan, Mohit Bansal, and Yun-Nung Chen (eds.), Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing, pp. 11143-161622 Miami, Florida, USA, November 2024. Association for Computational Linguistics. doi: 10. 18.18653/ http3///12222.mmmm/// waammp-main..... main. 903. URL https: /ttp:://acl antho ogy. log..oog/......... emn p- ma i n.. 9037
12

---

Published as a conference paper at ICLR 2026
Vineet Gupta, Tomer Koren, and Yoram Singer. Shampoo: Preconditioned stochastic tensor optimiza- tion. In International Conference on Machine Learning, pp. 1842-1850. PMLR, 2018.
Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. Measuring mathematical problem solving with the math dataset. NeurIPS, 2021.
Jingcheng Hu, Yinmin Zhang, Qi Han, Daxin Jiang, Xiangyu Zhang, and Heung- Yeung Shum. Open-reasoner-zzro: An open source approach to scaling up reinforcement learning the base model. arXiv preprint arXiv:2503.24299, 2025.
Gautier Izacard, Patrick Lewis, Maria Lomeli, Lucas Hosseini, Fabio Petroni, Timo Schick, Jane Dwivedi-YY Yu, Armand Joulin, Sebastian Riedel, and Edouard Grave. Atlas: few-shot learning with retrieval augmented language models. J. Mach. Learn. Res., 24(1), January 2023. ISSN 1532-4435.
willchat-cceaiiee-iiitig--k--- 2025. URL https::/hugginggtttt aatasstt///evviaahhh/aaaacccccaa www.eeetttttttiiiiiie.ee wwitiiiig rft
kevinshin.
co/
Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu. Joseph Gonzalez, Hao Zhang, and Ion Stoica. Efficient memory management for large language model serving with pagedattention. In Proceedings of the ACM SIGOPS 29th Symposium on Operating Systems Principles, 2023.
Derek Li, Jiaming Zhou, Amirreza Kazemi, Qianyi Sun, Abbas Ghaddar, Mohammad Ali Alomrani, Liheng Ma, Yu Luo, Dong Li, Feng Wen, et al. Omni-thinker: Scaling comssn generalization in llms via multi-task rl with hybrid rewards. arXiv preprint S.KVv:2507.14788. 2025a.
Yu Li, Zhuoshi Pan, Honglin Lin, Mengyuan Sun, Conghui He, and Lijun Wu. Can one domain help others? dentrce study on multi-domain reasoning via reinforcement learning. arXiv preprint arXiv:2507. 17512, 2025b.
Xize Liang, Lin Yang, Jie Wang, Yiyang Lu, Runyu Wu, Hanzhu Chen, and Jianye Hao. Boosting multi-domain fine-tuning of large language models through evolving interactions between samples. In International Conference on Machine Learning, pp. 37427-37441. PMLR, 2025.
Bill Yuchen Lin, Ronan Le Bras, Kyle Richardson, Ashish Sabharwal, Radha Poovendran, Peter Clark, and Yejin Choi. Zebralogic: On the scaling limits IIms for logical reasoning. In Forty-second International Conference on Machine Learning.
Bo Liu, Xingchao Liu, Xiaojie Jin, Peter Stone, and Qiang Liu. Conflict-averse gradient descent for multi-task learning. Advances in Neural Information Processing Systems, 34:18878-1888 2021.
Bo Liu, Yihao Feng, Peter Stone, and Qiang Liu. Famo: Fast adaptive multitask optimization. Advances in Neural Information Processing Systems, 66:57226-572433 2023.
Qiang Liu, Mengyu Chu, and Nils Thuerey. Config: Towards conflict-rre training of physics informed neural networks. arXiv preprint arXiv:2408.11100 2024.
Weiting Liu, Han Wu, Yufei Kuang, Xiongwei Han, Tao Zhong, Jianfeng Feng, and Wenlian Lu. Automated optimization modeling via localizable error-driven perspective. arXiv preprint arXiv:2602. 1164, 2026.
Xueguang Ma, Qian Liu, Dongfu Jiang, Ge Zhang, Zejun Ma, and Wenhu Chen. Generel-reasoner: Advancing Ilm reasoning across all domains. arXiv preprint arXiv:2505. 14652, 2025.
MAA. American mathematics competitions, 2023.
URL
https: patp::://maa.o0/
studeent-programs
James Martens and Roger Grosse. Optimizing neural networks with coonecker-factored approximate curvature. In International conference on machine learning, pp. 2400-2417. PMLR, 2015.
13

---

Published as a conference paper at ICLR 2026
Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Ar- cas. communication-cffiitie learning of deep networks from decentralized data. Artificial intelligence and statistics, pp. 1273-1283 PMLR, 2017.
Alex Nichol, Joshua Achiam, and John Schulman. On first-order meta-learning algorithms. arXiv preprint 1oXiv:8803.02999 2018.
Nitral-AI, Creative writing-shareept, 2024. URL wwt.:ittss///ngggiggae...../////
NitrallAA//Creative__rrr
Sha reGPT
David Rein, Betty Li Hou, Asa Cooper Stickland, Jackson Petty, Richard Yuanzhe Pang, Julien Dirani, Julian Michael, and Samuel R. Bowman. Gpqa: A graduate- google-proof q&a benchmark, 2023. URL www.h:tsss...mmm http:///axxvvvrrrg.e//
12022
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint urXiv:17070063411 2017.
Ozan Sener and Vladlen Koltun. Multi-task learning multi-objectivee optimization. Advances in neural information processing systems, 31, 2018.
Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song, Xiao Bi, Haowei Zhang, Mingchuan Zhang, YK Li, et al. Deepseckmath: Pushing the limits of mathematical reason- ing in open language models. arXiv preprint arXiv:2402.033000 2024.
Guangming Sheng, Chi Zhang, Zilingfeng Ye, Xibin Wu, Wang Zhang, Ru Zhang, Yanghua Peng, Haibin Lin, and Chuan Wu. Hybridflow: A flexible and efficient rlhf framework. In Proceedings of the Twentieth European Conference on Computer Systems, pp. 1279-1297, 2025.
Mustafa Shukor, Louis Bethune, Dan Busbridge, David Grangier, Enrico Fini, Alaaeldin El-Nouby, and Pierre Ablin. Scaling laws for optimal data mixtures. In NeurIPS, 2025. URL https:
//arxiv. org/ wwt.orr///bbbbeeee attpaab///.//.2 09404
Petru Soviany, Radu Tudor lonescu, Paolo Rota, and Nicu Sebe. Curriculum learning: survey. International Journal of Computer Vision, 13266))1526-1565 2022.
Richard S Sutton, Andrew G Barto, et al. Reinforcement learning; An introduction, volume 1. MIT press Cambridge, 1998.
Team, Xinrun Du Yifan Yao, Kaijing Ma, Bingli Wang, Tianyu Zheng, King Zhu, Minghao Liu, Yiming Liang, Xiaolong Jin, Zhenlin Wei, Chujie Zheng, Kaixin Deng, Shawn Gavin, Shian Jia, Sichao Jiang, Yiyan Liao, Rui Li, Qinrui Li, Sirun Li, Yizhi Li, Yunwen Li, David Ma, Yuansheng Ni, Haoran Que, Qiyao Wang, Zhoufutu Wen, Siwei Wu, Tyshawn Hsing, Ming Xu, Zhenzhu Yang, Zekun Moore Wang, Junting Zhou, Yuelin Bai, Xingyuan Bu, Chenglin Cai, Liang Chen, Yifan Chen, Chengtuo Cheng, Tianhao Cheng, Keyi Ding, Siming Huang, Yun Huang. Yaoru Li, Yizhe Li, Zhaoqun Li, Tianhao Liang, Chengdong Lin, Hongquan Lin, Yinghao Ma, Tianyang Pang, Zhongyuan Peng, Zifan Peng. Qige Qi. Shi Qiu, Xingwei Qu Shanghaoran Quan, Yizhou Tan, Zili Wang. Chenqing Wang, Hao Wang, Yiya Wang, Yubo Wang, Jiajun Xu, Kexin Yang, Ruibin Yuan, Yuanhao Yue, Tianyang Zhan, Chun Zhang, Jinyang Zhang, Xiyue Zhang, Xingjian Zhang, Yue Zhang, Yongchi Zhao, Xiangyu Zheng, Chenghua Zhong, Yang Gao, Zhoujun Li, Dayiheng Liu, Qian Liu, Tianyu Liu, Shiwen Ni, Junran Peng. Yujia Qin. Wenbo Su, Guoyin Wang. Shi Wang, Jian Yang, Min Yang, Meng Cao, Xiang Yue, Zhaoxiang Zhang, Wangchunshu Zhou, Jiaheng Liu, Qunshu Lin, Wenhao Huang, and Ge Zhang. Supergpqa: Scaling 1lm evaluation across 285 graduate disciplines, 2025. URLhetpe
wtt.xxxxivv...//.. abs / 25..250222.0 11255.2.44733 14739
Nelson Vithayathil Varghese and Qusay H Mahmoud. A survey of multi-task deep reinforcement learning. Electronies, 9(9):1363, 2020.
Nikhil Vyas, Depen Morwani, Rosie Zhao, Itai Shapira, David Brandfonbrener, Lucas Janson, and Sham M Kakade. Soap: Improving and stabilizing shampoo using adam for language modeling. In The Thirteenth International Conference on Learning Representations, 2025.
Sifan Wang, Ananyae Kumar Bhartari, Bowen Li, and Paris Perdikaris. Gradient alignment in
neural networks: A second-ordes optimization perspective. arXiv preprint arXiv:2502. 006604, 2025.
14

---

Published as a conference paper at ICLR 2026
Yuning Wu, Jiahao Mei, Ming Yan, Chenliang Li. Shaopeng Lai. Yuran Ren, Zijia Wang, Ji Zhang, Mengyue Wu, Qin Jin, and Fei Huang. Writingbench: A comprehensive benchmark for generative writing, 2025. URL
05244.
An Yang, Beichen Zhang, Binyuan Hui, Bofei Gao, Bowen Yu, Chengpeng Li, Dayiheng Liu, Jian- hong Tu, Jingren Zhou, Junyang Lin, et al. Qwen2. 5-math technical report: Toward mathematical expert model via ieffimpprveeent. arXiv preprint arXiv:2409.12122, 2024.
Qingyue Yang, Jie Wang. Xing Li, Yinqi Bai, Tong Xialiang. Hui-Ling Zhen, Jianye Hao, Mingxuan Yuan, and Bin Li. Why attention patterns exist: A unifying temporal perspective analysis. In The Fourteenth International Conference on Learning Representations, 2026. URL https:////
openreview net forum? -iaoeeidddoohouoiiin
Yufan Ye, Ting Zhang, Wenbin Jiang, and Hua Huang. Process-supervised reinforcement learning for code generation. arXiv preprint arXiv:2502.01715. 2025.
Qiying Yu, Zheng Zhang, Ruofei Zhu, Yufeng Yuan, Xiaochen Zuo, Yu Yue, Weinan Dai, Tiantian Fan, Gaohong Liu, Lingjun Liu, et al. Dapo: An open-source IIm reinforcement learning system at scale. arXiv preprint ar Arx:v:2503. 144466 2025a.
Tianhe Yu, Saurabh Kumar, Abhishek Gupta, Sergey Levine, Karol Hausman, and Chelsea Finn. Gradient surgery for multi-task learning. Advances neural information processing systems, 33: 5824-5836, 2020.
Tianyu Yu, Bo Ji, Shouli Wang. Shu Yao, Zefan Wang, Ganqu Cui, Lifan Yuan, Ning Ding, Yuan Yao, Zhiyuan Liu, et al. Rlpr: Extrapolating rivr to general domains without verifiers. arXiv preprint arXiv:2506.182. 18254, 2025b.
Huaye Zeng, Dongfu Jiang, Haozhe Wang, Ping Nie, Xiaotong Chen, and Wenhu Chen. Acecoder: Acing coder rl via automated test-case synthesis. arXiv preprint arXiv:2502.01718. 2025.
Xiangxin Zhou, Zichen Liu, Anya Sims, Haonan Wang. Tianyu Pang, Chongxuan Li, Liang Wang, Min Lin, and Chao Du. Reinforcing general reasoning without verifiers. arXiv preprint MrXiv:2505.21493 2025.
15

---

Published as a conference paper at ICLR 2026
A LLM USAGE STATEMENT
In preparing this manuscript, we used a large language model (LLM) in two distinct ways. First, we employed LLMs assistive tool for text refinement, including improving grammar, wording, and clarity. Second, LLMs themselves are the primary subject this research: we study reinforcement learning (RL) training for LLMs. Accordingly, all experiments involve using large models for training, inference, and scoring, as part of the methodology under investigation.
All scientific content, including problem formulation, methodology, experiments, and conclusions, was developed and verified entirely by the authors. The authors take full responsibility for the integrity and accuracy of the manuscript. No LLM was credited an author, and all substantive research contributions are attributable exclusively to the human authors.
B
MATHEMATICAL DERIVATIONS
B. DETAILED DERIVATION OF EQ.
Eq.
P2
in Section 3.2 states:
gy
( post
gy
((+)
pre
H
((x) pre
(0
{($)
pre
H
A(L) pre
g:
(6)
Derivation: Assuming the gradient function (0) is smooth, we apply a first-order Taylor expansion around 0f00
(1) post
gj
( pre
+
H
( pre
(e())
post
g piz
+ (11
where ()) = (0) is the Hessian matrix for domain j. and AO higher-order terms and rearranging gives:
oftie
Ope.
Neglecting
(
((e)
H
((x)
(t)) post
((a)
(8)
post
pre
pre
pre
In policy optimization, parameters are updated via gradient ascent (maximizing rewards):
(e())
post
pre + vg:
pre
(9)
where 1) is the learning rate. Substituting this into the previous equation yields:
post
pre
n/:
(0()) pre
(10)
and therefore,
gj
(000) post
gy
0
H
( pre
g:
(11)
which is Eq. 2 This approximation shows that the gradient update from domain influences the gradient of domain through the curvature of domain j.
B.2 DETAILED DERIVATION OF EQ. 33) AND EQ.
Eq. 3 and Eq.
in Section state:
D)
(0 = Bo(H)(0o)
H (0(0)(00)8011)00))) + O(1²)
(12)
1=1 - |o(s))
(0( = (o())(0o)
(o (ro))doo)g.ll (60) + O(72)
(13)
Derivation:
Consider the randomized sequential update: domains are processed the order The parameter update for each domain (using gradient ascent) is:
Øk Φ 7h ())8=()((0x1))
(14)
16

---

Published as a conference paper at ICLR 2026
where 7/ //)D Σx= (o((s) is the scaled learning rate.
For domain (kk), its gradient is evaluated at Φ Using a Taylor expansion around do:
(k) (#k-1) (k) +
Φb) + (17²)
(15)
Now compute Φh-1 do. Note that:
Φk-1 Φ0 +
(or 01 = do Σ
(
(16)
1=1
1=1
To first order, we approximate (r) (01-1) go(t) (0o) (error O17²)))
Φk-1 Φ0 2
(1) (o))
(17)
1=1
Substituting into the Taylor expansion:
Bo (() (o)) (R)
Σ (g (0o) O01))
(18)
Substituting n -/) / (D(s)) gives Eq. 3
we assume uniform batch sizes, i... Da(i)// Σx Da(s) = 1/K, then m = 1//K, which simplifies to Eq. 4
B.3 DETAILED DERIVATION OF EQ. (5)
Eq. (5) in Section 3.3 states:
077²
aløk 0o)
(0o)
H (k)) (0)()))))((000) + O
(19)
K
K2
k=
Derivation: The total parameter change is:
ØK do Σ (Øk 2
Σ (k) (0k-1).
(20)
Using the approximation from Eq. (uniform batch sizes):
go(k) _o(k) (o(k)
1 K
Ho(k) (o))())))
(21)
and substituting ?i- = 1//K K:
K
k-1 1
ΦK - Φn ~
Bo(ks) Bo(k) (do)
H (o0
K
K
k=1
1=1
1/²
K
8a (K)
H
(a) (o))
(22)
K
K2
k= 1=1
Multiplying by the mixing coefficient a:
K k-1
alok 20) 8o(h)
an
ΣΣ
H. Hot (do) (o)))(((())))
(23)
K2
Note that ン E (() Ek g& (do) (permutation invariant), yielding Eq. 5
17

---

Published as a conference paper at ICLR 2026
B.4 DERIVATION OF GRADIENT ALIGNMENT SYMMETRIZATION
In Section 3.3 it is mentioned that after randomization, the cross-term expectation symmetrizes as:
0
H. () (óo) + (00))) = dán
(g. (00) g) (60)).
(24)
Derivation: The key mathematical insight is the following identity concerning the gradient of the inner product between two gradients.
Consider the inner product S( gi()o) (T gy(())) The gradient of this scalar function S with respect to Φn is given by:
Ven S = Vo (8:())) B, (do)) = (oo (6o) + H, g:
(25)
where we have used the product rule and the symmetry of the Hessian matrices, H, = HJ. This result can be seen by noting that the derivative of g, w.r.t. do is (0g:/00o) gj + g (0g,/ - H,g, + Hy. Since g, H2 is row vector, its transpose is H The gradient (as a column vector) therefore H,g, + Hggi
Under a randomized ordering the expectation of the cross-term involving (o(k))o() for k > I will involve pairs (i. symmetrically. The update term derived from the second-order expansion is proportional to H:gj- The symmetric form H,g; + Hgg; appearing in the gradient of the inner product V do (g. gy) indicates that, in expectation, the update encourages an increase in the inner product between the gradients of different domains, thus promoting their alignment.
Remark. We would like clarify the intended meaning of Eq. (5) and the role of the expectation over permutations, in order to avoid possible ambiguities and to keep the presentation self-contained.
(1) Interpretation of Eq. (5) Eq. (5) is obtained from deterministic Taylor expansion of one sequential update pass conditioned on a fixed permutation a. The resulting parameter change decomposes into: (i) a first-order term corresponding to aggregated gradients, and (ii) second-order interaction term involving fessian-gradicn products. These Hessian-gradies interaction terms arise deterninistically from executing sequential update under a specific ordering, they do not rely on randomness or averaging. The expression makes explicit the structural cross cross-domain second-org -oddr interactions induced by sequential updates.
(2) Role of the expectation over The expectation over permutations is used to express a sym- metry property. To make this more concrete, imagine that at the same parameter 0., we were able Able-hypothetically, since the algorithm does not actually do this-to to sample M independent R
this hy- pothetical scenario, M oc, the events "i appears before j" and "j appears before i" would occur with essentially equal frequency for every pair (i.j). This limiting symmetry is exactly what our expectation argument intended express, and it is what leads to the symmetric combination H,g, + H in the discussion following Eq. (
permutations {d!m)) M
mm1 each corresponding to an ordering T, (mm)
(m)) (k)) カ-1
In the actual algorithm, of course, we sample only one permutation at each iteration. This introduces sampling error-but not bias in the expectation sense- we do not average over multiple permutations.
Importantly, this sampling error does not accumulate in a harmful way in practice. A helpful way to view this is through an analogy with standard SGD: each stochastic gradient is. in expectation, equal to the true gradient (just as the contributions of H,g, and H are symmetric in expectation), yet in practice we use only one stochastic gradient per step rather than averaging many samples just as our algorithm samples only permutation per iteration rather than averaging over many permutations at the same parameter. This practice in SGD does introduce variance and error, but it does not undermine either the effectiveness of SGD the usefulness of the statement that "the stochastic gradient equals the true gradient in expectation". The same phenomenon appears in our algorithm.
Therefore, when refer to expectation, mean the conditional expectation taken at fixed 6,, ie., conditional on the past history F.-1-just as the expectation a stochastic gradient in SGD is interpreted conditional on the current parameter value.
18

---

Published as a conference paper at ICLR 2026
C MORE DETAILS OF EXPERIMENTS
C.I TASKS AND DATASETS
We focus on enhancing LLMs' overall capabilities across four domains-matthmmaical reasoning. code generation, scientific QA, and creative writing. These domains not only represent core areas of current research interest but also span four distinct forms of reward feedback, thereby ensuring both comprehensivenes: and diversity.
Mathematics: we construct a subset of 6,250 samples from the Guru dataset (Cheng et al., 2025) This includes the 5,000 casiest problems (ranked by the pass rate of Qwen2.5-78 - 1,250 more challenging ones, ensuring a balance between accessible and difficult problems.
Instruct) and
Code generation: select total of 4,740 samples from Guru. Specifically, we take all 3,791 problems with Qwen2.5-7B-Instrul 's pass rate of at least 25% and add 949 problems randomly sampled from the remainder, yielding an approximate 4:1 ratio between easier and harder samples.
Scientific QA: we include the entire STEM split of Guru, resulting in 3,591 samples. This preserves the full coverage of seience-related reasoning tasks while maintaining consistency with prior benchmarks.
Creative writing: we randomly sample 2,000 samples each from three popular Hugging- face Litesch (Fein al. (0225), SaaiiveWrWiiinSSaaeeee (Nitral-Al (024). and wildchat- wwwattterreatvvviiiiiee -rrfi (kevinshin 2025) to construct dataset of 6,000 samples, ensur- ing stylistic variety and broad coverage of open-ended writing abilities.
C.2 BASELINES
We compare our CGPO against four representative baselines: joint learning, Omni-Thinker (Li et 2025a), Self-Paced CL, and FAMO (Liu et al. 2023
Joint learning. Joint learning the most basic paradigm MTL. It aggregates the loss functions of all tasks into single objective, enabling simultaneous optimization. As straightforward training strategy without any task-specific adjustments, joint learning serves as a reference point for evaluating improvements brought by more advanced methods.
Omni- Thinker, Tnii-Thinker belongs to progressive CL methods as categorized in (Soviany et al. 2022). introduces the backward transfer (BWT) metric to quantify the extent of catastrophic forgetting across domains. Based on BWT analysis, [i et al. (2025a) proposes a fixed training order-code - math >> scientific QA - creative writing with the goal of minimizing forgetting induced by multi-domain learning.
Self-paced CL. Self peled CL enables the model to adaptively select training samples according employ O0we2.55-7B-Instruue
to its learning state. In our implementation, winrate from easy to difficult, and train sequentially following this order. This curriculum reduces the risk of being misled by difficult samples in the early stages, thereby improving stability and promoting better generalization.
rank samples by
FAMO. FAMO is the improvement rate of the task that progresses the slowest, ensuring that all tasks advance at comparable pace. This balanced optimization strategy suppresses task dominance and guides the model toward solutions that are both fairer across tasks and stronger in overall performance. FAMO approximates weight updates using historical loss values instead of explicitly computing multi-task gradients, reducing per-iteration time and memory complexity to O(1). This efficiency makes particularly suitable for large-scale LLM training.
approach for MTL. adjusts loss weights maximize
C.3 REWARD FUNCTIONS
For all domains, we require the model to enclose its reasoning process within <think></thinks tags. The reward functions for the four domains as follows.
19

---

Published as a conference paper at ICLR 2026
Math, We adopt rule-based reward function:
1.0,
if has valid format and verify math math (Pans, a) = true,
mmtth (0, a)
0.5, if has valid format but verify math (Cans; a)) false,
1.0, o has an invalid format,
where Oans denotes the predicted answer extracted from structured tags (e.g.. <answer><< </nnw//// in the model output o, and verifyman (:.)) (:
checks symbolic equivalence
between Pans and the ground-truth answer via deterministic parser (e.g., handling equivalent forms of expressions or equations).
Code generation. We adopt a sandbox-based unit test reward:
1.0,
(Oans, if has a valid format and exec(Oans) } unittest (Oans, test_case),
0.5, if has a valid format but exec f unittest (Oans, test_case), (O, test case) =
1.0, if has invalid format (syntactically invalid),
where Oans is the generated code, executed in a sandbox and validated against the unit tests associated with the sample; += denotes logical satisfaction.
Scientific QA. We employ 1.5B General Verifie CCeen et al. 2025 to assess consistency between the model's output and the ground-truth answer:
1.0 0.05 min(
if o has valid format and Oans = a, |0],,
Typ "0 a) =
0
o has valid format but Oans 7 a, if has an invalid format,
1.0,
where gass the extracted answer content. Here, "valid format" means the response adheres to QA conventions (e.g., garbled text, complete sentences).
Creative writing. We adopt an -ude strategy, scoring the model's output against reference Oer via pairwise comparison:
1.0,
if has a valid format and > Otefs
0.25,
if has a valid format and o ~ O
Twriting (o, Oref)
- 0.5,
if has a valid format and < Opets
1.0,
if has invalid format,
where o > Our (preferred), (tie), and < Ozr (worse) are determined by a fixed evaluator (Qwen2.5-72B-Instrue serving as the judge.
C.4 HYPERPARAMTERS
We use learning rate of x 10-6, prompt batch size of 128, mini-batch size of 64, group size of 8, a rollout temperature of 1.0, flow = 0.2, &high 0.28, and 3 = 0.001 for CGPO and all baselines. All methods trained for one epoch. For the mixing coefficient a, tune it within the range of 0,5-1.5, and provide an ablation study on a Section 4.3
C.5 EVALUATION
To comprehensively evaluate cross-domain capabilities, adopt authoritative benchmarks spanning four domains: Math, Coding, Scientific QA, and Creative Writing. The evaluation settings are detailed below:
Math domain
MATH500 (Hendrycks et al. 202T): A set of 500 challenging problems sampled from the full MATH dataset, covering seven areas: elementary algebra, algebra, geometry. number theory, combinatorics, probability, and calculus. Problems presented in open-ended form and require precise solutions. This benchmark widely adopted for assessing LLMSs' mathematical reasoning and problem-solving abilities.
20

---

Published as a conference paper at ICLR 2026
AMC 2023 (MAA 2023): A set of 50 questions taken from the AMC 12A and 12B (2023) contests, spanning algebra, geometry, number theory. combinatorics, and probability. Multiple- choice options removed, requiring models to directly output the final answer. This benchmark focuses on higher-order reasoning, problem analysis, and accurate calculation.
Coding domain
Consisting of 164 human-written Python programming tasks, ranging from basic algorithms to medium-level function implementations. It evaluates whether models can generate correct and executable code from natural language descriptions.
HvmanEval (Chen et al. 2021:
.BPPP (Austin et al., )021): A collection of 974 beginner-level Python problems designed to test the ability to synthesize short programs from natural language instructions. It is standard benchmark for fundamental code generation.
Scientific QA domain
GPQA (diamond split) (Rein et al. 2023 Graduate-level QA items written and verified by domain experts across physics, chemistry, biology, and earth sciences. The diamond split represents the most difficult and highest-quality subset, specifically constructed to prevent shallow memorization or pattern matching. To ensure consistent evaluation, we reconstruct ordered option lists using randomized indexing.
SuperGPQA (Team al. 2025): Comprising 285 interdisciplinary graduate-level reasoning problems, curated to prevent direct solutions via search engines. To reduce computational cost, we use random seed 42 sample 200 problems, ensuring both representativeness and reliable measurement of deep reasoning ability.
Creative Writing domain
WritingBench (Wu et al. 2025 : A benchmark of 1000 real-world writing tasks spanning 6 domains and 100 sub-themes, covering diverse styles, task types, and difficulty levels. It evaluates generated text on quality, coherence. creativity, and task alignment through structured scoring framework. For efficiency, we sample 200 requests using random seed 42, and apply the official critic model Wr t ch cri Criti c-Model Metodel-@wen- -Qwen-7p automated scoring, striking a balance between evaluation cost and representativeness.
D MORE DISCUSSIONS
D.1 APPLICABILITY CGPO MULTI-DOMAIN PRE-TRAINING
Although our experiments focus on the RL post-training stage, the underlying mechanism of CGPO naturally extends to the multi-domain setting of L.LM pre-training. Pre-training corpora are inherently heterogeneous, and the aggregation of losses across diverse domains can lead to complex optimiza- tion landscape. Since CGPO designed to alleviate such difficulty by leveraging curvature-informed interactions induced by sequential updates, the framework is conceptually agnostic to the specific form of the loss and can, in principle, be applied during pre- training without modification.
It is also worth noting that CGPO is developed to address challenges unique RL for LLMs, many of which are absent in the pre-training stage. As a result, the design space for multi-domain optimization during pre-training is substantially broader. When the entire corpus is available offline, practitioners may employ a wide range of well-established approaches, including data mixture and sampling strategies (Shukor al.] 20255 Gu et al.) 2024), continual staged domain- specific pre-training
(Chen et al., (025a), and retrieval-augmented pre-training (Izacard et al.) 2023 Borgeaud et al. 2022). These techniques are not directly applicable in RL4LLMS but can be highly effective during pre-training, making the relative advantage of CGPO in this setting an open empirical question.
Model Qwen-
21

---

Published as a conference paper at ICLR 2026
D.2 NON-UNIFORM DOMAIN IMPORTANCE IN MULTI-DOMANE TRAINING
In practical multi dolli- applications, different domains may carry different levels of importance. While the main paper focuses on the uniform-weight objective
1 K km
π(0)
(z),),
this choice primarily for conceptual clarity and to highlight the core contribution of CGPO-namely, its ability mitigate cross-domain optimization conflicts and improve multi-domain reasoning performance.
The CGPO framework can be naturally extended to settings in which domains assigned non- uniform importance. Let each domain be associated with a user-defined weight wy satisfying
Σ kml Wk = 1. The training objective can then be written as
1 K J(0:w) = Σ Tx) = K ((KW))J(())
K
(26)
This formulation is equivalent to scaling each fomainn-specini loss and its corresponding gradient by a factor proportional to its importance. Crucially, no modification the CGPO algorithm required: the sequential updates, geometric interactions, and final interpolation behave identically in the uniform uniform-weight weight case, with the only difference being the importance-adjuster gradient contributions. This property allows CGPO seamlessly accommodate prioritized tasks, enabling model practical multi-domain scenarios in which some domains or skills must be emphasized more heavily than others.
D.3
WHY JOINT LEARNING CANNOT REPRODUCE OUR CROSS-DOMAIN MECHANISM
In this section, we provide additional analysis comparing joint learning with the proposed sequential mechanism, clarifying why joint learning cannot recover the same cross-domain geesiinn-gradeent interactions.
D.3.1 SEQUENTIAL UPDATES INDUCE CLEAN CROSS-DOMANE INTERACTIONS
As shown in Eq. (f) of the main paper, single sequential pass over the domains Jamains-correspoddiig to one parameter update spplae--yelds,, up to (v7)),
K
K k-1
ΦK 20 =
7 K
BA (xn) +
(60))) + O
(27)
K2
k=1
where a is the random permutation sampled at this iteration.
Crucially, the expression above describes the update for a fixed permutation a. Since our algorithm re-samples a independently at each iteration, the relevant quantity for understanding the behavior of the sequential mechanism is the expectation over a. Taking expectation symmetrizes the pairwise interactions: each ordered pair (i, appears with equal probability. After symmetrization,
have (do)g, (6)) (go)g (dn) Jog (gr())) go()) (see Appendix B.4), yielding
interpretable alignment effect across domains.
D.3.2 TWO-STEP JOINT LEARNING YIELDS MIXED SECOND-ORDER TERMS
To analyze why joint learning cannot replicate this mechanism, consider two consecutive joint- tran--trainig updates. Let
K
1 K
(())
Σ (O)
g()) VL(0) =
Σ
ga (0), H(O) = v²c(0) K k=1
1 Σ Ha
K
K
k=1
1=1
22

---

Published as a conference paper at ICLR 2026
Performing gradient- steps with step size 77- note that unlike the sequential pass above, these constitute two separate parameter updates and expanding up to second order gives
0++2 -2
1)
(()))((())
K (o (x))
21 K
Tf²
(0.) + K
Σ ISEJSK
H,(())g, +
Σ K
(28)
This expression reveals three types of contributions:
1. Single-domain gradients (O)):
2. Cross-domain Hessian-gradient interactions (,)0)),, for #
3. Self-curvature terms (o))
The presence the selftcurv terms the key structural difference from Eq. Because both updates in joint learning are taken with respect to the same aggregated loss, these self-curvature components naturally arise and are typically of comparable magnitude the cross-domain terms. As result, they can partially or fully cancel cross-domain contributions depending curvature structure. Thus, joint learning does not isolate cross-domain interactions. Its second-order structure an inseparable mixture of self- and cross-terms, lacking the clean symmetry and interpretability obtained under the sequential scheme.
D.3.3 IMPLICATIONS FOR GRADIENT ALIGNMENT
Because joint learning yields both H:g, and Hggs terms, effective update cannot be reduced to a symmetric pairwise structure. In particular, it cannot be rewritten as the gradient of inter-domain alignment quantity such as g; g,. The self-curvature terms disrupt this symmetry, preventing the simplification that underlies the alignment interpretation in our method.
By contrast, our sequential scheme avoids H&gs entirely: each domain is updated once per sequential pass, and its gradient is evaluated only after perturbations induced by other domains. Combined with the expectation over random permutations, this yields a clean, symmetric second-order term capturing cross-domain interactions.
D.4 WIIY APPROXIMATE VARIANTS OF NEWTON'S METHOO ARE INFEASIBLE FOR RL TRAINING OF LLMS
Second-order optimization methods broadly aim to exploit curvature information typpiilll through matrix-based preconditioning- enable more geometrically informed parameter updates. These ap- proaches span a wide family of techniques, including coonecker-factoriae layer-wise matrix preconditioners, and approximate Newton-style updates. To illustrate why such methods become impractical in RL. training of LLMs, we examine three of the most representative and advanced instances in this family-K-FAC K-FAC (Martens & Grosse 2013). Shampoo (Gupta et al.) 2018), and SOAP (Vyas et al.) 20225) and and analyze the computational and memory implications of applying their core mechanisms LLM scale.
natural-gradien methods,
D.4.1 K-FAC
Cronecker-factoree approximation to natural gradient descent. For a fully-connected (or linear) layer with weight matrix W e R&daa input activations a E 1R and backpropagated output gradients g Rd the gradient be written (for single sample) as V = ga If we vectorize W into w = vec( E R Rdadi, the Fisher information block corresponding to
K-FAC (Martens & Grosse 2015)
F
V C
Under the standard K-FAC independence assumptions (approximately independent a and g and certain factorization properties), this block is approximated as Kronecker product
Fw A
A A=
G
Elgg
(29)
23

---

Published as a conference paper at ICLR 2026
where A Rd and G moving) averages over bati-b
are the Kronecker factors maintained as running (exponential
Preconditioned update. Natural gradient descent would apply Fw approximation in Eq. 229 and the Kronecker identity (A x G) " A
to the gradient W. Using the
G-1 one obtains the
K-FAC preconditioned update for the weight matrix:
AW - G VWL A
W41 W. AW,
(30)
where 1) is the learning rate. In practice, A and G G-1 are not formed explicitly: K-FAC performs eigendecompossiiins
A = UAUAA
G G=
and then applies inverse (or inverse square-root) scalings in these eigen-bases. This requires storing the factors A, (and often their eigenvectors U.. Ug) and repeatedly computing or reusing their eigendecompositions.
Memory cost at LLM scale. Consider a typical transformer block with hidden size dia R doat s For modern LLMs, is in the range [4096, 4096,8192]. Each Kronecker factor A or G is then dense d x d matrix. It is important to note that, for numerical stability, K-FAC implementations typically store curvature factors in least FP32, even when the model itself uses FP16/BF16. A dense d FP32 matrix requires 4d bytes. For d = 4096, d = 4096² = 16,777,20 216 entries, which leads to
size of one factor (A or G) ~ 16,777,216 x 4 bytes = 64 MB.
Thus storing both A and G for weight matrix consumes about 2 x 64 MB 2 128 MB. A transformer block at this width typically has multiple large projection matrices, such as self-attention projections W W K, Wv, v, Wo plus two large forward matrices. Even if we conservatively apply K-FAC only to four matrices per block (e.g.. We, WK. Wv, Wo) and ignore the FFN, the curvature state per block is already
curvature per block R 128 MB 512 MB.
For a 7B LLM with roughly 80 transformer blocks, this yields
extra K-FAC curvature memory Rs 512 MB x 80 ss 40 GB per device,
only for storing A and G in FP32, without caching cigenvectors.
In practice, many K-FAC variants also cache eigendecompositions, i.e., U UG for each factor. Each eigenvector matrix U or UG is again d d FP32 matrix (another ~ 64MB for d = 4096), effectively doubling the curvature state:
curvature per weight (A, G, UA, Un) 4 x 64 MB = 256 MB,
curvature per block (4 weights) = 4 256 MB 1 GB,
curvature for 80 blocks x 80 GB per device.
Thus, for a realistic configuration (FP32 factors + cached eigen-bases), even a 7B model with 4096 requires the order of 40-80GB of additional curvature memory per device.
This curvature memory is replicated across data-parallel workers: each device maintains its own copy of the K-FAC state and participates in all-reduce operations aggregate the factors. The cost is therefore not amortized across 8 devices; it is incurred independently on each device.
Interaction with memory budget of 80GB. On devices with 80GB memory, RL training of LLMs already pushes device memory close to saturation due to:
model parameters (for a 7B model in FP16, parameters alone occupy ~ ~14-16GB),
optimizer states (Adam or AdamW typically add at least another ~224 parameter size, though sharding/ZeRO may partially mitigate this),
activations and KV caches from long-context rollouts (often tens GB for sequence lengths the thousands).
24

---

Published as a conference paper at ICLR 2026
Even under optimistic assumptions with aggressive activation checkpointing and optimizer sharding, reserving an extra 40-80GB purely for K-FAC curvature is incompatible with the 80GB memory budget. There is simply room left for long-context RL rollouts or for scaling to larger models.
Moreover, this overhead scales quadratically with the hidden size d. If we increase to = 8192 (typical of larger L.Ms),, then d? 8192 07,108,864 entries, which leads to
size of one FP32 factor = 67,108,888 104 x 4 bytes = 256 MB.
Repeating the above estimates, even storing only A and G (no cigenvectors) for four matrices per block across L blocks yields
extra curvature memory 44. 2d") tens 100 GB per device
for realistic depths and widths. Thus, at LLM scales, K-FAC curvature storage alone can easily demand 50-100GB or more per device, making it infeasible on current 80GB accelerators, especially in RI. settings where rollout activations are also resident in memory.
Computation cost. K-FAC's main computational bottleneck computing and updating the eigende- compositions of A and G for each layer. The complexity of eigendecompooitton for dense d x d matrix is (dd), and this dominates the cost forming the inverse (or inverse square-root) factors.
For d = 4096, d3 = 4096 = (68,,1,,4, (119,,76,33 736 R 6.9 10th FLOPs. Each K-FAC update of single factor (either A or G) therefore costs on the order of 10 floating-Pon point operations when accounting for constant factors. For four large matrices per block and L ss 80 blocks, full curvature refresh (updating both and for all K-FAC blocks) involves on the order of
factors) matrices) x x6.9 10
4.4 10 FLOPss
number of cigendecompssiiions
curvature update.
In classical applications of K-FAC, these expensive updates are amortized by refreshing curvature only every steps (e.g., J E [50, 200]) and reusing the same eigendecomposition in between. Even with such amortization, empirical reports on convolutional and recurrent networks show that K-FAC updates make each optimization step at least few times more expensive than first-order step when curvature refreshed regularly. At LLM scale, with many more and much wider layers, the (dd) factor makes this overhead more severe.
When we combine:
the (dd)) eigendecompositions required for each K-FAC factor,
the need to aggregate curvature statistics across data-parallel workers (extra communication),
the already high per-step cost of LLM RL training (due to long-context rollouts and large models),
a realistic deployment of K-FAC at LLM scale would very plausibly induce a three- to five-fold slowdown in effective optimization throughput compared to standard Adam or AdamW, even curvature is updated only every steps. Such slowdown, on top of the massive memory overhead outlined above, renders K-FAC effectively infeasible for RL training of modern L.MMs.
D.4.2 SHAMPOO
Shampoo Gupta et al. 2018 is second-order preconditioning method that keeps Kronecker-factored curvature statisties for each weight tensor, and then applies matrix inverse p-th roots of these statisties to precondition the gradient. We focus on the matrix case, which already captures the scaling issues LLM widths.
Consider matrix parameter W, E Rdouxda xda and its (per-minihatch) gradient
G, $ VW B Xdu
25

---

Published as a conference paper at ICLR 2026
Shampoo maintains two symmetric positive semidefinite (PSD) matrices per weight matrix,
L, edu.. ==1
Σ
G,G, E Rdaxda
(31)
R² cId, + E. G. e Rdadd,
(32)
s=1
where e > 0 is a small damping constant. In practice, L, and R, are updated by rank-dn and rank-dout increments of the form G.G, and G, G2 every optimization step.
The Shampoo update preconditions the gradient with inverse p-th powers L, and RŞ. For matrix parameter (order-2 tensor), the original analysis leads p 4:
1
VwL. L. G. R.
(33)
W W.- L..
(34)
where 7) > 0 the step size. The fractional powers are implemented via cigendecoopposiion: L, U, with A, diag(M ....., )da.), then
L, U U U. with A diag(A, A1
and similarly for R. For numerical stability, both the preconditioners and their eigen- decompositions are typically kept least 32-bit floating point precision, even when W, and G, are stored in FP16/BF16.
Memory cost at LLM scale. Assume transformer block where all large matrices have approximately square shape dn x dout x d, with d E [4096, 8192] typical for 7B-70B models. For each weight matrix W, Shampoo maintains:
Two curvature accumulators L,, R2 Rdxd;
In most practical implementations, the corresponding inverse fourth roots L, R. are also stored, to avoid recomputing matrix roots every step.
Thus, per weight matrix we have roughly four dense d matrices FP32:
#floats per curvature state zu 4d
memory per curvature state 4d 4 bytes 16d bytes.
For 4096, we have
d' = 4096 16,777,216 1.68 10?,
16d 2.68 1 bytes 256 MB.
So single large weight matrix requires on the order of Shampoo curvature memory per weight 256 MB.
A transformer block typically contains six large matrices (e.g., We, WK. Wv, Wo and two feed- forward matrices), so per block we obtain
curvature memory per block & 6 256 MB 1536 MB R 1.5 GB.
For 7B-scale model with 4096 and about Nutect 40 transformer blocks, the total Shampoo curvature memory device is
curvature memory per device 1.5 GB x Noteck & 1.5 GB 40 60 GB.
(35)
For a larger 13B-scale model with d zz 5120 and the same number of blocks, the d2 scaling yields d = 5120 curvature memory z 90-100 GB per device,
26

---

Published as a conference paper at ICLR 2026
and for even wider 70B-scale models with d = 8192, the full-matrix Shampoo preconditioners alone would require several hundred GB of memory.
Crucially, these curvature statistics optimizer state: in a standard data-paralle, RL fine-tuning setup without dedicated sharding of optimizer states (such as Distributed Shampoo), each device replica keeps its own copy of Le, R,, and their inverse roots for its local shard of parameters. This memory is in addition to:
Model parameters (often stored in FP16/BF16 together with first/second moment optimizer states),
Activations and attention KV caches required both for backpropagation and for collecting long- context trajectories,
The auxiliary models typically involved RLHF pipelines (e.g., reward/scoring models and reference policies), even in setups that do not maintain an explicit critic network.
Empirically, even first-order RLHF baselines (Adam/AdamW) already bring a 7B policy close to the 80 GB limit once the policy, reward/scoring model, and reference model are all active, especially with sequence lengths 2 1024 and realistic batch sizes. teck-of-the-envelope estimates and open-source RLHF reports indicate that 7B RLHF pipeline can easily consume 60-70 GB of memory on each device without any second-order optimizer states, Combining this with the ~ 60 GB of additional curvature memory estimated in Eq. (35) would clearly exceed the 80 GB device capacity. other words, full full-matrix Shampoo at LLM scale effectively leaves no headroom for rollouts, auxiliary models, or even storing the policy itself on a single device (80GB).
Computation cost. The two main sources of extra compute in Shampoo are:
Updating curvature accumulators L, and Ra;
Computing matrix inverse 1/4-powers L, and R,
(1) Curvature updates. For each weight matrix.
L L G.G., Forming the products G.G.T and G. G, costs
R2 1--R111 Do.I+G. G..
OGb d dadd Old²)
FLOPs when dm s dout s d. For a transformer with Notock blocks and roughly six large matrices per block, the per-step curvature update cost scales as
FLOP Pe Ceurv Notoct d3,
(36)
for some modest constant Caury (approximately O(10) when counting all GG T and G G computa- tions per block).
For d = 4096 and Nolock 80:
d 40963 = 68,711,476,,36 &s 6.87 x 10
so Eq. 360 gives
FLOPS Scurv ~ 10 80 x 6.9 x 1 5-7 10 FLOPs per optimization step,
just to update Shampoo's 's recond-moment statistics for the large matrices in the network.
(2) Inverse fourth roots. Computing L, and R, requires either:
Figendecomposition or SVD (apply - 11/4 to eigenvalues), or
Iterative snuarsetsquare-rooo schemes (e.g., Newton-Schulz),
both of which cost (dd)) per factor. One transformer block with six large matrices has twelve such
factors (L and R. for each weight), giving
FLOPs Spots,
mpots per update Cross 12 d
27

---

Published as a conference paper at ICLR 2026
where Croot depends solver details.
Root updates are typically amortized by refreshing them every T optimizer steps. With Notock blocks,
Croot X 12 x Nbluck X d3
FLOPs LLOOsS Syouis, step
(37)
T
For 4096, Notock = 80, and T = 100:
12 x 80 x 6.9 X 10th
FL.OP per step 2d
3 x 10 FLOPs per step.
100
The total extra work per optimization step therefore
FLOPSStampoo extra ~ FLOPSeurr + F-OP Syoots, per step
5 x 10
+ 3 x 10 = O(1018) FLOPs per step.
(38)
Relative slowdown from FLOPs. Large-scale RLHF pipelines for LLMs already require substantial per-step compute due to multiple forward/backwarc passes (policy, reference, reward/scoring, etc.) and long-context sequences. Eq. (38) shows that full-matrix Shampoo introduces additional ((1018) FL.Oss per optimization step, which is typically comparable to-or larger than-the cost of the remainder of the RL update.
Thus, even without invoking any specific algorithmic details, full-matrix Shampoo is expected to induce multi-x reduction optimization throughput solely from its second-order computations.
When combined with the ~ 60 GB curvature memory from Eq. (35). the method becomes impractical for RL training of LLMs on 80GB systems:
The curvature state alone exceeds the available memory once policy, reference, and reward models are included;
The extra (11033) FLOPs per step impose several-fold slowdown relative to standard first-order optimizers.
In short, full-matrix Shampoo cannot be used for RL training of modern LLMs currently available hardware.
D.4.3 SOAP
SOAP (Vyas et al. 2025) a second-order optimizer built on top Shampoo. For a fully-connected (or linear) layer with weight matrix W I2 d Xdi, and gradient
G
a V C E gduuxda
Shampoo maintains two curvature matrices that approximate second-moment information along the output and input dimensions:
L = 32 L,-1 B2) G.G.T,
(39)
R₂ = 32 R. + (1- 32) G2) G.T.G. G.,
(40)
where L, E Rdau and R, E dd are updated as exponential moving averages, and [3 E (0, 1)
a decay coefficient.
Preconditioned update in the eigenbasis. SOAP periodically (every t steps) computes cigendecom-
positions of the Shampoo preconditioners:
L Q Qu A Q
Q7 E Rauxdda
(41)
R = QR AR QT
QR E R dduud
(42)
where Ar. and An are diagonal matrices of eigenvalues, and Qr., Qn collect the corresponding eigenvectors. SOAP then rotates the gradient into this slowly changing eigenbasis:
G QT, G. Qu.
(43)
28

---

Published as a conference paper at ICLR 2026
and runs Adam-style first- and second-moment moment updates in the rotated coordinates:
M; B1 M (1 - 31) G
(44)
Vi B2 V (1 -1) O
(45)
where M. V E R dd 82 E (0. 1) are Adam Adam-style coefficients, and @ denotes element-wise multiplication. The preconditioned update in the eigenbasis is
U = M. 2 (vV, + te,
(46)
where @ is element-wise division and >> 0 is small numerical constant. Finally, SOAP rotates this update back to the original parameter space:
AW, = 17 Q2 U2 Q
(47)
W W++ = W, + AW,,
(48)
where 1) > 0 is the learning rate. Thus SOAP can be viewed as running Adam on a rotated version of the gradient, where the rotation is given by the Shampoo preconditioner eigenbasis.
Optimizer state and memory cost at LLM scale. Consider a transformer block where dout R dm R and the weight matrices are of size d d. At SOAP scale, the optimizer state associated with single such matrix W includes:
Shampoo curvature matrices L,, R. (each d x d),
cigenvector matrices Q2, QR (each x d).
rotated Adam moments M,, V2 (each d x d).
Altogether, this is six dense d x d matrices per weight matrix.
For numerical stability, these matrices are typically stored in least FP32, even when the model weights and activations are in BF16/FP16. A single dense d x d FP32 matrix requires 4d bytes. Therefore, the SOAP-related optimizer state per weight matrix is
bytes per weight (SOAP state) x 4d 24d bytes.
(49)
Let us instantiate this for a modern LLM width of 4096:
d² 40962 16,777,216,
24d 24 X 16,777,216 = 102,653,184 bytes.
Dividing by 1024 convert to MiB, have
102,653,184
SOAP state per weight =
= 384 MiB.
10242
Thus, each single 4096 4096 weight matrix carries roughly 384 MiB of SpAP-specific state.
typical transformer block at this width has at least four large projection matrices (for self- seff-attention: WQ. W K, Wv, Wo), not counting the feed-forward network. Even if we conservatively apply SOAP only to these four matrices, the curvature and moment state per block is
SOAP state per block x= 4 x 384 MiB = 1536 MiB s 1.5 GB.
(50)
For a 7 parametee LLM with roughly total SOAP state /= 1.5 GB x 80 = 120 GB per device,
80 transformer blocks, we obtain
(51)
only counting the FP32 matrices listed above, and ignoring any additional buffers or implementation overhead.
Crucially, this optimizer state is replicated across data- data-parallel devices: each worker maintains its own copy of L, R,, Q2, Qn, My V for its local parameters, and participates in all-reduce operations for gradient aggregation. The 120 GB figure in Eq. (51) is therefore per-device requirement; is not amortized across multiple devices.
Interaction with the memory budget in RL training, On devices with 80GB, RL training of LLMs already pushes device memory close saturation due
29

---

Published as a conference paper at ICLR 2026
model parameters (for 7B model in FP16, parameters alone occupy ~1 14-B6GGB. optimizer states (Adam or AdamW typically add at least another ~224x parameter size, though sharding/ZeRO may partially mitigate this),
activations and KV caches from long-contex, rollouts (often tens of GB for sequence lengths in the thousands).
Even under optimistic assumptions with aggressive activation checkpointing and optimizer sharding, is common consume on the order of 60-70 GB out of the 80 GB budget.
Adding the SOAP state from Eq. 51 would require around 120 GB per device purely for curvature and moment information, i...
60-70 GB
120 GB
180 GB per device.
existing RL pipeline
SOAP state
This exceeds the GOGB memory capacity by more than factor of two, even before accounting for safety margins and additional framework overhead. practice, there simply no configuration (batch size, sequence length, number of rollout trajectories) that allows both realistic RL training of a 7B LLM and full SOAP optimizer state to coexist on an 80 GB device.
Moreover, the SOAP memory overhead scales quadratically with the hidden size d. If we increase to 8192 (typical for larger LLMs), then
d² 81922 67,108,866
24d 24 x 67, 101,,18886 864 = 1,610,612,736 bytes ss 1536 MiB.
Thus, one 8192 x 8192 weight would carry about GB block many blocks would push the per-device optimizer state well beyond 200 GB. Therefore, at realistic LLM widths and depths, the SOAP memory requirements are incompatible with the fixed 80 GB budget in RL settings.
GB of SOAP state, and four such matrices per
Computation cost and slowdown in RI. SOAP inherits two major computational overheads:
periodic eigendecompositioon of L, and R, (every steps), and
per-step rotations of gradients and updates into and out of the preconditioner eigenbasis.
The eigendecompossition of a dense d d matrix has complexity O(d)) For d 4096,
d 4096 68,719, 088,719,476,0 476,776 6.9 x 10 FLOPs.
Each SOAP curvature refresh requires two such eigendecompasitions per weight (for L, and R2). so the cost per weight matrix is on the order of
FLOPs per weight (eigs) R 2 x 6.9 x 10 1.4x x 10t1
With four large matrices per block and L Rs 80 blocks, a full curvature refresh involves
FI.OPs per SOAP refresh (4 matrices)
(80 blocks) 1.4 x 10
x 4.5 x 100 FLOPs.
(52)
Even if this cost is amortized by updating the eigenbasis only every r = 100 steps, the amortized overhead is on the order of 4.5 x 10t1 Fl.OPs per training step, comparable to or exceeding the cost of the forward- backwadd pass itself for a 7B model at moderate sequence lengths.
In addition, at every step (not just every steps), SOAP performs the rotations
G' QJ G,QR. R
(53)
U = Q2 Q'O
(54)
which each involve two dense x matrix multiplications (left and right multiplication) and therefore have complexity (dd) per large weight matrix. For d 4096, these rotations add another substantial multiple of 6.9 x 10 FLOPs per weight per step.
When combined across all large matrices and blocks, these extra (dd)) operations typically make each SOAP step several times more expensive than a standard adam//Aamm step. At LLM scale, and especially in RL-style fine-tuning where:
30

---

Published as a conference paper at ICLR 2026
Table 5: Wall-clock time comparison between joint learning and CGPO extremely large LLMs.. Each device has 140GB.
Model
#Device
Method
50 steps (hours)
Per step (min)
16
Joint Learning
7.1
8.52
Qwen2.55-3BBInnttruet
CGPO
7.5
9.00
32
Joint Learning
12.6
15.12
Qwen2..5-7BBIInsrucc
32
CGPO
12.8
15.36
Table Performance of models (Qwen2.5-7B-Instrucce trained on the multi-domain dataset (math + code) with different methods, evaluated on multiple benchmarks. The bold font indicates the best result.
Math
Code Generation
Methods
AVG
MATH500
AMC
HumanEval
MBPP
FAMO
76.25
57.37
84.01
71.40
72.26
CGPO
76.15
60.81
84.66
72.60
73.56
rollouts require long sequences and sufficiently large batch sizes for stable training,
multiple model passes (policy, reference, reward/scoring, etc.) are performed per update. environment interaction and cross-device communication already contribute substantially to the per-step cost,
this optimizer overhead becomes dominant bottleneck. A conservative estimate is that SOAP would induce at least a 3-5x slowdown relative to AdamW; for long-conters RL training of large LLMs with many wide layers, the combined effect of repeated eigendecompositions and per-step rotations can easily push this into the 5-10x range in terms of effective tokens tokens-per-sedd per-second throughput.
Therefore, although SOAP attractive optimizer moderate scales, its quadratic memory footprint and cubic-time eigen-computation render it infeasible for RL training of modern LLMs on current hardwares.
E MORE EXPERIMENTS
E. TIMING EXPERIMENTS 32B AND 72B MODELS
Extremely large LLMs place substantial computational demand on rollout generation, since the cost of producing each token grows with model size. As model scale increases, rollout generation becomes the dominant component of end-to-end training time, while variations in gradient-update scheduling (e.g., sequentil updates vs. a single aggregated update) account for only a small fraction of the total compute.
To quantify this effect, we conduct timing experiments on two large models, Qwen2.5-32B Instruttt and Diwen2.5-72B-Instteee (140GB per device), respectively. For both models, we measure the total wall-clock time and average per-step time over the first 50 steps. These results provide a representative comparison of computational overhead under realistic large-scale RL training conditions.
The 32B and 72B experiments
clusters of 16 and 32 devices
across both model scales, the difference between joint learning and CGPO remains marginal relative to the overall training time. This supports the observation that, at extremely large scales, rollout generation dominates end-to-end runtime, and the additional gradient steps used in CGPO do not introduce a meaningful computational bottleneck.
As shown in Table 5
E.2 DISCUSSION ON THE SENSITIVITY TO THE NUMBER OF DOMAINS
In this section, we provide additional analysis and experiments regarding how CGPO behaves as the number of K increases. As discussed in the main paper, CGPO's effectiveness is driven primarily by
31

---

Published as a conference paper at ICLR 2026
Table Performance of models (Qwen2.5-7B-Instrus) trained on the multi-domain dataset (math + creative writing) with different methods, evaluated on multiple benchmarks. The bold font indicates the best result.
Math
Creative Writing
Methods
AVG
MATH500
AMC
WritingBench
74.85
54.72
CGPO
75.10
58,94
67.01
* 67.02
Table Performance of models (Qwen2.5-7B-Inss, -nssttutt trained on the multi-domain dataset (math + code + scientific QA + creative writing + logic + tabular) with different methods, evaluated on multiple benchmarks, The bold font indicates the best result.
Math
Code Gemeration
Scientific
Creative Writing
Logic
Tabular
Methods
AVG
MATH500
AMC
HumanEval
MBPP GPQA-diamond SuperGPQA
WritinBeenh
Zebra
HiTab
FAMO
75.30
55.02
82.93
68.60
22.64
31.58
63.09
36.84
68.71
56,
CGPO
74.90
59.84
83.88
70.80
26.91
31.72
65.08
37.63
69.57
*5787
the degree of cross-domaie conflict, rather than by K itself. Here, we elaborate this claim and present new experimental evidence.
Key Observation. CGPO's sequential curvature-informed mechanism is designed to mitigate cross- domain conflicts. Therefore, its benefit scales with how much the domains disagree. Across all experiments conducted--includingg those with substantially heterogeneous domain mixtures- did not observe any evidence of performance plateau or degradation as K increases.
Experimental Settings. To empirically verify this, condueted three groups of experiments, varying either the number of domains or the strength of cross-domain conflict:
1. Math + Code (moderate conffict): The datasets and evaluation benchmarks identical to those used in the main experiments.
2. Math Creative Writing (high conflict): The datasets and evaluation benchmarks are identical to those used in the main experiments.
3. Math + Code + Scientific QA + Creative Writing + Logic + Tabular:
The datasets and evaluation benchmarks for math, code, scientific QA, and creative writing follow the same setup as in the main experiments. For the logic domain, we train on Zebra Puzzle (1.3k samples) (Lin et al.)) and Ordering Puzzle (1.9k samples). and evaluate on the test set of Zebra Puzzle. For the tabular domain, train on HiTab (4.3k samples) (Cheng al. 2022) and evaluate on its test set. Both the logic and tabular training and test splits use the filtered versions provided by (Cheng et al. 2025). The reward functions for the logic and tabular domains are male-based.
We select FAMO for comparison in these experiments because it is the best-performing baseline at the 7B scale in our main experiments. All other training details, reward functions, and evaluation protocols follow the same setup in the main paper.
Results. Across all configurations, CGPO remains stable and effective, as shown in Tables 6 Importantly:
The performance improvement in the math creative writing setting (high conflict) is noticeably larger than in the math + code setting (moderate conflict), confirming our claim that CGPO's advantage grows as cross-domain conflict increases.
In the six-domain experiment, CGPO continues to deliver clear, consistent gains, showing that its benefits persist even when K becomes large and the domain mixture is highly heterogeneous.
These results confirm that CGPO's performance does not degrade the number of domains increases. Instead, its effectiveness is governed by the level of cross- domain conflict, and CGPO remains robust in large, diverse multi-domain training scenarios.
32

---

Published as a conference paper at ICLR 2026
F LIMITATIONS
Although CGPO demonstrates consistent performance improvements across multiple domains, several broader limitations remain. First, similar existing multi-domain RLALLMs approaches (Li et al. 2025a), we employ external reward models for certain domains, which may themselves be constrained by current LLM-based evaluation paradigms. For instance, in the creative writing domain, using single LLM-as MMMMaas-a as-a-judge may introduce stylistic biases that reflect the limits of automated evaluation. Second, similar to existing studies (Cheng et al.) 2025 Li al. 2025b), the overall effectiveness depends on the coverage and granularity of domain-specifie rewards, and future advances in reward modeling may naturally enhance performance. Finally, while the randomized sequential update scheme encourages cross-domain interaction, exploring more elaborate scheduling strategies or structured coordination mechanisms remains an open direction for future work. We view these limitations reflecting broader challenges shared across current RLALLMs research, and hope that our work can contribute to the community's continued progress on addressing them.
33


Reviews:
[{"review_id": "MADHH2GHIK", "reviewer": "Reviewer_f3Af", "summary": "This paper studies multi-domain optimization in RL. The authors propose a simple training design: at each (mini-batch) step, prompts are randomly sampled from all tasks. A random permutation of the task order is generated, and the model performs sequential updates on each task according to this permuted order. They show that the sequential update scheme has a theoretical connection to Newton's method where sequential updates induce cross-domain Hessian–gradient interactions, which encourages cooperation across domains  and steers optimization toward coordinated improvements over all tasks.", "questions": "NA", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "* Originality: The authors establish a novel connection between task-wise sequential updates within a mini-batch in GRPO and Newton's method and Hessian approximation. This insight and perspective are neat and enlightening. \n* Quality/Clarity: The entire paper is very well written, clearly structured, and easy to follow. Overall, it demonstrates high technical and presentation quality.", "weaknesses": "* The paper focuses on optimizing the average task accuracy. But in practice, some tasks may be more critical than others, and the proposed method does not have a mechanism to account for task-specific importance or weighting. Extending the framework to handle non-uniform task importance could make the approach more broadly applicable.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "kC7IfwNX41", "reviewer": "Reviewer_Fgcd", "summary": "1．\tThis paper presents Curvature-Guided Policy Optimization (CGPO), a novel training framework designed to enhance the multi-domain reasoning capabilities of Large Language Models (LLMs) using reinforcement learning (RL). \n2．\tThe core challenge addressed is the presence of cross-domain conflicts, where improvements in one domain (e.g., coding) can lead to degradation in another (e.g., creative writing). \n3．\tCGPO is inspired by the geometric principles of Newton's method but avoids its prohibitive computational cost by using a randomized sequential update mechanism. This mechanism implicitly approximates Hessian-gradient interactions across domains, fostering gradient alignment and steering the model towards parameters that perform well across all domains simultaneously. \n4．\tThe method is evaluated on a diverse dataset spanning math, coding, science, and creative writing, and demonstrates superior performance over several strong baselines.", "questions": "1.How sensitive is CGPO to the number of domains (K)? Does performance plateau or degrade as K becomes very large?\n 2. The method relies on multiple gradient steps per \"outer\" iteration. While the overhead is shown to be low, are there scenarios (e.g., with extremely large models) where even this could become a bottleneck compared to a single-update method?", "limitations": "", "rating": 4, "confidence": 5, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1.  The core insight of leveraging the geometric structure of the reward landscape (inspired by Newton's method) to mitigate multi-domain conflicts is both innovative and principled. The motivation is clearly laid out, connecting the limitations of existing first-order and gradient-manipulation methods to the potential benefits of second-order information.\n2.  A significant contribution is the design of a method that captures the benefits of curvature information without explicitly computing the Hessian. The proposed mechanism of randomized sequential updates is elegant and computationally lightweight, introducing only negligible overhead compared to standard joint training. This makes CGPO highly practical and scalable for large-scale LLM RL training, a critical consideration in this field.\n3.  The paper provides a solid mathematical foundation for the proposed method. The derivations (in the main text and appendix) convincingly show how the sequential update scheme approximates cross-domain Hessian-gradient products and, in expectation, encourages gradient alignment. This moves beyond a purely empirical contribution.\n4.  The choice of four distinct domains (math, code, science, creative writing) effectively demonstrates the method's ability to handle varied types of reasoning and reward signals.\nThe comparison includes a representative set of baselines (joint learning, curriculum learning, gradient balancing), showing that CGPO advances the state-of-the-art.\n5．Evaluating on both 3B and 7B parameter models strengthens the claims, showing that the benefits are consistent and potentially scale with model size. Using seven established benchmarks provides a robust assessment of multi-domain capability.\n6. The ablation studies on domain order randomization and the mixing coefficient `α` are crucial for validating the design choices. The analysis of computational overhead directly addresses a key potential concern for adoption.", "weaknesses": "1.  The role and intuition behind the final interpolation step (`θ_new ← φ_0 + α(φ_K - φ_0)`) could be explained more clearly in the main text. While the ablation study shows its importance, a more intuitive explanation of why this interpolation stabilizes training would be helpful. Is it primarily a form of trust region enforcement or a step-size damping mechanism?\n2.  The paper would be strengthened by a more explicit discussion of its limitations. For instance, the creative writing reward is based on a single LLM-as-a-judge (Qwen2.5-72B). The potential biases of this judge and how they might influence the overall multi-domain optimization are not discussed.\n3.  While the baselines are well-chosen, the related work section mentions other second-order approximation methods like Shampoo, K-FAC, and SOAP. A discussion on why CGPO is a more suitable path for the RL-for-LLMs setting compared to adapting these existing approximate second-order methods would provide valuable context and further justify the novelty. The argument is implicit but could be made more explicit.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "pWNwZwObth", "reviewer": "Reviewer_SFav", "summary": "This manuscript introduces a multi-domain post-training algorithm for LLMs. This method is motivated by Newton's method, where the Hessian matrix can reshape the objective landscape and accelerate convergence. In practice, the authors propose that sequentially training on each domain in a randomized order will naturally simulate the cross-domain Hessian-gradient interaction, and thus is better than joint training on all the domains.", "questions": "* Please refer to the weaknesses.\n* I wonder whether the proposed mechanism can be utilized for pre-training of LLMs, because typically the pre-training corpus is also multi-domain. The reviewer would love to hear some discussions on this point, and what is even better, some experimental results.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "* The algorithm is simple in practice and grounded in well-motivated theory.\n* The paper is well written; the authors clarify their ideas with a thorough background and intuitive discussion.", "weaknesses": "* The author has some concerns about Equation (5) and the discussion following it. The authors claim that \"since $\\sigma$ is randomized, in expectation every pair will contribute equally\". However, for each **specific** training iteration, the permutation is fixed, which means half of the $(i,j)$ pairs will not be present. Therefore, it is not an unbiased estimation of the Hessian matrix, and it is not clear that this effect will cancel out across multiple iterations. The reviewer would like to hear more analysis and discussions on this point. \n* Regardless of the above point, let's accept the proposal that the training mechanism in the manuscript can simulate the cross-domain Hessian-gradient interaction. However, it is still unclear why joint-training on all domains cannot. Of course, there are no such interactions for a mini-batch, but there might be some cross-domain interaction stemming from two consecutive mini-batches. To convince the readers of the superiority of the training mechanism proposed in this paper, the authors should provide some analysis of the trivial joint-training mechanism in addition to experimental results.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)