E2E-AT: A Unified Framework for Tackling Uncertainty
in Task-aware End-to-end Learning

In supervised learning, a parametric model $\bm{f}(\bm{x};\bm{\theta})$ can be defined to map from input $\bm{x}$ to label $\bm{y}$ for $(\bm{x},\bm{y})\in\mathcal{D}$ by minimizing the following empirical loss:

In this paper, we use $(\bm{x},\bm{y})\in\mathcal{D}$ and $i\in\mathcal{D}$ interchangeably to denote a sample of the dataset.

It has been well studied that the deep neural network is prone to small perturbations on its input. In security-critical applications, robustness has become an emerging factor. Adversarial training has been shown to improve the robustness of NN, which transforms (1) into a robust optimization (Madry et al. 2017):

where $\Delta=\{\bm{\delta}:\|\bm{\delta}\|_{\infty}\leq\epsilon\}$ is the attack budget set for some small $\epsilon>0$.

Robust optimization (2) is practically solved by iterative approaches in which the adversarial perturbation is first solved by inner minimization through gradient ascent. To keep the attack within $\Delta$, projected gradient descent (PGD) (Madry et al. 2017) is adopted:

where $\mathcal{P}_{\Delta}$ is the projector on $\Delta$ and $\gamma$ is the step size.

Although adversarial training is effective in improving robustness, it may give the wrong sign of security as inner maximization is inexactly solved. A certification approach can be made on the feedforward neural network with ReLU activations, e.g., a piecewise linear neural network³³3Convolution layer and other piecewise linear activations, such as leaky ReLU, are also piecewise linear.. Using the big-M method, the neural network with $d$ layers can be represented by the following set (Tjeng, Xiao, and Tedrake 2017)

where $\bm{\theta}_{i}=(\bm{W}_{i},\bm{b}_{i})$ is the weights of the $i$-th layer. $\bm{u}_{i}$ and $\bm{l}_{i}$ are the upper and lower bounds of the output of $i$-th layer, which can be efficiently estimated by interval bound propagation (IBP) (Gowal et al. 2018). $\bm{v}_{i}$ is an integer vector that controls the activation of ReLUs.

Based on (4), the inner maximization of (2) becomes:

As $\mathcal{C}_{\text{nn}}(\bm{x};\bm{\theta})$ is defined by mixed integer linear constraints, (5) can be solved exactly for certain types of loss function. When cross-entropy loss is used, we can choose to maximize the output of each logit, and if all the logits are not greater than the ground-truth logit within the attack budget, the NN is said to be certified robust at the candidate sample. For a regression task, (5) can be formulated as mixed-integer linear programming (MILP) to maximize or minimize the forecast value (Xu and Teng 2023).

End-to-end learning takes the fusion of prediction (ML) and decision making (CO), which aims to find the map from the data to the optimal decision such that the learned ML model can optimally reflect the CO via training. Meanwhile, the E2E learning distinguishes itself from separate supervised learning followed by the CO in which the prediction divergence and the task-aware cost can have a large mismatch (Sadana et al. 2023).

Referring to a recent review (Kotary et al. 2021), the CO can be modelled as

where $\ell(\cdot)$ is the objective function, $\bm{z}$ is the decision variable, $\bm{y}$ is the parameter, and $\mathcal{C}$ is the constraint set.

Let $\hat{\bm{y}}$ be the output of $\bm{f}(\bm{x};\bm{\theta})$ and let $\hat{\bm{z}}^{\star}$ be the minimizer of (6) parameterized by $\hat{\bm{y}}$. There are two options to learn the E2E model guided by (6). In the supervised learning setting, a suitable loss function $\mathcal{M}$ can be chosen to minimize the difference to ground-truth decision, e.g. $\mathcal{M}(\hat{\bm{z}}^{\star},\bm{z}^{\star})$ or $\mathcal{M}(\ell(\hat{\bm{z}}^{\star};\hat{\bm{y}}),\ell(\bm{z}^{\star};\bm{y}))$ (Kong et al. 2022) . Alternatively, the regret function (Elmachtoub and Grigas 2022) can be implemented:

The ground truth decision $\bm{z}^{\star}$ in (7) is not necessary to compute, as $\ell(\bm{z}^{\star};\bm{y})$ is a constant (Wilder, Dilkina, and Tambe 2019). Therefore, we argue that the regret function (7) is more intrinsic as the training loss has the same format as the CO objective (6).

Consider the uncertainty of the input $\bm{x}+\bm{\delta}_{x}\in\mathcal{X}$ and the unpredictable parameter $\bm{\phi}+\bm{\delta}_{\phi}\in\Phi$. Denote $\bm{\psi}=(\bm{x},\bm{\phi})\in\Psi:=\mathcal{X}\times\Phi$. The worst scenario, which maximizes the task-aware objective, can be formulated as

in which $\bm{\theta}$ is fixed.

Consequently, a robust optimization can be formulated. A unified E2E adversarial training (E2E-AT) considering both input and CO uncertainties becomes

where the constrains are subject to both minimization and maximization. Adversarial training can be adopted to solve (16). Similarly to the implicit function theorem (14), the gradient of the constraint exists, regardless of the minimization or maximization. Therefore, Danskin’s theorem can be used by first solving the inner maximization through gradient ascent (with $\bm{\theta}$ fixed) and then for the outer gradient descent (with $\bm{\psi}$ fixed). Although using the Danskin theorem requires one to exactly solve the inner maximization, it can give a descent direction for suboptimal $\bm{\psi}$, e.g. solved by PGD (3) and is applicable to various adversarial training algorithms (Madry et al. 2017; Zhang et al. 2019; Dong et al. 2020).

Although PGD (3) is effective for E2E-AT, it cannot verify the robustness, as it only finds the local maximum (Xiao et al. 2018). In particular, a robustness certification on (11) verifies if an adversarial example exists within the budget $\Delta$ such that the task-aware objective is altered by a certain amount. The key is to find the exact adversarial attack in (15). We show that for specific type of objective and COs (e.g. affine-parametric QPs), optimal solution to (15) can be solved exactly, which extends the certified robustness in piecewise linear neural network (4).

The proof can be found in the appendix.

Due to the complexity of integer programming, we restrict the original settings in (Amos and Kolter 2017) by assuming linearity in uncertain terms for certified robustness. For example, when the uncertainty in the input feature is considered, the parameter $\bm{z}_{i}$ that represents the optimal decision of the previous task is decoupled from the variable $\bm{z}$ and is affine so that the reformulation is linear. When the uncertainty of CO is considered, we assume that the uncertain unpredictable parameter is decoupled from the variable as well. We note that this setting follows the disciplined parametrized programming (DPP) (Agrawal et al. 2019).

Since both affine-parametric QP (17) and neural network (4) can be represented by mixed-integer linear constraints, maximizing an affine function subject to these constraints becomes mixed-integer linear programming (MILP), which can be effectively solved and used to certify the worst possible cost.

Previously, the uncertainty involved in COs has been considered in many E2E learning algorithms. From the perspective of contextual optimization, E2E learning applies ML to predict the uncertain parameter (Sadana et al. 2023). In such setting, probabilistic forecast and stochastic program can be implemented (Donti, Amos, and Kolter 2017). However, we view the entire E2E learning as an integrated model such that the uncertainty of the intermediate variable can be merged within the E2E training objective. Alternatively, we separate the parameters of COs into two parts. The first part (predictable parameter) is forecasted by the ML while the second part is unpredictable. Indirectly, the uncertainty of the predictable parameter is tackled by adversarial training on the input feature, while the uncertainty of unpredictable parameter needs to be tackled as well.

We note that the physical meaning of the uncertainty of the unpredictable parameter is different from that of the adversarial attack in the input feature space. Although the ultimate goal is to improve the robustness of the ML, the adversarial attack assumes that there exists a malicious party that can find the worst attack. However, the uncertainty of the unpredictable parameter always exists regardless of the malicious party. Inspired by the previous work (Donti et al. 2021; Agarwal et al. 2022), we can alternatively view the E2E-AT on the unpredictable parameter as the following optimization problem:

where $\bar{\bm{z}}^{i}$ is the decision variable from the COs parameterized by the nominal unpredictable parameter $\bar{\bm{\phi}}$.

It can be argued that the NN forecaster is trained by considering the expected task-aware cost due to uncertainties. The learning objective (19) takes the nominal unpredictable parameter (denoted as $\bar{\bm{\phi}}$) and the expected uncertainties over $\bm{\Phi}^{i}$ into account, which are balanced by the hyperparameter $\alpha$. As shown in (19), this stochastic program can be solved by sampling $\bm{\phi}^{i}\in\bm{\Phi}^{i}$ during training. In addition, if the uncertainty set of the unpredictable parameter is independent of the sample, $\bm{\Phi}$ is not subject to index $i$.

In E2E-AT, (19) is solved by robust optimization by finding the maximum over $\bm{\psi}^{i}\in\Psi^{i}$, as in (16). We also take the input uncertainty into account:

The new adversarial training objective provides an upper bound on the expectation part of (19). Meanwhile, similar to adversarial training on image tasks (Zhang et al. 2019), $\alpha$ can be used to balance the clean and adversarial accuracies. When $\alpha\rightarrow 1$, (20) becomes the original E2E learning and when $\alpha\rightarrow 0$, it becomes pure adversarial training. In addition, clean and adversarial training losses may not directly reflect clean and robust accuracy for image tasks, causing an unbalanced training objective. In E2E-AT, the two objectives are defined by the task, which is the exact metric during decision making.

Previously in (Donti et al. 2021; Agarwal et al. 2022), the authors reformulate a bilevel optimization problem into robust optimization and use the implicit function theorem for constraint satisfaction. Connecting (19) to (20), we extend (Donti et al. 2021) to training ML models. It can be seen that for each mini-batch in E2E-AT, a similar robust optimization is solved as (Donti et al. 2021). The implicit function theorem is also used to learn the task-aware objective while satisfying the constraints.

In E2E-AT, the number of forward and backward passes is equal to $\texttt{no\_batch}\times(\texttt{no\_pgd}+1)\times\texttt{no\_epoch}$ (assuming that all minibatches have the same size). Adversarial training is computationally ineffective due to intensive backpropagation (controlled by the complexity of the neural network). The computational burden is even higher in E2E-AT as in each forward pass, the COs need to be solved (controlled by the complexity of the COs). To save training time, we adopt the gradient reuse strategy. In the adversarial training for free (Shafahi et al. 2019), the attack vector and the model parameter are repeatedly updated in the same mini-batch for no_pgd times. Then, epoch_no is divided by no_pgd to maintain the total number of model updates unchanged. This results in $\texttt{no\_batch}\times\texttt{no\_epoch}$ numbers of forward and backward passes, which is the same as the clean E2E training.

A practical power system operation problem, named as network constrained economic dispatch (NCED), is considered, which has been widely used in the US and can be formulated as two-stage QP or LP (Conejo and Baringo 2018). In stage one (also known as dispatch), the set points of the generator are determined based on the forecast load. The goal of the first stage is to minimize the generator cost while meeting the physical constraints of the grid. When the generator has been dispatched, we consider a realization on the actual load by solving the second stage problem (also known as redispatch), in which any mismatches on the load and generation from stage one, as well as the violation of the physical constraints of the grid, will be penalized by extra cost:

where $\bm{P}_{g}^{\star}$ is the optimal generator dispatch, $\bm{P}_{ls}^{\star}$ is the load shedding and $\bm{P}_{gs}^{\star}$ is the power storage. $\bm{b}$ is the susceptance of the transmission line (when the resistance is close to zero, the susceptance is reciprocal to the reactance). The task-aware objective is defined as

where $\bm{c}_{g}$, $\bm{c}_{ls}$, and $\bm{c}_{gs}$ are the coefficients such that $\bm{c}_{ls}\gg\bm{c}_{gs}\gg\bm{c}_{g}$. That is, the load shedding is more costly.

We use an open source load forecasting dataset from the Texas Backbone Power System (Lu et al. 2023) on a modified IEEE bus-14 system. We randomly collect 1.0k samples and use a feedforward neural network with three hidden layers to forecast the load⁴⁴4More experiment can be found in the appendix.. We do E2E-AT with a). input feature uncertainties $\bm{\delta}_{x}$, b). uncertainty of the unpredictable parameters $\bm{\delta}_{\phi}$, and c). integrated uncertainties of both $(\bm{\delta}_{x},\bm{\delta}_{\phi})$. We first implement natural (or clean) E2E learning, based on which we warm-start the E2E-ATs. Adam optimizer is used, and ‘Adversarial training for free’ (Shafahi et al. 2019) is applied to reuse the gradients for PGD.

In detail, a). Since meteorological features have been normalized into $[0,1]$, we attack with a normalized budget $\epsilon_{x}\in\{0.02,0.03\}$. The inner maximization is solved with 7 PGD steps and the step size is dynamically set as $\epsilon_{x}/7\times 2$. We summarize this setting from (Zhang et al. 2019). We clamp the attacked input into $[0,1]$ whenever it is updated. b). We consider uncertainties on the susceptance $\bm{b}$ in the redispatch problem. Since each transmission line can have different nominal susceptance, we set the budget $\epsilon_{\phi}\in\{0.05,0.15\}$ as the proportion to the individual nominal value, which is consistent with the common operation range of susceptance (Xu, Jaimoukha, and Teng 2022). c). We do E2E-AT with $(\epsilon_{x},\epsilon_{\phi})\in\{(0.02,0.05),(0.03,0.15)\}$ for the integrated uncertainties. The other settings are the same as in (a) and (b).

Multi-run adversarial attacks are evaluated in Table 1. For each sample and attack scenario, we randomly select three starting points within the attack budget and report the worst task-aware cost (22) to reduce the variance. Specifically, five training algorithms are compared. NAT: natural training with task-aware loss; AT-MSE: adversarial training with MSE loss; AT-INPUT: adversarial training with task-aware loss and input uncertainties; AT-PARA: adversarial training with task-aware loss and unpredictable CO uncertainties; and AT-BOTH: adversarial training with task-aware loss and integrated uncertainties.

We first highlight that uncertainties in COs can significantly increase the cost, e.g. 15 times higher than the clean cost when $\epsilon_{\phi}=0.15$ for NAT. The AT-MSE performs better for input attacks, compared to the NAT, but performs poorly on the CO attacks. This is because AT-MSE is only trained with input uncertainties. Second, the performance of E2E-AT is similar to conventional adversarial training for image classification (Madry et al. 2017; Zhang et al. 2019). For instance, training with a larger attack budget can result in better robustness on attacks with a smaller budget but can inevitably increase the clean cost. Meanwhile, the hyperparameter $\alpha$ gives a trade-off between clean and robust accuracy in most cases.

In addition to common findings on adversarial robustness of image tasks, some unique findings of E2E-AT are highlighted. First, AT-INPUT and AT-PARA are more effective in the uncertainty with which they are trained. However, it is observed that E2E-AT based on one source of uncertainty can also improve the robustness of the other. For example, in AT-PARA, the cost of the input attack is even lower than that trained by AT-INPUT, when $\epsilon_{x}$ is small. Moreover, AT-PARA is more effective in the integrated attack than AT-INPUT. In fact, any input uncertainty eventually feeds into the COs, which becomes the uncertainties of the predictable parameters in COs. Finally, AT-BOTH not only improves the robustness of integrated uncertainty, but improves the robustness of the individual’s. All of the findings demonstrate that the uncertainties of both sources can be treated together in a unified way.

Certification on the robustness of the input space is considered, as susceptance $\bm{b}$ is coupled with the decision variable and Proposition 2 is not applicable. Using Proposition 2 and the mixed integer reformulation of NN (4), the exact input space adversarial attack on sample $(\bm{x},\bm{y})$ can be found by a mixed integer linear program (MILP). Interval bound propagation (IBP) (Gowal et al. 2018) is used to estimate the bounds of layers in NN. As for COs, we set $M=10^{5}$ by experience. Due to the large computation burden, we randomly sample 15 same samples and solve the robust certification using Gurobi.

First, the branch-and-bound algorithm is applied whose optimality is guaranteed if feasible. Meanwhile, the MILP formulation is verified as the same task-aware cost is achieved when solving the downstream COs parameterized by the optimal attack vector. Second, it can be observed that the exact attack vector causes the worse cost degradation, compared to the PGD-30 attacks. Finally, AT-INPUT can effectively reduce the task-aware cost on the exact input space attack.

As shown in (19), the uncertainties of CO can be modeled by stochastic CO. To model stochastic susceptance $\bm{b}$, we randomly alter the susceptance (random attack) for each sample and report the task-aware cost in Table 2. Under the random attack, the task-aware cost of NAT increases, clearly demonstrating the need for adversarial training. After E2E-AT, task-aware cost under random attacks is significantly reduced and is equal to the corresponding clean cost when the attack budget is small. Although the clean cost increases, it is still lower than the average cost under random attacks. The task-aware cost is also empirically upper bound by the adversarial attack (e.g. PGD-7), which verifies our argument on connecting the stochastic COs with E2E-AT (16) and (20).

The stage-one problem is a generator dispatch problem, which can be modelled as:

In (24), $\bm{P}_{g}$ and $\bm{\vartheta}$ are the vector of generator dispatch and voltage angle in each bus. A linear cost function is considered. All the equality and inequality signs in the constraints are element-wise. The first constraint represents the upper and lower bounds for each generator. The second constraint represents the thermal limit on the power flow in each transmission line where $\bm{b}$ is the line susceptance, $\bm{A}$ is the incidence matrix of the power system. The third constraint represents the requirement for power balance on each bus, where $\bm{C}_{g}$ and $\bm{C}_{l}$ are the incidence matrices of the generator and the load, respectively. $\hat{\bm{y}}$ represents the forecast load which is parameterized by the neural network model, for example, $\hat{\bm{y}}=f(\bm{x};\bm{\theta})$. We also add a slack variable $\bm{s}\geq\bm{0}$ as a compensation for infeasibility with large cost coefficient $c_{ls}$. The fifth constraint indicates that the system is referenced to the slack bus with constant phase angle.

When the generator has been dispatched as $\bm{P}_{g}^{\star}$, we consider a realization on the actual load $\bm{y}$ by solving the second stage problem (also known as redispatch problem):

The objective of stage two is to balance the load and to solve any violation of physical constraints of power grid using load shedding $\bm{P}_{ls}$ if the generator dispatch is lower than the actual load and energy storage $\bm{P}_{gs}$ if the generator dispatch is higher than the actual load. Since load shedding (similar to blackout) is more critical and should be avoided as much as possible, it is assigned by a larger penalty, i.e. $\bm{c}_{ls}\gg\bm{c}_{gs}$.

We highlight that our formulation on power system operation is more realistic than (Donti, Amos, and Kolter 2017), in which the behavior of different loads and network constraints are ignored. This results in more complex COs. In detail, the two COs have more than 60 decision variables and 150 constraints in total, which needs to be exactly solved for every forward pass.

Using Proposition 2 and the mixed integer reformulation of NN (4), an exact adversarial attack on input $(\bm{x},\bm{y})$ can be found by MILP. It is known that the value of the ‘big M’ used for NN linearization (e.g. the lower and upper bounds of the output of each NN layer (4), the upper bounds of the active inequality dual variables, and the lower bounds of the inactive inequality constraints of the COs) is essential to the performance of MILP solution. Therefore, interval bound propagation (IBP) (Gowal et al. 2018) is used to estimate the layers’ bounds in NN. As for COs, we set $M=10^{5}$ by experience.

Based on the experiment results in Table 1 and Table 3, we speculate that the robustness of one uncertainty can improve the robustness of the other, that is, the two uncertainties are not contradictory in E2E-AT. We thereby give a more detailed discussion.

To start, the two sources of uncertainties, e.g., input and CO uncertainties, are handled by training a same robust NN, which is a parametric model to forecast the predictable parameter of CO. According to the sensitivity analysis (Boyd and Vandenberghe 2004), a robust E2E model under uncertain COs implies that the robust NN can forecast a parameter such that the activations of constraints will not be significantly changed. From the structure of E2E learning, the input uncertainty is amplified by the NN (Gowal et al. 2018), which becomes an additional uncertain parameter of the COs. Similarly, an input-robust E2E model also requires that the NN forecast does not trigger significant constraint violations. Therefore, we speculate that the E2E-AT training objectives under input and CO uncertainties are not contradictory, as they are both reflected and controlled by the behavior of the COs.

However, the effectiveness of AT-INPUT on CO uncertainties (or AT-PARA on input uncertainties) depends on the individual attack budget, which will be discussed later by the extra experiment.