Preventing Model Extraction in Contextual Linear Optimization

Adam N. Elmachtoub; Tianyi Lin; Wenxuan Liu

cs.LGOct 21, 2026

Preventing Model Extraction in Contextual Linear Optimization

Authors: Adam N. Elmachtoub, Tianyi Lin, Wenxuan Liu

Organizations: Columbia University

Abstract

Consider a contextual linear optimization problem, where a predict-then-optimize (PTO) system returns optimal decisions without revealing the cost vectors used in its downstream linear program. Our focus is on how to prevent attackers, who can submit context queries, from extracting the underlying prediction model. Although the cost-to-decision map is many-to-one, observations across contexts can reveal the underlying prediction model. For linear models, we provide sufficient conditions where the attacker can learn the model up to positive scaling by only observing the decision. These conditions motivate an adaptive attack that alternates inverse optimization with boundary sampling, i.e., using context queries leading to decision boundaries of the current attacker model. To prevent extraction, we propose a mechanism that perturbs decisions only when the predicted cost vector is near a decision boundary. We establish a regret bound in terms of the perturbation radius. For the attackers that fit the defended responses as optimal decisions under the evaluation distribution, we also prove a lower bound on expected extraction error. Experiments on synthetic instances and real data quantify the security-performance trade-off of the defense mechanism and illustrate that our defense can increase the extraction error.

Figures & tables

Panel of Figure 1 — Figure 1: Illustration of prediction misalignment. Arrows show negative cost vectors. The two models return different decisions in ( 1(a) ) and the same decision in ( 1(b) ).

Figure 2: Attack performance on synthetic data as the decision dimension

n

(left) and feature dimension

d

(right) vary. The first and third panels report prediction loss. The second and fourth report the number of context-decision observations required to attain loss below

0.001

. Prediction loss points and error bars show medians and

25^{\textnormal{th}}

75^{\textnormal{th}}

percentiles over 100 repetitions per parameter setting.

Panel of Figure 3 — Figure 3: Prediction loss under increasing query budgets for BndS ( $M=20$ ), Unif, and MIX, with $\epsilon=0$ (left) and $\epsilon=0.0015$ (right). Curves show medians over five repetitions with $d=10$ and $n=5$ .

Panel of Figure 4 — Figure 4: Empirical prediction loss-regret trade-offs: (a) decision perturbation against different attacks on synthetic instances; (b) prediction versus decision perturbation on synthetic instances; and (c) decision perturbation on the short-video recommendation task. Points show median regret and median prediction loss, and shaded bands show the $25^{\textnormal{th}}$ - $75^{\textnormal{th}}$ percentiles of prediction loss over 100 repetitions per parameter setting.

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 5 — Figure 5: Convergence of the ADMM algorithm and Tang and Khalil’s method. The red line represents the optimization loss, defined as the proportion of incorrect decisions $w^{\star}(P^{a}x_{i})\neq w^{\star}(P^{DM}x_{i})$ over the dataset.

Algorithm	CaVE loss	predLoss	optLoss
ADMM	0.0007018	0.0356	0.1450
Tang and Khalil’s method	0.0000032	0.0539	0.1150
Global	0.0000030	0.0378	0.0650

Appendix

Table 1: Comparison of three approaches for minimizing the CaVE loss.

Panel of Figure 6 — Figure 6: (a) the tight constraints (also the boundaries of normal cone) incident to the a normal cone $C_{k}$ . (b) the normal cones (also the vertex) incident to the constraint $a_{i}$ .

Panel of Figure 7 — Figure 7: The out-of-sample prediction loss and iterations for the attacking framework in different attacker’s parameter settings.

Panel of Figure 8 — Figure 8: (a) Prediction loss; (b) Sample size for BndS and Unif against MLP and SVR predictors fitted to the linear synthetic data; and (c) Prediction loss of a linear attacker against MLP and SVR predictors as the quadratic coefficient $\gamma$ varies.

Panel of Figure 9 — Figure 9: Prediction loss (a) and runtime per iteration (b) of Unif on fractional Knapsack instances of different dimensions.

Figure 10: Security-performance trade-off under the fractional Knapsack problem with

n=20

, using at most three neighboring vertices per query.

Panel of Figure 11 — Figure 11: Prediction loss (left) and decision-maker’s regret (right) as functions of the defensive parameter $\epsilon$ .

Model Extraction Attacks

Model extraction attacks infer a machine-learning model’s behavior, architecture, parameters, or hidden assets from queries and system leakage.

See all 49 papers

Decision-Focused On-Policy Learning for Contextual Linear Optimization with Partial Feedback
May 31, 2026Wyame Benslimane, Tinghan Ye, Pascal Van Hentenryck +1Predict-Then-Optimize Decision-Focused Learning
Department of Industrial Engineering and Operations Research University of California, Berkeley · H. Milton Stewart School of Industrial and Systems Engineering Georgia Institute of Technology
Adversarial Training of Linear Models under Stealthy Attacks
Aug 26, 2026Lovisa Eriksson, Dave Zachariah, André M. H. TeixeiraAdversarial Attacks Adversarial Robustness
Department of Information Technology, Uppsala University, Uppsala, Sweden
Best-of-Both Worlds for linear contextual bandits with paid observations
Oct 8, 2025Nathan Boyer, Dorian Baudry, Patrick RebeschiniMulti-Armed Bandits Follow-the-Regularized-Leader
ENS Paris-Saclay, France · Univ. Grenoble Alpes, Inria, CNRS, Grenoble INP, LIG, 38000 Grenoble, France · Department of Statistics, University of Oxford

May 31, 2026J/K move · Enter open · S save

Decision-Focused On-Policy Learning for Contextual Linear Optimization with Partial Feedback

Wyame Benslimane, Tinghan Ye, Pascal Van Hentenryck, Paul Grigas

Department of Industrial Engineering and Operations Research University of California, Berkeley · H. Milton Stewart School of Industrial and Systems Engineering Georgia Institute of Technology

Predict-Then-Optimize Decision-Focused Learning Policy Gradient Methods Contextual Bandits Policy Learning

Open paper AskPDF arXiv

Decision-focused learning (DFL) trains predictive models by optimizing downstream decision quality rather than standalone prediction accuracy. For contextual linear optimization, most existing DFL methods assume offline data and full observations of the objective cost vector. We develop an on-policy learning method for sequential contextual linear optimization under partial feedback, generalizing the standard bandit feedback setting. Our method learns a stochastic predict-then-optimize policy that samples a cost-vector prediction from a conditional distribution and solves the resulting downstream linear optimization problem. To update this distributional model, we introduce a two-component hybrid gradient estimator. The first component is a score function estimator, which provides an unbiased but potentially high-variance policy gradient estimate. The second is a decision-focused plug-in component that uses an auxiliary nuisance estimate of the latent cost vector to exploit the downstream optimization structure, becoming more informative as the estimate improves. We prove an O(T^-1/2) bound on the average squared policy-gradient norm, matching the standard non-convex SGD rate. Experiments on top-k selection, shortest path, combinatorial pricing, and a real-data energy-scheduling benchmark show that, under bandit feedback, the hybrid gradient approach achieves lower mean cumulative regret than the contextual-bandit baselines on all four benchmarks and than linear Thompson sampling on three, and that it also works with richer conditional generative cost models. Code is available at https://github.com/Joeyetinghan/on-policy-bandit-dfl.

Abstract

Model Extraction Attacks

Explore similar work

Decision-Focused On-Policy Learning for Contextual Linear Optimization with Partial Feedback

Adversarial Training of Linear Models under Stealthy Attacks

Best-of-Both Worlds for linear contextual bandits with paid observations