Optimization Proxies using Limited Labeled Data and Training Time – A Semi-Supervised Bayesian Neural Network Approach

Constrained optimization problems are fundamental in the optimal operation of various engineering systems, such as supply chains, transportation networks, and power grids. Learning a forward mapping between the inputs and outputs of these problems can significantly reduce computational burdens, especially when rapid solutions are required, such as in electricity market clearing or real-time transportation planning.

Recent advancements in machine learning (ML) have led to considerable efforts to solve optimization problems using deep neural networks (DNNs) Khadivi et al. (2023); Kotary et al. (2021); Fajemisin et al. (2024). The idea of learning input-to-output mappings has been explored via supervised and unsupervised methods, particularly in power system applications Zamzam & Baker (2020); Donti et al. (2021); Park & Van Hentenryck (2023); Fioretto et al. (2020); Kotary et al. (2021); Rolnick et al. (2022). Additionally, constraint penalization approaches have been proposed to enforce feasibility in predicted outputs within DNN loss functions AI4OPT (2023).

Supervised DNN models rely on labeled datasets obtained by solving numerous instances of optimization problems. This data generation step poses a significant limitation due to the prohibitive computational time required for moderate to large problem instances. For example, Park & Van Hentenryck (2023) report that generating labeled data to train supervised DNN models for a medium-sized power grid problem takes over three hours¹¹1See Table 4 in Park & Van Hentenryck (2023) for the computation time comparison.. Unsupervised methods aim to address the labeled data generation issue Donti et al. (2021); Park & Van Hentenryck (2023); however, they often have high training time requirements. Methods requiring projection steps to recover feasible solutions involve time-consuming constraint correction within the framework Donti et al. (2021); Gupta et al. (2022); Zamzam & Baker (2020), resulting in slower training and prediction. Moreover, unsupervised methods still require testing data and consequently the associated data generation time in order to perform validation and provide confidence bounds on error with respect to true solution. Thus, it is important to note that that total time requirement to deploy these ML based optimization proxies is: $\texttt{Training Data Generation}+\texttt{Testing Data Generation}+\texttt{Training}$.

In practice, the input space of the problem in many applications is dynamic and changes constantly, necessitating frequent retraining and evaluation for ML models to remain effective amid distributional shifts in the input data. In the context of longer horizon planning problems involving bi-level optimization problems, these ML proxies serve as subroutines to simulate stochastic decision-making processes for given first-level decisions Ibrahim et al. (2020). Consequently, these models must be adaptable in both training and testing to changing problem inputs and parameters. Minimizing or limiting the time required for learning input/output mappings in optimization problems is thus crucial.

Estimating the generalization error over the testing dataset is another critical aspect, and is particularly relevant to planning and operations of engineered systems where the system must obey the physical and safety limits. As a result, when limited labeled data is available, one must rely on concentration bounds such as Hoeffding’s inequality Sridharan (2002); Hoeffding (1994) to develop expected error bounds using finite out-of-sample testing data. However, these bounds are often loose and impractical, creating the need for tighter expected error bounds with limited labeled testing data.

We start by defining a suitable likelihood function for the unsupervised learning process. To that end, we augment the unlabeled data $\mathcal{D}^{u}$ using the necessary feasibility conditions that the vector $\mathbf{y}$ must satisfy to be a solution of equation 1a. We propose a function $\mathcal{F}(\mathbf{y},\mathbf{x})$ to measure the feasibility of a candidate solution $\mathbf{y}$ for a given input $\mathbf{x}$. This function consists of two terms: one measuring the equality gap and the other measuring the one-sided inequality gap or violations. The relative emphasis on each term is determined by the parameters $\lambda_{e}$ and $\lambda_{i}$, respectively, i.e.,

This section focuses on providing bounds on the expected absolute error of our method, i.e., testing error. We explore probabilistic confidence bounds (PCBs) for optimization proxies. The core concept of PCBs is to first evaluate models on a labeled testing dataset with $M$ samples, compute the empirical mean error, and then probabilistically bound the error for any new input. Specifically, PCBs assert that the expected error will be within $\varepsilon$ of the empirical errors computed from $M$ out-of-sample inputs, with a high probability (usually 0.95). Mathematically, we aim to provide a guarantee on the error $e=y-y^{t}$, where $y^{t}$ is the BNN prediction and $y$ is the true value, as

where $\mathbb{E}[|e|]$ represents the expected absolute error, $1-\delta$ is the confidence level, and $\varepsilon$ is the allowable prediction error.

Ideally, we would like to evaluate our model on a large number of samples since, as $M\rightarrow\infty$, the error bound $\varepsilon\rightarrow 0$. However, increasing $M$ leads to a higher requirement for labeled data, which defeats the purpose of training using low labeled data ⁴⁴4Note that the total labeled data requirement is the sum of training and testing samples, i.e., $N+M$. To address this issue, confidence inequalities are commonly used to provide PCBs, with Hoeffding’s inequality (Hoeffding (1994)) being one of the most widely used bounds. As stated in Appendix D, Hoeffding’s inequality assumes that the error is bounded (i.e., $|e_{i}|\leqslant R$ for all $i$) and provides PCBs whose tightness is governed by $M$, with the relationship $\varepsilon=R\sqrt{\frac{\log\left({2}/{\delta}\right)}{2M}}.$ However, the Hoeffding bound can often be too loose to be practically relevant.

To improve upon this, we propose using Bernstein’s inequality (see Audibert et al. (2007)) as the concentration bound, which utilizes the total variance in error (TVE) information along with $M$, under the same bounded error assumption. The main challenge in using the Bernstein bound is obtaining the TVE without extensive out-of-sample testing. One possible solution is to use the empirical Bernstein bound as in Audibert et al. (2007); Mnih et al. (2008), which employs the empirical variance of the error $\widehat{\mathbb{V}}_{e}$, obtained from the same $M$ testing samples, and accounts for the error in variance estimation by modifying the theoretical Bernstein inequality. Mathematically, the PCB using the Empirical Bernstein inequality is

as given in Audibert et al. (2007); Mnih et al. (2008). Since the term under the square root depends on the empirical TVE $\widehat{\mathbb{V}}_{e}$ rather than $R$, the empirical Bernstein bound becomes tighter more quickly with increasing $M$ if $\widehat{\mathbb{V}}_{e}\ll R$.

To further tighen the bound, we propose using the Theoretical Bernstein bound Sridharan (2002) (Theorem 3 in Appendix D) with the Mean Predictive Variance (MPV) as a proxy for the TVE $\mathbb{V}_{e}$. The MPV is the mean of the predictive variance of testing samples, i.e., $\text{MPV}=\frac{1}{M}\sum_{k=1}^{M}\mathbb{V}_{W}\left[\mathbf{Y}_{i\cdot}^{k}\right],$ where variance $\mathbb{V}_{W}\left[\mathbf{Y}_{i\cdot}^{k}\right]$ for the $k^{th}$ test sample using entries of the posterior prediction matrix across columns, generated using posterior weights, for $i^{th}$ output variable. In principle, the MPV captures the expected variance in the predictions due to the posterior distribution of BNN weights. We hypothesize that with a constant multiplier $\alpha>1$,

where $\mathbb{E}_{M}$ and $\mathbb{E}_{W}$ denote expectations with respect to $M$ testing samples and posterior weight samples, respectively, and $\mathbb{V}_{M}[e]$ and $\mathbb{V}_{W}[e]$ denote the variance of the error with respect to $M$ testing samples and posterior weight samples, respectively. The equality in equation 5 follows from the law of total variance Blitzstein & Hwang (2019). $\mathbb{V}_{|e|}$ represents the variance of the absolute value of the error, which is lower than the variance of the error.

Notice that the first term of the TVE, $\mathbb{E}_{M}\left[\mathbb{V}_{W}[e]\right]$, is independent of the labeled testing dataset because the true output $y$ is constant with respect to posterior weight samples; thus, $\mathbb{V}_{W}[e]=\mathbb{V}_{W}[y-y^{t}]=\mathbb{V}_{W}[y^{t}]$. Furthermore, from the definition of MPV, we have $\text{MPV}=\mathbb{E}_{M}\left[\mathbb{V}_{W}[e]\right]$. As an example, if $\mathbb{E}_{W}\left[\mathbb{V}_{M}[e]\right]\leqslant\text{MPV}$, our hypothesis in equation 5 holds with $\alpha=2$. Consequently, we can use $2\times\text{MPV}$ as an upper bound for $\mathbb{V}_{e}$ in the Theoretical Bernstein bound (see Theorem 3 in Appendix D), which gives the error bound

which is better than the Empirical Bernstein bounds. The hypothesis in equation 5 and the corresponding constant $\alpha$ can be computed by using application specific information or performing a meta-study like in Section 5.

The strength of this approach is that we do not require labeled testing samples to calculate MPV, thus incurring no additional computational burden from generating labels. Also, note that this constitutes an advantage of BNNs since MPV information is readily available with BNNs but cannot be obtained from DNN-based optimization proxies.

In the next section, we perform a meta-study using different BNN models on different test cases to show the performance of our proposed learning architecture as well as demonstrate that hypothesis equation 5 indeed holds for the proposed optimization proxy learning problems.

We test the proposed method on the Alternating Current Optimal Power Flow (ACOPF) problem, essential for the economic operation of electrical power grids Molzahn et al. (2019). Efficient ACOPF proxies can mitigate climate change by enabling higher renewable energy integration, improving system efficiency by minimizing losses and emissions, and enhancing grid resiliency against extreme weather conditions Rolnick et al. (2022). ACOPF is a constrained optimization problem with nonlinear equality constraints and double-sided inequality bounds. It aims to find the most cost-effective generator set points while satisfying demand and adhering to physical and engineering constraints. The inputs are active and reactive power demands; outputs include generator settings, voltage magnitudes, and phase angles at each bus. We adopt the standard ACOPF formulation Babaeinejadsarookolaee et al. (2019); Park & Van Hentenryck (2023); Coffrin et al. (2018) and benchmark our method using the open-source OPFDataset from Torch_Geometric, which provides numerous solved ACOPF instances (see Lovett et al. (2024)).

In our results, ‘Gap%’ denotes the average relative optimality gap compared to the objective values in the labeled testing instances. ‘Max Eq.’ and ‘Mean Eq.’ represent the maximum and mean equality gaps over all equality constraints, while ‘Max Ineq.’ and ‘Mean Ineq.’ indicate the same for inequality gaps in per unit, all averaged over testing instances. We compare our proposed method with the following state-of-the-art baseline supervised learning models available in the literature using the same labeled dataset and training time constraints, utilizing AI4OPT’s ML4OPF package AI4OPT (2023) (network and hyper-parameter details are in Appendix B):

• •

  Naïve MAE and Naïve MSE (Supervised): Use $l_{1}$-norm and $l_{2}$-norm loss functions, respectively, to measure differences between predicted and actual optimal solutions Park & Van Hentenryck (2023), incorporating a bound repair layer with a sigmoid function. The bound repair layer ensures that inequality constraints are always satisfied.

• •

  MAE + Penalty, MSE + Penalty, and LD (Supervised): Add penalty terms for constraint violations to the naïve MAE or MSE loss functions Park & Van Hentenryck (2023). The Lagrangian Duality (LD) method applies the $l_{1}$-norm as outlined in Fioretto et al. (2020); Park & Van Hentenryck (2023), and also uses a bound repair layer with a sigmoid function.

We exclude self-supervised constrained optimization methods like Primal-Dual Learning (PDL) Park & Van Hentenryck (2023) and DC3 Donti et al. (2021) due to their significantly higher training times and computational demands, which violate the premise of this paper. For example, PDL requires over 125 minutes of training time for the ACOPF problem on a 118-node power network using a Tesla RTX6000 GPU. Methods that require solving alternating current power flow to recover solutions (e.g., Zamzam & Baker (2020)) are also excluded, as they result in high prediction times compared to DNN or BNN forward passes⁵⁵5see prediction time studies in Donti et al. (2021) which suggest 10 times higher prediction time with power flow based projections (0̃.080 sec. compared to 0.001 sec. for DNN and 0.003 sec. per testing instance for proposed BNNs)..

To demonstrate the effectiveness of the proposed BNN learning methods, we conduct simulation studies on both our models and the ML4OPF models using an M1 Max CPU with 32 GB RAM, without any GPU. This setup highlights performance improvements due to the learning mechanism rather than computational power. We propose two classes of different models as:

• •

  Supervised BNN and Supervised BNN SvP: Standard BNN learning with labeled data, utilizing mean prediction and Selection via Posterior (SvP), respectively. The network uses ReLU activation and no bound repair layer.

• •

  Sandwich BNN and Sandwich BNN SvP: The proposed Sandwich BNN trained with labeled and unlabeled data as discussed in Section 3. Unsupervised training utilizes four times the number of labeled data samples. The network uses ReLU activation and no bound repair layer.

Importantly, we intentionally keep the BNN architecture unoptimized, constructing it using four different sub-networks (one for each of the four ACOPF outputs: real power generation, reactive power generation, voltage magnitude, and voltage angle). Each sub-network has two hidden layers with the number of hidden neurons equal to 2 $\times$ input size. In the Sup stage of the Sandwich BNN, both weights and biases are updated, whereas in the UnSup stage, only the weights are modified via SVI. Best model out of five random trials is selected.

Tables 1 and 2 present the comparative performance of different methods for solving the ACOPF problem on the ‘case57’ and ‘case118’ test cases, containing 57 and 118 nodes, respectively. For ’case57’, the Sandwich BNN SvP method achieves the best Max Eq. performance (0.027), outperforming all other methods, including the standard Sandwich BNN, without compromising other metrics. All the methods proposed in this paper outperform the best DNN results, typically achieved with the LD+MAE model (last row in the tables). Similar trends are observed for the ‘case118’ as well.⁶⁶6See Appendix C for similar tabular results and discussion of trends on larger test instances.

It is important to contextualize the significance of these numerical improvements. In ACOPF problems, cost values are in USD, with the mean cost for ‘case118’ being $97,000 or 9.7 in the per-unit system. Therefore, a 1% Gap corresponds to an expected difference of $970 across the testing instances. A ‘Max Eq.’ value of 0.08 implies a maximum expected power imbalance of 8.0 Megawatts among all 118 nodes of ‘case118’. Thus, reducing the Max Eq. from 1.284 with the LD+MAE model to 0.089 with our Sandwich BNN SvP model represents a significant improvement.

Figure 2 illustrates the growth of various metrics with increasing system size while keeping training resources constant. The proposed BNN methods exhibit significantly lower scaling in the expected maximum equality gap. Although Sandwich BNN SvP shows slightly higher scaling inequality gaps, the relative improvement in Max Eq. far outweighs these minor drawbacks. Notably, for ‘case500’ (see Table 6 in the Appendix C), the expected maximum power imbalance is below 0.5% of the mean real power demand ($1.7\times 10^{4}$ MW) and power grids already are equipped with spinning reserves (see Ela et al. (2011)) that have reserve capacity to handle these imbalances. This is a significant improvement over the DNN models which have much higher ‘Mean Eq.’ values. Moreover, the ‘Max Ineq.’ growth could be easily suppressed by incorporating bound repair layers, as used in DNN models in ML4OPF AI4OPT (2023).

Next, we present results for Proabilistic Confidence bounds, described in Section 4. Figure 3 shows that $2\cdot\text{MPV}$ consistently serves as an empirical upper bound for total variance in error, validating its robustness across models and system configurations⁷⁷7Total variance in error is assumed to stabilizing with 1000 testing samples (see Figure 6 in Appendix C).. In Figure 4, the theoretical Bernstein bounds using $\mathbb{V}_{e}=2\cdot\text{MPV}$ provide tight, practical bounds, whereas Hoeffding’s bounds are overly conservative and not useful for grid operations. For example, in ‘case118’, the Bernstein bound ensures a probabilistic guarantee on voltage constraint satisfaction, such that a predicted voltage between 0.91–1.09 pu guarantees no violations within the ACOPF limits of 0.90–1.10 pu, i.e., the maximum value of the Bernstein bound on the error is 0.010 pu across all nodes. Compared to Hoeffding’s bound ($0.064$ pu) or the empirical Bernstein bound ($0.018$ pu), the Bernstein bound ($0.010$ pu) is far tighter and more practical, highlighting the benefits of BNNs for optimization proxies. The error bounds for the ‘case500’ is provided in Fig. 6 of the Appendix C. Finally, we note that both Hoeffding’s and empirical Bernstein bounds can also be obtained by testing DNN models across $M$ samples⁸⁸8A form of MPV can be obtained via Ensemble DNNs Ganaie et al. (2022), however, it will lead to very high computational requirement compared to the BNN..

In conclusion, this paper introduces a semi-supervised Bayesian Neural Network (BNN) approach to address the challenges of high labeled data requirements and limited training time in learning input-to-output maps for constrained optimization problems. The proposed Sandwich BNN method incorporates unlabeled data through input data augmentation, ensuring constraint feasibility without relying on a large number of labeled instances. We provide tight confidence bounds by utilizing Bernstein’s inequality, enhancing the method’s practical applicability. Results show that BNNs outperform DNNs in low-data, low-compute settings, and the Sandwich BNN more effectively enforces feasibility without additional computational costs compared to supervised BNNs.