GauS: Differentiable Scheduling Optimization via Gaussian Reparameterization

2026-02-23

AI Summary

This work addresses the limitations of traditional operator scheduling methods in efficiency and scalability, particularly the inability of existing differentiable approaches to effectively model temporal ordering and their prohibitively large parameter spaces. The paper proposes GauS, a novel framework that introduces Gaussian reparameterization into scheduling optimization for the first time, formulating the scheduling problem as a stochastic relaxation based on Gaussian distributions. By leveraging continuous Gaussian variables, GauS naturally captures the ordinal nature of time while drastically reducing the parameter space. This enables differentiable modeling of complex pipeline schedules and seamless integration of diverse objectives and constraints. Combined with GPU-accelerated parallel computation, GauS efficiently generates Pareto-optimal schedules across multiple benchmarks, demonstrating its generality and superior performance.

Technology Category

• Planning, Routing, and Scheduling: Scheduling
• Search and Optimization: Mixed Discrete/Continuous Search
• Constraint Satisfaction and Optimization: Distributed CSP/Optimization

Application Category

• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

Efficient operator scheduling is a fundamental challenge in software compilation and hardware synthesis. While recent differentiable approaches have sought to replace traditional ones like exact solvers or heuristics with gradient-based search, they typically rely on categorical distributions that fail to capture the ordinal nature of time and suffer from a parameter space that scales poorly. In this paper, we propose a novel differentiable framework, GauS, that models operator scheduling as a stochastic relaxation using Gaussian distributions, which fully utilize modern parallel computing devices like GPUs. By representing schedules as continuous Gaussian variables, we successfully capture the ordinal nature of time and reduce the optimization space by orders of magnitude. Our method is highly flexible to represent various objectives and constraints, which provides the first differentiable formulation for the complex pipelined scheduling problem. We evaluate our method on a range of benchmarks, demonstrating that Gaus achieves Pareto-optimal results.

Problem

• operator scheduling
• differentiable optimization
• Gaussian reparameterization
• ordinal time
• pipelined scheduling

Innovation

• differentiable scheduling
• Gaussian reparameterization
• operator scheduling
• stochastic relaxation
• pipelined scheduling