Decision-Focused Learning

End-to-end learning that optimizes for decision quality, not prediction accuracy.

We develop theory and algorithms for decision-focused learning — training ML models whose outputs feed directly into optimization problems, trained end-to-end with respect to decision quality rather than prediction error.

Key directions:

  • Fully first-order methods for bilevel and linearly-constrained optimization
  • Differentiable optimization layers for neural networks
  • Predict-then-optimize for sequential decision problems (MDPs, RL)
  • Compact quality-aware surrogate learning

Selected publications:

(Zhao et al., 2026; Wang et al., 2021; Wang et al., 2020; Kornowski et al., 2024; Shah et al., 2022)

References

2026

  1. A Fully First-Order Layer for Differentiable Optimization
    Zihao Zhao, Kai-Chia Mo, Shing-Hei Ho, Brandon Amos, and Kai Wang
    In International Conference on Machine Learning (ICML), 2026
    Spotlight
    arXiv

2024

  1. Fully First-Order Methods for Linearly Constrained Bilevel Optimization
    Guy Kornowski, Swati Padmanabhan, Kai Wang, Zhe Zhang, and Suvrit Sra
    In Advances in Neural Information Processing Systems (NeurIPS), 2024
    arXiv

2022

  1. Decision-Focused Learning without Decision-Making: Learning Locally Optimized Decision Losses
    Sanket Shah, Kai Wang, Bryan Wilder, Andrew Perrault, and Milind Tambe
    In Advances in Neural Information Processing Systems (NeurIPS), 2022
    arXiv

2021

  1. Learning MDPs from Features: Predict-Then-Optimize for Sequential Decision Problems by Reinforcement Learning
    Kai Wang, Sanket Shah, Haipeng Chen, Andrew Perrault, Finale Doshi-Velez, and 1 more author
    In Advances in Neural Information Processing Systems (NeurIPS), 2021
    Spotlight
    arXiv

2020

  1. Automatically Learning Compact Quality-aware Surrogates for Optimization Problems
    Kai Wang, Bryan Wilder, Andrew Perrault, and Milind Tambe
    In Advances in Neural Information Processing Systems (NeurIPS), 2020
    Spotlight
    arXiv