SCALE: Scalable Cross-Attention Learning with Extrapolation for Agentic Workflow Scheduling

📄 arXiv: 2606.06820v1 📥 PDF

作者: Zhifei Xu, Jierui Lan, Zixuan Liang, Aiji Liang, Jinxi He

分类: cs.LG, cs.AI

发布日期: 2026-06-05

备注: Submitted to Computer Networks


💡 一句话要点

提出SCALE以解决异构集群调度的可扩展性问题

🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture) 支柱九:具身大模型 (Embodied Foundation Models)

关键词: 深度强化学习 调度算法 异构集群 交叉注意力 结构化表示正则化 规模泛化 资源优化

📋 核心要点

  1. 现有的深度强化学习调度器在集群规模变化时需要重新训练,限制了其灵活性和可扩展性。
  2. SCALE通过交叉注意力指针网络设计,能够在不同规模的集群上进行有效调度,无需微调。
  3. 在实验中,SCALE在48个节点上相较于未使用SRR的架构平均响应时间减少了8.9%,显示了其优越性。

📝 摘要(中文)

代理大型语言模型(LLM)系统将复杂任务分解为工作流有向无环图(DAG),其原语必须在异构集群上调度。现有的深度强化学习(DRL)调度器固定于特定集群规模,且在服务器数量变化时需重新训练。我们提出了SCALE(可扩展的交叉注意力学习与外推),这是一种DRL调度器,能够在不进行微调的情况下推广到未见过的集群规模。SCALE采用交叉注意力指针网络,任务特征与服务器特征进行查询,从而使得该架构在构建上能够接受任意数量的服务器。然而,我们观察到,仅有置换不变的架构并不能保证在新规模下的良好性能,因此我们引入了结构化表示正则化(SRR),通过去相关损失结合KL惩罚,保持特征统计在输入规模变化时的稳定性。在16个节点上训练,并直接在32和48个节点上测试,SCALE在N=48时相较于没有SRR的同一架构平均响应时间减少了8.9%,确认了显式正则化在缩小规模泛化差距中的必要性。

🔬 方法详解

问题定义:本论文旨在解决现有深度强化学习调度器在集群规模变化时需要重新训练的问题,这限制了其在动态环境中的应用。

核心思路:SCALE的核心思路是通过交叉注意力机制,使得调度器能够在不同规模的集群上进行有效调度,而无需进行微调。

技术框架:SCALE的整体架构包括交叉注意力指针网络,任务特征与服务器特征进行查询,此外还引入了结构化表示正则化(SRR)以稳定特征统计。

关键创新:最重要的创新点是引入了结构化表示正则化(SRR),通过去相关损失和KL惩罚,解决了特征统计在集群规模变化时的不稳定性问题。

关键设计:在网络结构上,SCALE采用了交叉注意力机制,损失函数中结合了去相关损失和KL惩罚,以确保特征统计的稳定性,适应不同的集群规模。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

在实验中,SCALE在16个节点上训练,并在32和48个节点上进行测试。结果显示,在48个节点上,SCALE相较于未使用SRR的架构,平均响应时间减少了8.9%,证明了其在规模泛化方面的有效性。

🎯 应用场景

该研究具有广泛的应用潜力,尤其在云计算和大规模分布式系统中,可以有效优化资源调度,提高系统的响应速度和效率。未来,SCALE可能会在智能调度、自动化运维等领域发挥重要作用。

📄 摘要(原文)

Agentic Large Language Model (LLM) systems decompose complex tasks into workflow Directed Acyclic Graphs (DAGs) whose primitives must be scheduled on heterogeneous clusters. Existing deep reinforcement learning (DRL) schedulers are tied to a fixed cluster size and require retraining whenever the number of servers changes. We propose SCALE (Scalable Cross-Attention Learning with Extrapolation), a DRL scheduler that generalizes to unseen cluster scales without fine-tuning. SCALE employs a cross-attention pointer network where task features query against server features, so the architecture accepts any number of servers by construction. We observe, however, that permutation-invariant architecture alone does not guarantee good performance at new scales - the attention feature undergoes distribution shift as the server count grows. To counter this, we introduce Structured Representation Regularization (SRR): a decorrelation loss combined with a KL penalty toward the standard normal, which keeps feature statistics stable regardless of input size. Trained on 16 nodes and tested directly on 32 and 48 nodes, SCALE reduces average response time by 8.9% at N=48 relative to the same architecture without SRR, confirming that explicit regularization is necessary to close the scale-generalization gap.