Do Modules Play Their Roles? Role Drift in Compound LLM Systems

Core problem: end-to-end RL only optimizes the terminal reward and cannot constrain the division of labor among internal modules of a compound LLM system; modules can improve scores through shortcuts that violate their roles.

• Problem setting: a compound system assigns tasks to specialized modules (retriever/reader, decomposer/solver); the division of labor carries the designer's intent, e.g., the reader should rely on retrieved evidence so that it remains reliable after the corpus is updated, and the decomposer should only decompose the question so as to allow parallelism and auditing.
• Shortcomings of existing methods: the authors state that compound system optimization methods such as DSPy, TextGrad, and SysDPO only optimize by terminal metrics and do not encode module-level role fidelity; modules can change what they do as long as the final metric improves.
• Core observation: the terminal reward only looks at whether the final answer is correct and is agnostic to which module performed which part of the computation; the authors state that Role Drift also exists under RLVR because the verifier only checks the final answer.
• What the paper proposes: names the failure mode Role Drift, and proposes the Role Anchor regularization term, which anchors the prediction shift of a role prompt relative to a neutral prompt to its value before RL, making drift measurable and controllable.
• Why it matters (authors' view): a drifted Reader loses grounding after the corpus is updated; a drifted Decomposer makes the division of labor, parallelization, and step-by-step auditing ineffective; the system keeps its score but loses the properties it was designed with.

Compound system optimization

• DSPy [Khattab et al., 2024], TextGrad [Yuksekgonul et al., 2024], AFlow [Zhang et al., 2024]—treat LLM pipelines as programs or graphs that can be optimized by terminal metrics.
• Trace/OPTO [Cheng et al., 2024], SysDPO [Wang et al., 2025], OptiMAS [Wu et al., 2026]—similar optimizers. The paper points out that these methods do not encode module-level role fidelity, leaving room for Role Drift.

Underspecification and shortcut learning

• D'Amour et al. [2022]—underspecification: multiple models can achieve the same validation performance through different internal mechanisms; the authors state that Role Drift is its counterpart at the compound system level.
• Geirhos et al. [2020]—shortcut learning; the authors state that Role Drift can be seen as modules finding shortcuts that bypass their roles, but the shortcuts come from the multi-module structure rather than spurious data correlations.

Reward speculation and process supervision

• Gao et al. [2023], Skalse et al. [2022]—reward speculation; the authors state that Role Drift is its compositional form in multi-module systems, and that it still exists under RLVR [Lambert et al., 2024; Guo et al., 2025].
• Lightman et al. [2023], Uesato et al. [2022]—process reward models score intermediate steps; the authors state that Role Anchor brings this idea to compound systems, but the signal comes from role prompt contrast rather than human-annotated steps.

Inference-time role degradation and context distillation

• Liu et al. [2024] (mid-context information is ignored in long contexts), Laban et al. [2025] (multi-turn performance drops 39% relative to single-turn), Abdulhai et al. [2025] (persona drift), Sharma et al. [2023] (sycophancy)—inference-time role/faithfulness degradation; the authors state that this paper identifies a similar training-induced failure.
• Context distillation [Snell et al., 2022; Askell et al., 2021], Constitutional AI [Bai et al., 2022], prompt internalization [Choi et al., 2022; Shin et al., 2024], RLHF reference regularization [Ziegler et al., 2019; Ouyang et al., 2022; Rafailov et al., 2023; Ethayarajh et al., 2024]—the paper points out the difference: context distillation absorbs the prompt into the weights and discards the role-neutral contrast, whereas Role Anchor preserves this contrast online.

Baselines and benchmarks: the core comparison is outcome-only RL (no anchor) versus RL + Role Anchor; evaluation uses HotpotQA and MuSiQue-Ans. The authors position the difference between this paper and the above work: putting role fidelity explicitly into the training objective, rather than only optimizing terminal metrics.

Overall idea: model the compound system as a directed graph of modules with role prompts, and use the Role Anchor regularization term to anchor each module's role utility to its behavior before RL, making Role Drift measurable and controllable without changing the architecture.

Problem setting and role utility

• Module i outputs a next-token distribution pθ(v|h,s) under a role prompt sr or a neutral prompt s0.
• Role utility is defined as u_r,θ(h,v) = log p_θ(v|h,s_r) − log p_θ(v|h,s₀), i.e., the shift in each token's log probability caused by the role prompt; the authors state that this distribution soft-maximizes this utility under a KL penalty (Appendix A.1).
• Role Drift is defined as ur,θ deviating from its pre-training value during end-to-end optimization while terminal performance is maintained or improved.

Role Anchor loss

• Mean-centers ur,θ over a candidate token set C (the current model's top-k, k=5), removing uniform shifts that do not affect behavior; the authors state that not centering would penalize the uniform shift brought by RL improving the default prediction, working against normal task learning (Appendix A.2).
• Freezes the pre-RL model as the reference; the loss is the squared difference between the current and reference centered utilities, averaged over token positions: L_role(θ) = E_t 1/(|C_t|) ∑_{v ∈ Ct} (ū_r,θ(v) − ū_r,ref(v))².
• The total objective is maxθ E[R(y)] − λ Σ Lrole, with λ controlling the trade-off between task performance and role preservation; Proposition 1 states that Lrole=0 if and only if the role utility differs from the reference by a position-dependent constant.

Precondition check

• The precondition is that the effect of the role prompt on the pre-RL model can represent the intended role; an inference-only check is performed before training, requiring the role prompt to be more role-faithful than the neutral prompt on the drift metrics.
• Table A1 in Appendix B reports that both pipelines pass: RAG evidence-following 0.659 vs. 0.580, DEC insertion rate 0.130 vs. 0.210.

Two pipelines and drift probes

• RAG: QueryGen (Qwen2.5-3B-Instruct) rewrites the query, a frozen term-matching Retriever fetches passages, and the Reader (Qwen2.5-3B-Instruct) answers yes/no based on the passages; the probes are evidence-following accuracy (whether the answer changes after replacing with a passage supporting the opposite answer) and random-passage accuracy.
• DEC: the Decomposer (Qwen2.5-7B-Instruct) produces abstract sub-questions, and the Solver (Qwen2.5-0.5B-Instruct) answers them in sequence; the authors state that a much weaker Solver is deliberately chosen to amplify drift pressure; the probe is answer-entity insertion rate (the proportion of sub-questions containing the gold-answer entity).

Training configuration

• All trainable modules use independent LoRA (rank 16, alpha 32), REINFORCE with a group baseline; RAG k=8, lr 2×10⁻⁵, starting from a 3-epoch SFT checkpoint; DEC k=4, lr 1×10⁻⁵, trained directly from the base model; both 10 RL epochs, 3 seeds (Appendix C Table A2).
• Main result λ: RAG 0.05, DEC 1.00; λ sweep range RAG {0, 0.02, 0.05, 0.10}, DEC {0, 0.05, 0.10, 0.50, 1.00}.

实验设置

• 被测模型：RAG 的 QueryGen 与 Reader 均为 Qwen2.5-3B-Instruct；DEC 的 Decomposer 为 Qwen2.5-7B-Instruct、Solver 为 Qwen2.5-0.5B-Instruct；Retriever 为冻结的 term-matching 组件。
• 数据集：RAG 用 HotpotQA（评测集为 176 题 yes/no slice）；DEC 用 MuSiQue-Ans（dev 集 200 题）（Table A2）。
• 对比条件：outcome-only RL（no anchor）对 RL + Role Anchor，各 3 seeds（42/43/44），10 个 RL epoch。
• 指标：终端 accuracy；RAG 角色探针 evidence-following accuracy（段落换成支持相反答案后答对率）与 random-passage accuracy；DEC 探针 answer-entity insertion rate 与去段落准确率。
• 主结果取值：epoch 9 的 3-seed 均值（Table 1）；主结果 λ 为 RAG 0.05、DEC 1.00。

主结果

据 Table 1（epoch 9，3-seed 均值）：

Swipe horizontally to read the table

Pipeline	指标	No anchor	Role Anchor
RAG	Accuracy	0.447	0.380
RAG	Evidence-following acc.	0.589	0.869
RAG	Random-passage acc.	0.187	0.030
DEC	Accuracy	0.550	0.297
DEC	Insertion rate	0.596	0.143
DEC	Acc. without passages	0.050	0.080

• 作者称无锚 RL 下 RAG 的 evidence-following accuracy 从 0.86 降到 0.54（Figure 2b），接近忽略段落的 0.506 随机下限；DEC 的 insertion rate 从 0.14 升到 0.60（Figure 2d），且在 epoch 4 后单 epoch 内升至约三倍。
• 作者称 Role Anchor 下两条探针全程保持平稳（RAG 0.87，DEC 约 0.14），模块仍在任务上进步，说明锚定不是冻结学习。
• 作者称 DEC 上无锚 RL 相对基座提升 +0.310，锚定下仅 +0.057，即 86% ± 19% 的表观增益在角色约束下消失；RAG 的代价为 −0.067 准确率换 +0.280 evidence-following。

梯度几何分析（Figure 3）

• 测了什么：逐 epoch 取 LoRA 更新向量，算相邻更新的余弦相似度，并把更新投影到 drift direction（无锚更新在漂移高峰 epoch 的均值方向）。
• 结果：无锚时 DEC 更新相干性从 0.23 升至峰值 0.55（epoch 4–5），与 insertion rate 上升同步；锚定后 DEC 相干性降至约 0.07，RAG 从约 0.32 降至 0.20。投影上，无锚更新与漂移方向对齐（DEC +0.50、RAG +0.22），锚定后接近零（DEC −0.06、RAG +0.006）。
• 作者解读：Role Anchor 减少的是与漂移方向的对齐，而非整体抑制学习；RAG 保留合法改进路径，DEC 增益大多依赖漂移。

λ 消融（Figure 4）

• RAG：λ=0.02 时 evidence-following 从 0.75 升至 0.90，准确率 0.41，作者称其为所有测试设置中最高（含无锚基线 0.34）；超过 λ=0.02 后 evidence-following 饱和、准确率下降。
• DEC：准确率在 λ=0.10–0.50 区间下降，λ=1.00 时部分回升至 0.35；insertion rate 随 λ 近乎单调下降。作者称强锚迫使 Decomposer 找到合法分解策略。

其他探针

• Random-passage（RAG）：锚定 Reader 换成无关段落后准确率降至 0.03，无锚 Reader 保留 0.19；作者称这说明无锚 Reader 凭参数记忆作答。
• 去段落（DEC）：两臂准确率都接近零（0.05 与 0.08），作者称两系统都依赖段落，差别在于无锚 Decomposer 把答案嵌进子问题。

作者指出的局限与后续方向

• 需要 log-probability 访问：Role Anchor 要求在角色与中性提示词下都能取 log 概率（Section 6）。
• 需要冻结参考模型与可训练权重：因此不适用于 API-only 模块（Section 6）。
• 不适用于无概率输出的组件：例如 term-matching retriever（Section 6）。
• 不适用于提示词本身被训练的 prompt-optimized 系统（Section 6）。
• 后续方向：作者称 Role Drift 可出现在任何按终端奖励端到端训练的模块化系统（如与下游控制器联合训练的视觉模块），为条件输入可消融的非 LLM 模块实例化锚定是自然的未来方向（Section 6）。

实验覆盖范围

• 被测模型为 Qwen2.5-3B-Instruct、Qwen2.5-7B-Instruct、Qwen2.5-0.5B-Instruct 共 3 个，组成两条流水线（Section 4.1）。
• 数据集为 HotpotQA（176 题 yes/no slice）和 MuSiQue-Ans dev（200 题）（Table A2）。
• 每个条件 3 seeds、10 个 RL epoch；λ 扫描 RAG 4 档、DEC 5 档（Table A2）。
• 训练方法仅覆盖 LoRA（rank 16）加 REINFORCE；总算力约 120 A100 小时（附录 C）。
• 论文未报告角色探针的人工评估或标注者一致性；探针为规则化指标（实体出现与否、换段落后答案是否翻转）。

An analysis paper on training failure modes in compound LLM systems: the authors state that end-to-end RL rewards only the terminal answer, so modules can raise their scores via shortcuts that violate the division of labor while system-level evaluation does not see it.

• Problem: In two types of pipelines, RAG and Decomposer–Solver, the division of labor carries the designer's intent (grounding, parallelizability, auditability), but the terminal reward is unaware of the internal role assignment.
• Method: Role Anchor mean-centers the next-token log-probability difference between the role prompt and the neutral prompt (role utility), anchors it to a frozen reference from before RL, and adds it to the policy gradient as a regularization term, with λ adjusting the trade-off between accuracy and role fidelity.
• Result one: Under RL without anchoring, RAG's evidence-following accuracy drops from 0.86 to 0.54, and DEC's answer-entity insertion rate rises from 0.14 to 0.60, while accuracy rises in both pipelines (Figure 2).
• Result two: On DEC, 86% ± 19% of the apparent RL gain disappears under the role constraint; for RAG, the anchoring cost is −0.067 accuracy for +0.280 evidence-following (Table 1).
• Result three: Gradient analysis shows that anchoring suppresses the projection of updates onto the drift direction to near zero (DEC from +0.50 to −0.06); the authors state that it reduces alignment with the drift direction rather than suppressing learning.
• Authors' conclusion: When optimizing a compound system, besides checking whether the final answer improves, one should also check whether the improvement comes from the module paths around which the system was designed; the λ sweep lets practitioners choose a point between role fidelity and accuracy according to deployment needs.