Publications

Publications from the KOALA Lab at Georgia Tech.

Kai Wang's name is underlined. See also Google Scholar.

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
  2. Diffusion-DFL: Decision-focused Diffusion Models for Stochastic Optimization
    Zihao Zhao, Christopher Yeh, Lingkai Kong, and Kai Wang
    In International Conference on Learning Representations (ICLR), 2026
    arXiv

2025

  1. Finding a Multiple Follower Stackelberg Equilibrium: A Fully First-Order Method
    April Niu, Kai Wang, and Juba Ziani
    2025
    In submission
    arXiv
  2. Bridging Constraints and Stochasticity: A Fully First-Order Method for Stochastic Bilevel Optimization with Linear Constraints
    Cac Phan and Kai Wang
    2025
    In submission
    arXiv
  3. Revealing Potential Biases in LLM-Based Recommender Systems in the Cold Start Setting
    Alexandre Andre, Gauthier Roy, Eva Dyer, and Kai Wang
    2025
    EARL 2025 Workshop
    arXiv
  4. One-Step Flow Policy Mirror Descent
    Tianyi Chen, Haitong Ma, Na Li, Kai Wang, and Bo Dai
    2025
    In submission
    arXiv
  5. Non-Stationary Restless Multi-Armed Bandits with Provable Guarantee
    Yu Heng Hung, Kai Wang, and Ping-Chun Hsieh
    2025
    In submission
    arXiv
  6. Neural Index Policies for Restless Multi-Action Bandits with Heterogeneous Budgets
    Himadri S Pandey, Kai Wang, and Gian-Gabriel P. Garcia
    2025
    In submission
    arXiv
  7. Networked Restless Multi-Arm Bandits with Reinforcement Learning
    Hanmo Zhang, Zenghui Sun, and Kai Wang
    2025
    PRL Workshop, AAAI 2025
    arXiv
  8. Efficient Online Reinforcement Learning for Diffusion Policy
    Haitong Ma, Tianyi Chen, Kai Wang, Na Li, and Bo Dai
    In International Conference on Machine Learning (ICML), 2025
    arXiv
  9. Opportunistic Screening of Type 2 Diabetes with Deep Metric Learning using Electronic Health Records
    Qixuan Jin, Haoran Zhang, Lukasz Szczerbinski, Jiacheng Zhu, Walter Gerych, and 10 more authors
    Scientific Reports, 2025
    HTML
  10. Primal-Dual Spectral Representation for Off-policy Evaluation
    Yang Hu, Tianyi Chen, Na Li, Kai Wang, and Bo Dai
    In International Conference on Artificial Intelligence and Statistics (AISTATS), 2025
    arXiv
  11. What is the Right Notion of Distance between Predict-then-Optimize Tasks?
    Paula Rodriguez-Diaz, Lingkai Kong, Kai Wang, David Alvarez-Melis, and Milind Tambe
    In Conference on Uncertainty in Artificial Intelligence (UAI), 2025
    arXiv
  12. What’s in a Query: Polarity-aware Distribution-based Fair Ranking
    Aparna Balagopalan, Kai Wang, Olawale Elijah Salaudeen, Asia Biega, and Marzyeh Ghassemi
    In The Web Conference (WWW), 2025
    HTML

2024

  1. Aligning Large Language Models with Representation Editing: A Control Perspective
    Lingkai Kong, Haorui Wang, Wenhao Mu, Yuanqi Du, Yuchen Zhuang, and 5 more authors
    In Advances in Neural Information Processing Systems (NeurIPS), 2024
    arXiv
  2. 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

2023

  1. Characterizing and Improving the Robustness of Predict-Then-Optimize Frameworks
    Sonja Johnson-Yu, Jessica Finocchiaro, Arunesh Sinha, Kai Wang, Yevgeniy Vorobeychik, and 2 more authors
    In Conference on Decision and Game Theory for Security (GameSec), 2023
    HTML
  2. Restless Multi-Armed Bandits for Maternal and Child Health: Results from Decision-Focused Learning
    Shresth Verma, Aditya Mate, Kai Wang, Neha Madhiwalla, Aparna Hegde, and 2 more authors
    In International Conference on Autonomous Agents and Multi-Agent Systems (AAMAS), 2023
    HTML
  3. Optimistic Whittle Index Policy: Online Learning for Restless Bandits
    Kai Wang, Lily Xu, Aparna Taneja, and Milind Tambe
    In AAAI Conference on Artificial Intelligence (AAAI), 2023
    arXiv
  4. Scalable Decision-Focused Learning in Restless Multi-Armed Bandits with Application to Maternal and Child Health
    Kai Wang, Shresth Verma, Aditya Mate, Sanket Shah, Aparna Taneja, and 3 more authors
    In AAAI Conference on Artificial Intelligence (AAAI), 2023
    arXiv
  5. Smoothed Online Combinatorial Optimization Using Imperfect Predictions
    Kai Wang, Zhao Song, Georgios Theocharous, and Sridhar Mahadevan
    In AAAI Conference on Artificial Intelligence (AAAI), 2023
    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
  2. Coordinating Followers to Reach Better Equilibria: End-to-End Gradient Descent for Stackelberg Games
    Kai Wang, Lily Xu, Andrew Perrault, Michael K. Reiter, and Milind Tambe
    In AAAI Conference on Artificial Intelligence (AAAI), 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
  2. Dual-Mandate Patrols: Multi-Armed Bandits for Green Security
    Lily Xu, Elizabeth Bondi, Fei Fang, Andrew Perrault, Kai Wang, and 1 more author
    In AAAI Conference on Artificial Intelligence (AAAI), 2021
    Best Paper Runner-Up
    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
  2. Robust Spatial-Temporal Incident Prediction
    Ayan Mukhopadhyay, Kai Wang, Andrew Perrault, Mykel Kochenderfer, Milind Tambe, and 1 more author
    In Conference on Uncertainty in Artificial Intelligence (UAI), 2020
    HTML
  3. Scalable Game-Focused Learning of Adversary Models: Data-to-Decisions in Network Security Games
    Kai Wang, Andrew Perrault, Aditya Mate, and Milind Tambe
    In International Conference on Autonomous Agents and Multi-Agent Systems (AAMAS), 2020
    HTML

2019

  1. DeepFP for Finding Approximate Nash Equilibrium in Continuous Action Spaces
    Nitin Kamra, Umang Gupta, Kai Wang, Fei Fang, Yan Liu, and 1 more author
    In Conference on Decision and Game Theory for Security (GameSec), 2019
    HTML
  2. Learning to Signal in the Goldilocks Zone: Improving Adversary Compliance in Security Games
    Sarah Cooney, Kai Wang, Elizabeth Bondi, Thanh Nguyen, Phebe Vayanos, and 5 more authors
    In European Conference on Machine Learning and Principles and Practice of Knowledge Discovery in Databases (ECML-PKDD), 2019
    HTML
  3. Deep Fictitious Play for Games with Continuous Action Spaces
    Nitin Kamra, Umang Gupta, Kai Wang, Fei Fang, Yan Liu, and 1 more author
    In International Conference on Autonomous Agents and Multi-Agent Systems (AAMAS), 2019
    Extended Abstract
    HTML
  4. Adversarial Machine Learning with Double Oracle
    Kai Wang, Bryan Wilder, and Milind Tambe
    In International Joint Conference on Artificial Intelligence (IJCAI) Doctoral Consortium, 2019
    HTML
  5. Improving GP-UCB Algorithm by Harnessing Decomposed Feedback
    Kai Wang, Bryan Wilder, Sze-chuan Suen, Milind Tambe, and Bistra Dilkina
    In ECML SoGood Workshop, 2019
    HTML

2018

  1. The Price of Usability: Designing Operationalizable Strategies for Security Games
    Sara Marie Mc Carthy, Corine Laan, Kai Wang, Phebe Vayanos, Milind Tambe, and 1 more author
    In International Joint Conference on Artificial Intelligence (IJCAI), 2018
    HTML
  2. Equilibrium Refinement in Security Games with Arbitrary Scheduling Constraints
    Kai Wang, Qingyu Guo, Phebe Vayanos, Milind Tambe, and Bo An
    In International Conference on Autonomous Agents and Multi-Agent Systems (AAMAS), 2018
    HTML
  3. Strategic Coordination of Human Patrollers and Mobile Sensors with Signaling for Security Games
    Haifeng Xu, Kai Wang, Phebe Vayanos, and Milind Tambe
    In AAAI Conference on Artificial Intelligence (AAAI), 2018
    HTML