GauS: Differentiable Scheduling Optimization via Gaussian Reparameterization¶

会议: NeurIPS 2026
arXiv: 2602.20427
OpenReview: https://openreview.net/forum?id=QfSSTcTPqe
领域: 优化/理论
关键词: 算子调度, 可微组合优化, 高斯重参数化, 模调度, 增广拉格朗日

一句话总结¶

GauS 提出用一维高斯连续分布对算子执行时间步进行松弛建模，保留时间序列的天然序数邻近先验并将优化参数量压缩至 $O(|V|)$，结合增广拉格朗日法与 GPU 向量化首次实现了流水线/模调度的高效端到端可微求解。

研究背景与动机¶

算子调度（Operator Scheduling）是软件编译指令重排与硬件高层综合（HLS）中的经典 NP-hard 核心问题，其本质是将计算图中的每个操作节点映射到离散的时间步上，在严格满足数据依赖、延迟界限和物理资源容量的前提下最小化执行开销或峰值资源占用。随着大模型和异构加速器（如 Groq 流式张量处理器及存内计算 PIM）的爆发式发展，调度图规模迅速由几百个节点扩展至数万个节点，现代系统迫切需要在极短时间内对海量异构功能单元完成精细的时空资源编排。

长久以来，该领域深陷于两难境地：以整数线性规划（ILP）和 SMT 求解器为代表的精确算法虽能保证最优解，但其指数级复杂度在面对中大规模真实图时经常超时崩溃甚至无法给出可行解；而以列表调度（List Scheduling）和力导向调度（FDS）为代表的贪心启发式算法虽然执行迅速，却因缺乏全局视野在应对多重竞争约束（如内存占用峰值与通信开销）时极易陷入次优局部解。近年来兴起的可微组合优化虽通过松弛离散决策为连续空间引入了基于梯度的全局搜索，但开创性工作 GS-Schedule 依赖 Gumbel-Softmax 分类分布建模，将调度步骤视为无序互斥的 One-Hot 标签，完全忽视了时间维度的序数特性（第 $d$ 步与 $d+1$ 步在优化空间中与 $d+100$ 步等距），导致参数规模随深度与节点数恶化为 $O(D \cdot |V|)$，在面对深层大图时频发显存爆炸且 GPU 利用率极低。

本文的切入角度是：算子在时间轴上的分布具有天然的局部聚集性与序数平滑性，若用连续单峰概率分布替代离散类别分布，不仅能大幅压缩搜索空间，还能利用概率密度的平滑输运引导梯度更新。核心 idea：将每个算子的调度时刻参数化为独立高斯分布 $X_i \sim \mathcal{N}(\mu_i, \sigma_i^2)$，将参数量大幅削减至 $2|V|$ 并天然保留时间邻近先验，通过闭式积分与 LogSumExp 统一表达依赖、内存与循环流水约束，配合增广拉格朗日法在 GPU 上实现全向量化可微求解。

方法详解¶

整体框架¶

GauS 的全流程以计算 DAG 为输入，将离散调度映射转化为针对每个节点的均值向量 $\boldsymbol{\mu} \in \mathbb{R}^{|V|}$ 和标准差向量 $\boldsymbol{\sigma} \in \mathbb{R}^{+|V|}$ 的连续约束优化问题。在初始化阶段，框架根据静态 ASAP 和 ALAP 计算各节点自由度赋予先验分布；在连续优化循环中，GPU 向量化前向计算目标函数（如峰值资源或通信开销）与依赖约束期望违规量，通过增广拉格朗日乘子动态调整惩罚力度并执行反向传播；在每轮迭代末尾，通过对连续均值取整提取离散调度，并在发现细微因果冲突时触发拓扑合法化算子安全修正，直至满足收敛条件。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入计算图 DAG 与最大延迟界限 D"] --> B["高斯重参数化初始化<br/>均值置中 + 标准差对齐 ASAP/ALAP 自由度"]
    B --> C["连续概率密度与单位区间离散对齐<br/>积分计算单步占用概率 P_i^d"]
    C --> D["可微目标与违规期望建模<br/>LogSumExp 峰值松弛 + 依赖/循环期望"]
    D --> E["增广拉格朗日优化与反向传播<br/>动态更新拉格朗日乘子与惩罚项"]
    E --> F["离散调度采样与拓扑合法化<br/>四舍五入整型映射 + 贪心无环微调"]
    F -->|未收敛且在预算内| C
    F -->|收敛退出| G["输出满足硬约束的全局优化调度方案"]

关键设计¶

1. 高斯重参数化与序数时间先验建模：从 $O(D|V|)$ 降维至 $2|V|$

传统类别松弛为每个节点分配长为 $D$ 的概率向量，各时间步在几何上相互正交，无法体现时间前后的连续过渡。GauS 针对性地为每个算子 $v_i$ 引入标量均值 $\mu_i$（期望执行步骤）和标准差 $\sigma_i$（松弛不确定度），形成随机变量 $X_i \sim \mathcal{N}(\mu_i, \sigma_i^2)$。优化初期 $\sigma_i$ 保持较大以充分探索解空间，随着训练收敛 $\sigma_i$ 自动收缩迫使概率质量凝聚向离散整数点。为了保持连续优化与离散舍入评估的一致性，算子 $v_i$ 落入离散时间步 $d \in \{0, \dots, D-1\}$ 的概率被定义为高斯概率密度在 $[d-0.5, d+0.5]$ 上的积分： $$P_i^d = P\left\{\lfloor X_i + 0.5 \rfloor = d\right\} = \Phi\left(\frac{d+0.5-\mu_i}{\sigma_i}\right) - \Phi\left(\frac{d-0.5-\mu_i}{\sigma_i}\right)$$ 其中 $\Phi(\cdot)$ 为标准正态累积分布函数（CDF）。此设计使得向邻近时间步微移 $\mu_i$ 时能够产生连续平滑的概率质量输运，梯度能够自然拉动算子向其在时间轴上的可行邻域移动，同时全局参数量严格锁死在 $2|V|$，不再随调度深度膨胀。

2. 任意目标与约束的连续期望闭式转化：支持平滑流水线/模调度

在建立 $P_i^d$ 概率表示后，离散的硬性图约束与非线性硬件目标被系统化转化为关于 $(\boldsymbol{\mu}, \boldsymbol{\sigma})$ 的可微期望函数。针对数据前向依赖约束（生产者 $v_i$ 必须先于消费者 $v_j$ 执行），违反约束的期望次数表示为所有逆序组合概率的二次求和： $$\widehat{\mathcal{V}}_{Dep} = \sum_{(v_i, v_j) \in E} \sum_{d_i=1}^{D-1} \sum_{d_j=0}^{d_i-1} P_i^{d_i} \cdot P_j^{d_j}$$ 针对硬件存储单元占用（寄存器/SRAM生命周期覆盖从算子启动到其所有后继启动的最大跨度），时间步 $d$ 上的期望存储压力表示为： $$\widehat{Reg}(d) = \sum_{v_i \in V} \left(\sum_{d_i=0}^d P_i^{d_i}\right) \cdot \left(1 - \prod_{v_j \in succ(v_i)} \sum_{d_j=0}^d P_j^{d_j}\right) \cdot b_i$$ 全局最大资源/内存峰值通过温度系数为 $\tau$ 的 LogSumExp 平滑算子逼近 $\widehat{\mathcal{L}}_{Reg} \approx \tau \cdot \log\left(\sum_d \exp(\widehat{Reg}(d)/\tau)\right)$。更进一步，对于高吞吐循环流水中的模调度（Modulo Scheduling，启动间隔为 $II$），GauS 首创将各时间步对 $II$ 取模进行折叠求和，通过 $P\{\lfloor X_i+0.5\rfloor \equiv t \pmod{II}\} = \sum_k P_i^{t + k \cdot II}$ 精准表达模资源碰撞与跨迭代回边（Recurrence Edge）循环时序约束违规，填补了可微模调度的理论空白。

3. 增广拉格朗日求解与轻量级拓扑安全网合法化：兼顾探索自由度与硬可行性

静态损失加权常导致硬约束欠满足或梯度主导失衡。GauS 采用增广拉格朗日方法（ALM），构造联合目标： $$\mathcal{L}_{total} = \widehat{\mathcal{L}} + \sum_i \left(\lambda_i \widehat{\mathcal{V}}_i + \frac{\rho}{2} \|\widehat{\mathcal{V}}_i\|^2\right)$$ 乘子 $\lambda_i$ 随违规量按规则 $\lambda_i \leftarrow \lambda_i + \rho \widehat{\mathcal{V}}_i$ 动态递增，自适应强化硬约束权重。在提取离散解 $s_i = \text{round}(\mu_i)$ 时，若舍入导致因果冲突，并不依赖优化器反复盲试，而是启动一趟轻量级拓扑排序合法化流程：在预设的 $[s_i^{ASAP}, s_i^{ALAP}]$ 范围内将冲突节点顺推至满足最小依赖延迟；对于模调度中的回边环路，则采用不动点松弛迭代推进。该机制绝不作为主要优化驱动源（实测触发频率低且微调节点比例极小），而是充当保障 100% 物理可行性的安全气囊。

实验关键数据¶

主实验¶

论文在标准硬件综合基准集（EPFL 逻辑基准）及大规模合成随机工作负载（RW）上进行了全面评测，涵盖通信/资源优化常规调度（Formulation A）、存储峰值优化常规调度（Formulation B）以及多约束模调度（Formulation C）。基线包含商用精确求解器（CPLEX、Gurobi）、传统工业级启发式（List Scheduling、FDS）以及前沿可微基准 GS-Schedule，运行时间统一限制为 15 分钟（900 秒）。

下表汇总了部分代表性大规模基准在 Formulation A 下的调度质量比值（以 GauS 成本为基准归一化为 1.000，比值越低越优；inf 表示无解或超出 4 倍；OOM 表示显存耗尽）：

| 基准图名称 | 节点数 $|V|$ | 深度 $D$ | GauS (本文) | GS-Schedule | CPLEX | Gurobi | | :--- | :--- | :--- | :--- | :--- | :--- | :--- | :--- | | router | 344 | 55 | 1.000 | 1.350 | inf (超时) | 1.020 | | adder | 1,276 | 256 | 1.000 | 1.820 | inf | inf | | bar | 3,471 | 13 | 1.000 | 1.410 | inf | inf | | RW_5 | 4,716 | 21 | 1.000 | 1.760 | inf | inf | | square | 18,550 | 251 | 1.000 | OOM | inf | inf | | multiplier | 27,190 | 275 | 1.000 | OOM | inf | inf | | div | 57,375 | 4,373 | 1.000 | OOM | inf | inf |

在全部共享可行解测试集上，GauS 相比前沿可微基线 GS-Schedule 实现了几何平均 71.8% 的解质量提升；面对节点数大于 1.8 万的巨型图，GS-Schedule 全部 OOM 崩溃，而 GauS 保持线性显存开销稳健收敛。

消融实验¶

为了严格证明高斯重参数化及各设计的内在有效性，作者在 EPFL 基准上开展了细致的对比消融：包括打乱时间步以验证“时间邻近先验”的独立贡献、更换对称单峰概率分布族验证机制泛化性、以及超参数鲁棒性扫描。

下表给出了时间序数邻近消融（固定重排时间步轴）与分布族替换在 EPFL 上的实验对比：

消融配置	核心机制设计	平均解质量比 (vs GauS)	最大单项退化 / 胜出次数 (15图)	结论说明
GauS (高斯基准)	保持自然时间轴邻近 + 标准高斯 CDF	1.000 (几何基准)	6 胜 (平均排名 2.23)	基础表现最优且稳定
打乱时间轴 (Shuffled)	破坏时间连续性，每节点独立置换步长	1.350 (+35.0% 成本恶化)	adder 退化 10.87×, priority 退化 5.98×	决定性证明时间邻近先验对引导搜索的关键价值
Logistic 分布	替换为 Logistic CDF 并匹配方差	1.075	1 胜 (平均排名 3.00)	单峰对称分布均能平滑优化，差距在 8% 以内
Laplace 分布	替换为 Laplace CDF 并匹配方差	1.002	2 胜 (平均排名 2.43)	性能与高斯极其接近，进一步印证连续松弛核心机制
Cauchy 分布	替换为柯西分布并匹配半高全宽 (FWHM)	1.066	6 胜 (平均排名 2.33)	重尾特性在部分稀疏图中具有探索优势

关键发现¶

时间邻近诱导偏置是收敛关键：打乱时间轴后所有图上的调度成本平均暴增 35%（在深层加法器 adder 上暴增近 11 倍），证明连续密度向邻近步的梯级流动直接决定了梯度方向的物理正确性。
参数扩展性支撑极高 GPU 利用率：GS-Schedule 的显存消耗随图深度和节点数爆炸，且随计算图复杂化 GPU 算力利用率跌破 40%；而 GauS 将参数量限制在 $2|V|$，在保持几乎恒定低显存开销的同时全程维持接近 100% 的 GPU 满载利用率。
泛化至非高斯单峰分布：分布族消融证明该框架的核心机理并非正态分布独占，任意具备一阶可微 CDF 的对称单峰分布均可达到类似优化水准，这为后续定制特定硬件时钟偏斜分布提供了数学空间。

亮点与洞察¶

将离散调度几何化为连续参数输运：巧妙利用一维单峰分布的均值位移承载时间维度的滑动，以最优雅的数学手段消除了类别松弛中所有时间步正交无序的病态空间设定。
首次攻克可微流水线模调度：通过对时间索引进行模映射，将离散周期性资源冲突与回边约束映射为平滑的模概率折叠，突破了以往可微方法只能处理单次前向无环图的边界。
以统计自由度指导冷启动：将经典的 ASAP/ALAP 静态时序分析结果直接映射为高斯初始标准差 $\sigma_i = \frac{1}{6}(s_i^{ALAP} - s_i^{ASAP})$，利用正态分布 $3\sigma$ 准则精确覆盖 99.7% 的合法时隙，实现了无缝高效的物理启发式预热。

局限与展望¶

节点间独立假设忽略局部边级耦合：模型假设各个算子的高斯随机变量完全独立，无法显式建模边所连接的生产者与消费者之间的联合概率分布；引入完全高斯过程（GP）虽能捕获协方差但会导致 $O(|V|^2)$ 复杂度，如何在稀疏局部邻域低成本引入相关性仍是未解问题。
极窄可行域下的边界受限：在 Formulation C（模调度）中当回边数量极大时，可行搜索空间极度狭窄，基于梯度的连续松弛在靠近硬不可行边界时容易产生震荡，导致对超大图的提升幅度收窄。
启发式合法化后的梯度非闭环：当均值舍入触发拓扑合法化调整后，修正动作是作为硬重置写回而非端到端反向传播，未来可探索可微投影层以彻底消解对后处理安全网的依赖。

评分¶

新颖性: ⭐⭐⭐⭐⭐ (首次将高斯重参数化引入算子调度并首创可微流水线模调度建模)
实验充分度: ⭐⭐⭐⭐⭐ (覆盖 EPFL 及大规模合成基准、消融完备、分析深入，含时间邻近与分布族实验)
写作质量: ⭐⭐⭐⭐⭐ (数学推导严密清晰，问题阐述与设计动机非常透彻)
价值: ⭐⭐⭐⭐⭐ (为现代深度学习编译器与EDA硬件高层综合提供了真正具备线性扩展力的实用调度新引擎)