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

arXiv: https://arxiv.org/abs/2602.02943 · HTML: https://mlsanity.com/paper/2602.02943/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Diffusion foundations and sampling, one of 14 directions in Diffusion models.
- Builds on 32 of its 52 references that are in the map.
- First posted on arXiv 3 Feb 2026.

## Placed in

- Diffusion models / [Diffusion foundations and sampling](https://mlsanity.com/diffusion/diffusion-foundations-and-sampling/)

## Builds on (most cited first, 30 of 32)

- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Denoising Diffusion Implicit Models](https://mlsanity.com/paper/2010.02502/) (2020)
- [Diffusion Models Beat GANs on Image Synthesis](https://mlsanity.com/paper/2105.05233/) (2021)
- [Generative Modeling by Estimating Gradients of the Data Distribution](https://mlsanity.com/paper/1907.05600/) (2019)
- [Convex optimization](https://mlsanity.com/paper/2106.01946/) (2021)
- [Maximum Likelihood Training of Score-Based Diffusion Models](https://mlsanity.com/paper/2101.09258/) (2021)
- [Splitwise: Efficient generative LLM inference using phase splitting](https://mlsanity.com/paper/2311.18677/) (2023)
- [Data-driven Distributionally Robust Optimization Using the Wasserstein Metric: Performance Guarantees and Tractable Reformulations](https://mlsanity.com/paper/1505.05116/) (2015)
- [Distributionally Robust Stochastic Optimization with Wasserstein Distance](https://mlsanity.com/paper/1604.02199/) (2016)
- [A Primer on Zeroth-Order Optimization in Signal Processing and Machine Learning](https://mlsanity.com/paper/2006.06224/) (2020)
- [Decision-Focused Learning: Foundations, State of the Art, Benchmark and Future Opportunities](https://mlsanity.com/paper/2307.13565/) (2023)
- [DynamoLLM: Designing LLM Inference Clusters for Performance and Energy Efficiency](https://mlsanity.com/paper/2408.00741/) (2024)
- [Distributionally Robust Bayesian Optimization](https://mlsanity.com/paper/2002.09038/) (2020)
- [The Best of Many Worlds: Dual Mirror Descent for Online Allocation Problems](https://mlsanity.com/paper/2011.10124/) (2020)
- [Dual Mirror Descent for Online Allocation Problems](https://mlsanity.com/paper/2002.10421/) (2020)
- [LLM Inference Serving: Survey of Recent Advances and Opportunities](https://mlsanity.com/paper/2407.12391/) (2024)
- [Distributionally Robust Learning](https://mlsanity.com/paper/2108.08993/) (2021)
- [Modeling the Second Player in Distributionally Robust Optimization](https://mlsanity.com/paper/2103.10282/) (2021)
- [Gen-DFL: Decision-Focused Generative Learning for Robust Decision Making](https://mlsanity.com/paper/2502.05468/) (2025)
- [Distributionally Robust Policy Learning via Adversarial Environment Generation](https://mlsanity.com/paper/2107.06353/) (2021)
- [Resource Scheduling in Edge Computing: A Survey](https://mlsanity.com/paper/2108.08059/) (2021)
- [Hybrid Heterogeneous Clusters Can Lower the Energy Consumption of LLM Inference Workloads](https://mlsanity.com/paper/2407.00010/) (2024)
- [Distributionally Robust Bayesian Optimization with varphi-divergences](https://mlsanity.com/paper/2203.02128/) (2022)
- [A Queueing Theoretic Perspective on Low-Latency LLM Inference with Variable Token Length](https://mlsanity.com/paper/2407.05347/) (2024)
- [Joint Online Learning and Decision-making via Dual Mirror Descent](https://mlsanity.com/paper/2104.09750/) (2021)
- [Differentiable Distributionally Robust Optimization Layers](https://mlsanity.com/paper/2406.16571/) (2024)
- [AI Augmented Edge and Fog Computing: Trends and Challenges](https://mlsanity.com/paper/2208.00761/) (2022)
- [Precise Energy Consumption Measurements of Heterogeneous Artificial Intelligence Workloads](https://mlsanity.com/paper/2212.01698/) (2022)

## Abstract (arXiv)

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.