Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

3D-Learning: Diffusion-Augmented Distributionally Robust Decision-Focused Learning

Jiaqi Wen, Lei Fan, Jianyi Yang

TL;DR

This work addresses the sensitivity of Predict-then-Optimize decisions to distribution shifts by introducing Distributionally Robust Decision-Focused Learning (DR-DFL) and its diffusion-augmented instantiation, 3D-Learning. By modeling the ambiguity set with diffusion processes and constraining them via score-matching loss $J(\theta)$, the approach enables tractable search for worst-case distributions while remaining aligned with training data through a reversed KL guidance. The algorithm alternates inner maximization over diffusion parameters with outer minimization over the predictor, using a dual-learning reformulation and PPO-based gradient estimates. Empirical results on LLM resource provisioning show that 3D-Learning outperforms traditional Wasserstein/KL-based DRO and data augmentation in both average and worst-case regimes, and exhibits robust performance under corrupted inputs, indicating practical value for robust PTO in computing and networking contexts.

Abstract

Predict-then-Optimize (PTO) pipelines are widely employed in computing and networked systems, where Machine Learning (ML) models are used to predict critical contextual information for downstream decision-making tasks such as cloud LLM serving, data center demand response, and edge workload scheduling. However, these ML predictors are often vulnerable to out-of-distribution (OOD) samples at test time, leading to significant decision performance degradation due to large prediction errors. To address the generalization challenges under OOD conditions, we present the framework of Distributionally Robust Decision-Focused Learning (DR-DFL), which trains ML models to optimize decision performance under the worst-case distribution. Instead of relying on classical Distributionally Robust Optimization (DRO) techniques, we propose Diffusion-Augmented Distributionally Robust Decision-Focused Learning (3D-Learning), which searches for the worst-case distribution within the parameterized space of a diffusion model. By leveraging the powerful distribution modeling capabilities of diffusion models, 3D-Learning identifies worst-case distributions that remain consistent with real data, achieving a favorable balance between average and worst-case scenarios. Empirical results on an LLM resource provisioning task demonstrate that 3D-Learning outperforms existing DRO and Data Augmentation methods in OOD generalization performance.

3D-Learning: Diffusion-Augmented Distributionally Robust Decision-Focused Learning

TL;DR

This work addresses the sensitivity of Predict-then-Optimize decisions to distribution shifts by introducing Distributionally Robust Decision-Focused Learning (DR-DFL) and its diffusion-augmented instantiation, 3D-Learning. By modeling the ambiguity set with diffusion processes and constraining them via score-matching loss , the approach enables tractable search for worst-case distributions while remaining aligned with training data through a reversed KL guidance. The algorithm alternates inner maximization over diffusion parameters with outer minimization over the predictor, using a dual-learning reformulation and PPO-based gradient estimates. Empirical results on LLM resource provisioning show that 3D-Learning outperforms traditional Wasserstein/KL-based DRO and data augmentation in both average and worst-case regimes, and exhibits robust performance under corrupted inputs, indicating practical value for robust PTO in computing and networking contexts.

Abstract

Predict-then-Optimize (PTO) pipelines are widely employed in computing and networked systems, where Machine Learning (ML) models are used to predict critical contextual information for downstream decision-making tasks such as cloud LLM serving, data center demand response, and edge workload scheduling. However, these ML predictors are often vulnerable to out-of-distribution (OOD) samples at test time, leading to significant decision performance degradation due to large prediction errors. To address the generalization challenges under OOD conditions, we present the framework of Distributionally Robust Decision-Focused Learning (DR-DFL), which trains ML models to optimize decision performance under the worst-case distribution. Instead of relying on classical Distributionally Robust Optimization (DRO) techniques, we propose Diffusion-Augmented Distributionally Robust Decision-Focused Learning (3D-Learning), which searches for the worst-case distribution within the parameterized space of a diffusion model. By leveraging the powerful distribution modeling capabilities of diffusion models, 3D-Learning identifies worst-case distributions that remain consistent with real data, achieving a favorable balance between average and worst-case scenarios. Empirical results on an LLM resource provisioning task demonstrate that 3D-Learning outperforms existing DRO and Data Augmentation methods in OOD generalization performance.
Paper Structure (31 sections, 1 theorem, 18 equations, 4 figures, 2 tables, 2 algorithms)

This paper contains 31 sections, 1 theorem, 18 equations, 4 figures, 2 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Formulation and Appications
  3. Decision-Focused Learning
  4. Distributionally Robust Decision-Focused Learning
  5. Applications in Computing and Communication Networks
  6. Workload-Aware Resource Provisioning for LLM Inference
  7. Demand Response in AI Data Centers
  8. Channel-Aware Edge Data Center Selection
  9. Diffusion Ambiguity Modeling
  10. Challenges in Ambiguity Modeling
  11. Distribution Modeling via Diffusion Models
  12. Diffusion-Based Ambiguity Set
  13. 3D-Learning Algorithm
  14. Inner Maximization of 3D-Learning
  15. Min-Max Solution of 3D-Learning
  16. ...and 16 more sections

Key Result

Lemma 1

Given the assumptions listed in Appendix A of score-based_diffusion_SDE_song2021maximumThe assumptions require that the expected squared norm over $P_0$ and $\pi$ are bounded by any finite value, the functions $k(\cdot, t)$, $\nabla_x\log P_t(x)$, and $s_{\theta}(\cdot, t)$ are Lipschitz continuous where $P_T$ is the final-step output distribution of the forward process and $P_T\approx \pi$ by th

Figures (4)

  • Figure 1: Framework of Decision Focused Learning
  • Figure 2: Robustness evaluation under diverse noisy corruptions on $\textbf{23D24D}$ dataset.
  • Figure 3: 3D-Learning with decision-focused training and MSE training.
  • Figure 4: Effect of budget $\epsilon$ on 3D-Learning performance.

Theorems & Definitions (1)

  • Lemma 1: Theorem 1 and Corollary 1 in score-based_diffusion_SDE_song2021maximum