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

Reply to Reviewer MUif

We thank the reviewer for their comments. 1. We have used the bibliographystyle\{iclr2025\_conference\} for citations and will re-check and update the reference format as per guidelines. 2. We would like to request the reviewer to expand more on this to specify what improvements can be made by the authors for the reviewer to adjust the score. Additionally, we have decided to update the manuscript with one motivation diagram and some additional results on different dataset comparing with other state of art methods like DC3. Detailed results, to be updated in the revised manuscript, can be found on: https://drive.google.com/file/d/1Y57JVPegi2HY2krnLb7qihMvQNdWvfdg/view?usp=share_link 3. Our focus was on ensuring feasibility without compromising optimality in settings constrained by limited training data and time. The goal was to develop an optimization proxy for practical yet challenging scenarios where labeled data and total training and testing times are restricted (see Figure 1: (https://drive.google.com/file/d/1Y57JVPegi2HY2krnLb7qihMvQNdWvfdg/view?usp=share_link)). This differs from scenarios with abundant training samples or extended training ($T_{\text{training}}$) and testing ($T_{\text{prediction}}$) times. Notably, $T_{\text{prediction}}$ for our model is at least **10 times lower** than approaches like DC3 or Zumzum & Baker's, as it avoids projection steps requiring power flow solvers during training or testing. Additionally, our results show up to an order of magnitude improvement in feasibility gaps without compromising optimality. For example, the optimality gap for case118 reduces to **0.089** with our method, compared to **1.284** for the best point prediction method (Table 2). Our BNN-based proxy also generates meaningful confidence bounds on prediction accuracy despite using only 1,000 validation samples. As shown in Figure 4, prior DNN-based models, even with the same validation data and using Hoeffding's or Empirical Bernstein inequalities, fail to provide meaningful bounds. Self-supervised methods like primal-dual learning (Park & Van Hentenryck, 2023) require significantly longer training times, e.g., **5,932.5 seconds** for the 57-bus system and **7,605.1 seconds** for the 118-bus system. In contrast, our method trains within **600 seconds**, making their training times 10 times longer. For this reason, we did not use such methods for benchmarking. 4. We included 1d to be explicit about the input to the optimization problem. We will move it to the paragraph below equation (1) in the revised manuscript. 5. We respectfully disagree with the reviewer’s concern regarding the computational intensiveness of constructing the feasibility dataset $D_f$. The process of creating $D_f$ is, in fact, computationally light because it relies solely on the necessary conditions for feasibility in constrained optimization problems. For any feasible solution to a constrained optimization problem, the constraints must be satisfied exactly—resulting in zero constraint violations or "gap" as defined in equation (2). Thus, by definition, the feasibility dataset $D_f$ is constructed such that each entry in $D_f $ corresponds to a constraint satisfaction value of zero. Specifically, given any input $ \mathbf{x}$ in $ D_f $, the output is identically zero, indicating that all constraints are met. Therefore, constructing $ D_f $ is as computationally inexpensive as simply sampling inputs, since there is no additional cost for evaluating or solving complex constraint satisfaction conditions. The feasibility requirement simplifies the dataset creation process, ensuring that the computational load remains minimal. This part is mentioned in paragraph below equation (2) of the manuscript. 6. As mentioned before, the proposed works motivation is to learn an ACOPF proxy under limited training data and training time situations. Therefore, methods which takes considerably large amount of time to train (approx. 6000 sec. for by Park & Van Hentenryck (2023) for 57-Bus system) does not fit within the motivation of this work. Moreover, an important difference between methods such as DC3 and Zamzam & Baker (2020) is presence of power flow in the pipeline during training and prediction. The inclusion of power flow implies that we need to solve nonlinear equations within the proxy to obtain accurate predictions. This leads to considerable increase in both training and prediction time. For example, in DC3 paper it is listed that it takes 0.089 sec. to predict the ACOPF solution for one instance, which is approximately 30 times higher than 0.003 sec prediction or inference time of the proposed method. This limits the use of use of these models in stochastic settings such as probabilistic risk quantification where a large number of (order of millions) predictions are needed.

Reply to Reviewer MUif (Point 7 onwards)

7. We would like to highlight that the manuscript includes five different state-of-art baselines, for all the systems. Thus, we disagree that baselines were omitted. Also, as noted in the paper and before, self-supervised methods are very time consuming in training (Authors in Park & Van Hentenryck (2023) report that **5932.5 seconds** are taken to train models on 57-bus system and **7605.1 seconds** are taken to train on 118-bus systems.) The primary motivation of this paper is to develop an Optimization proxy under a practical but challenging situation where total labeled data for training as well as total time for training and testing are constrained. Moreover, to strengthen the discussion, we will update the comparative results in the revised manuscript, which can be found here: https://drive.google.com/file/d/1Y57JVPegi2HY2krnLb7qihMvQNdWvfdg/view?usp=share_link 8. Yes we agree that we have not accounted for time for generating 512 training samples. However, we do define the total time to be data generation + training time, in which data generation time is constant for all different models. We have not commented on data generation time as we have used open source **torch_geometric** dataset which does not provide generation time. However, in the revised manuscript, we will update the training data generation time using **PowerModels.jl** package. It takes on average 0.15 sec. to solve an instance of ACOPF for 57-Bus and 0.45 sec. for 118-Bus system ACOPF. Therefore, it takes 45.36 sec to generate training + validation set for case57 and 136.08 for case118, on a five core CPU. In practice, this data is can also be available from historical operations. Further, it is clear that these data generation time is much slower than training time of 600 sec. opted in this work. 9. Yes, the proposed method requires validation samples to generate the gap and constraint violation. The statement was written to highlight that unsupervised methods are not completely free from labeled data requirements. We understand that this might be misleading, hence we will qualify the statement as follows: ``Moreover, unsupervised methods, similar to supervised and proposed method also require validation data and consequently incur data generation time. This validation is needed to provide confidence bounds on error with respect to true solution." However, we want to highlight that for the same set of validation data, DNN based models produce trivial confidence bounds due to the use of only Hoeffding's inequality as highlighted in section 4 of the paper. Our method is able to leverage the sampling of weights inside the BNN to obtain better confidence bounds through the use of Bernstein inequality. On Mistakes: We thank the reviewer for careful reading. We will proof read and revise the manuscript. A standard approach in the context of of mean prediction is written to reflect that we do not opt for weighted mean approach, instead go with most used approach of exception calculation for mean posterior. We will qualify the statements to reflect the same in revised manuscript.

Reply to Reviewer KQjj

We thank the reviewer for their comments. Point 1: Reviewer is correct to say that semi-supervised training approach proposed in this work can be applied to any NN. However, we want to argue that with limited training data, A BNN is better suited for learning ACOPF than NNs and DNNs. This argument can be corroborated from results presented in the paper where simple BNN always outperforms a NN with both MSE and MAE loss function (Comparative result Table 1 and 2 in Section 5 and Table 6 and 7 in Appendix C). Moreover, when using NN we do not have access to predictive uncertainty and thus cannot use the superior theoretical Bernstein inequality for error bounds. -- Yes, training BNN is more expensive than regular NN. However, in the present works setting, the training is constrained by the time and number of data points available are very limited. In such settings, the BNN outperforms NNs as shown in the results. We have adopted a setting motivated by practical considerations where there is a hard upper bound on the total time available. This is a different setting where the theoretical complexity of training may not be as relevant compared to what is the best that can be done in limited time. Please find details of motivation and additional results here: https://drive.google.com/file/d/1Y57JVPegi2HY2krnLb7qihMvQNdWvfdg/view?usp=share_link -- Firstly, we want to highlight that for selection via posterior strategy, reduction of the constraint violation is only one of the criteria. One can easily design a similar strategy for minimum cost, minimum feasibility error or a weighted combination of both. BNNs provide a principled way to quantify the uncertainty in the predicted output, which is not the same as adding arbitrary noise to NNs. In BNNs, the weights are probabilistic and each weight of the network has a non trivial posterior distribution which is updated during training. The distribution of the output is then a consequence of the weight distribution. The significantly larger probabilistic search space afforded by the weight distributions (often bimodal in our case) can lead to an output distribution that very different from Gaussian, even when the weight distributions started with a Gaussian prior. Therefore, using Gaussian noise to the output in NNs will not correspond it the same thing. For uncertainty quantification through NNs, a prominent approach is to use Ensemble of NNs, which provide output distribution via multiple possible weights. However, training large number of NNs require a lot of training time+ resource along with variations in training dataset. Thus, in the current works limited time and data setting it won't be feasible and BNN based approach is the preferred method both for accuracy and confidence bound generation. -- The hypothesis of using MPV as a proxy for TVE is motivated by our empirical observations. For ACOPF proxies, we observed that two times MPV is always greater than the total variance in error. Furthermore, as explained from line 294 onward, the first term of the total variance in error is the same as MPV. If the second term of the total variance (expectation over weights and variance over samples) is lower than MPV, our hypothesis holds. For ACOPF, we observed that the second term of TVE (Equation 5) is significantly lower than MPV (the first term of TVE). Figure 3 has been included to validate the hypothesis for various power system test cases. Note that Figure 3 shows results for both supervised and sandwich BNNs, covering a total of six instances where the hypothesis holds.We have also observed the same characteristic of MPV in other test cases. We will clarify and update the hypothesis discussion in the revised manuscript. -- Our confidence bound is similar to Proposition 5 in [1]. In the premise of Proposition 5, an absolute bound on the error (C) and on the second moment ($\gamma^2$) is assumed and the result is derived using the theoretical Bernstein inequality. In our setting, by using the automatically generated MPV from the training phase of BNNs along with hypothesis in line 286 or Eq 5 (clarified in our response to the previous question) we are able to obtain a similar confidence bound. Note that in either case, this bound is to help quantify the accuracy over i.i.d. out of sample test data. Section 4 in our paper is not a attempting to generalize the training accuracy, that is complicated due to dependence between model and the training samples. We will cite [1] in the revised document.

Reply to Reviewer KQjj (Point 2 Onwards)

Point 2: We want to clarify that one cannot combine supervised and unsupervised losses in Bayesian Neural Networks (BNNs) without having a tight lower bound on the cost function with unlabeled input data. This is because BNNs are trained using the Bayesian inference principle, where a likelihood function needs to be defined for the output. In the unsupervised learning stage, the proposed method relies on the fact that the observation for the feasibility gap will be zero for any input. This allows the definition of a likelihood for constraint satisfaction that attempts to achieve a distribution centered around zero (although a perfect delta distribution is not achievable due to numerical issues). To construct a similar likelihood for the cost (as done in the supervised training phase), we require tight bounds on the cost for unlabeled data. The setting of traditional NNs is different, as one can define a loss function that combines cost and weighted feasibility terms. These NNs are often referred to as Physics-Informed Neural Networks (PINNs). We have compared our proposed approach with these models (MAE + penalty and MSE + penalty) across various test cases presented in the paper. The results show that the proposed BNN approach outperforms these PINNs. Point 3: The section 5 of the manuscript provide separate contribution of each of the components. For example, first row of table 1 shows improvement via Sandwich BNN SvP compared to the Sandwich BNN where both cost function and feasibility gap improves. Similarly, comparing the Sandwich BNN and supervised BNN results indicate benefits of Sandwiching. We will update the discussion on the revised paper to highlight contribution of each part separately. Point 4: As mentioned in the manuscript, the proposed work's motivation is to learn an ACOPF proxy under limited training data and training time situations. Therefore, methods which takes considerably large amount of time to train (approx. 6000 sec. for by Park & Van Hentenryck (2023) for 57-Bus system) does not fit within the motivation of this work. Moreover, an important difference between proposed approach and methods such as DC3 and Zamzam & Baker (2020) is presence of power flow in the pipeline during training and prediction. The inclusion of power flow implies that we need to solve nonlinear equations within the proxy to obtain accurate predictions. This leads to considerable increase in both training and prediction time. For example, in DC3 paper it is listed that it takes 0.089 sec. to predict the ACOPF solution for one instance, which is approximately 30 times higher than 0.003 sec prediction or inference time of the proposed method (with 100 weight samples). This limits the use of use of these models in stochastic settings such as probabilistic risk quantification where a large number of (order of millions) predictions are needed. Detailed results, to be updated in the revised manuscript, can be found on : https://drive.google.com/file/d/1Y57JVPegi2HY2krnLb7qihMvQNdWvfdg/view?usp=share_link Similar to methods such as DC3 and Zamzam & Baker (2020), the work Huang, W., & Chen, M. (2021) also have power flow in the pipeline during prediction. The inclusion of power flow implies that we need to solve nonlinear equations within the proxy to obtain accurate predictions. This leads to considerable increase in both training and prediction time. For example, in DC3 paper it is listed that it takes 0.089 sec. to predict the ACOPF solution for one instance, which is approximately 30 times higher than 0.003 sec prediction or inference time of the proposed method (with 100 weight samples). This limits the use of use of these models in stochastic settings such as probabilistic risk quantification where a large number of (order of millions) predictions are needed. Questions: 1. We have observed that the proposed method also follows the general principle where the prediction error decreases with increase in training data. However, when the training time is limited, increasing the number of training samples beyond a certain threshold does not provide additional improvements. In our experiments, we observed that for systems larger than 118-Bus, increasing training samples beyond 1000 does not provide further benefits when the training time is limited to 600 sec. We will append these results to the revised manuscript. 2. The effect of number of posterior samples on output is presented in the Figure 5, Appendix C. It shows that errors in various parameters stabilizes around 200 posterior samples.

Reply to Reviewer KQjj : On Motivation

The key advantage that BNNs—or Bayesian methods in general—offer over DNNs and other deterministic methods in any learning problem (not just optimal power flow) lies in their performance under limited data conditions. With a sufficiently large dataset, DNNs or other universal approximation methods should perform equally well, if not better. However, as stated in the Introduction, the realistic regime for power grid optimization problems involves limited data and constrained training time. Another perspective when comparing BNNs and DNNs is that BNNs can separate two types of uncertainty **Epistemic uncertainty**, which reflects the uncertainty in model parameters ($p(w|D)$, where $w$ represents the parameters and $D$ is the training data). **Aleatoric uncertainty**, which represents the inherent uncertainty in the data itself ($p(y|x, w)$, where $y$ is the output and $x$ is the input). This separation allows BNNs to be highly data-efficient, enabling them to effectively learn from small datasets without overfitting [R1–R2]. During prediction, instead of confidently producing an incorrect result, BNNs assign high epistemic uncertainty to points far from the training dataset, signaling that the model lacks sufficient knowledge about them. This eliminates the need to rely on **Ensemble DNNs** to account for weight distributions and build model robustness. Consequently, BNNs save significant computational resources that would otherwise be required to train multiple DNN models to create an ensemble. Additionally, we have demonstrated that the inherent randomness in BNNs allows for the rapid development of much better confidence bounds on predicted outputs compared to DNNs on similarly sized test datasets. This is primarily because the mean predicted variance of a BNN better approximates the true variance of outputs than the empirical variance computed by DNN models. While this observation is still empirical, our results confirm its validity across a range of power grid networks, including large networks that have not been addressed in existing DNN-based studies in this domain. Assuming that the variance predictions are reliable, we show that our model reduces confidence error scaling from $1/\sqrt{M}$ to $1/M$, where $M$ is the number of test samples. Finally, we argue that the proposed work offers a novel perspective on ACOPF and optimization proxies. The ability of BNNs to quantify uncertainty in predictions facilitates effective active learning and Bayesian Optimization [R3]. This work lays the foundation for using BNNs as surrogates in various engineering problems, where active learning and/or Bayesian Optimization can be employed to sample more informative datasets. Therefore, this study should not only be evaluated from a state-of-the-art standpoint but also as a contribution that opens new directions in the field. **We will ensure that the manuscript is updated with these discussions and additional details.** [R1] Laurent Valentin Jospin et. al., "Hands-on Bayesian neural networks—A tutorial for deep learning users." IEEE Computational Intelligence Magazine 17, no. 2 (2022): 29-48 [R2]Stefan Depeweg et. al. "Decomposition of Uncertainty in Bayesian Deep Learning for Efficient and Risk-sensitive Learning", Proceedings of the 35th International Conference on Machine Learning, PMLR 80:1184-1193, 2018. [R3] Yucen Lily Li et. al. "A Study of Bayesian Neural Network Surrogates for Bayesian Optimization" ICLR 2024.

Reply to Reviewer uTQN

We thank the reviewer for their comments. Weakness: 1. The authors would kindly like to understand the reviewer's perspective on why the proposed work, which introduces a novel semi-supervised learning approach to solve constrained optimization problems in practically significant limited data and time settings, is perceived as an applied work with limited technical novelty. 2. Authors would like to know what uncertainty calibration analysis the reviewer is referring to. Figure 3 and subsequent results within the paper provide extensive evidence on the effectiveness of proposed method in providing variance information and variance in error. We have also used a multiplier (2 specifically) to avoid the issues for underestimation of variance, if any, due to the mean-field assumption in variational Bayesian inference. 3. We would like to argue that it is standard practice to assume that at least one feasible solution exists for every input when studying machine learning models for constraint optimization problems. All previous works on ML proxies including DC3, Park & Pascal, Zumzum & Kyri etc. work under this assumption. Even in practice, this is not a limiting assumption for ACOPF as power system is operating at a feasible solution at every instance. 4. A limitation of this work stems from its motivation that it is targeted to a specific setting of limited training data and time, and does not provide a general method which minimizes violation in general ACOPF proxy settings. Besides, a well known limitation is the high training time requirement by BNN, compared to DNNs. We will update the manuscript in the revision, explicitly mentioning these. We do however believe that despite its higher training complexity BNN is the better choice in the limited data and time setting. We would like to highlight that the paper already presents a lot of comparisons with existing methods, especially the ones that are relevant to the power systems setting. Additionally, we will update the manuscript with additional discussion and results, which can be found on this link: https://drive.google.com/file/d/1Y57JVPegi2HY2krnLb7qihMvQNdWvfdg/view?usp=share_link Again, we would like to emphasize that methods like ensembles will require a lot of training time, which is beyond the problem setting of this paper. We will include additional discussion in the introduction to clarify this point.

Reply to Reviewer Vowh

We thank the reviewer for their thoughtful comments. 1. Yes, our primary focus was on feasibility without compromising on optimality, in learning time and data constrained settings. The primary motivation of this paper is to develop an Optimization proxy under a practical but challenging situation where total labeled data for training as well as total time for training and testing are constrained, as shown in Figure 1 on this link: https://drive.google.com/file/d/1Y57JVPegi2HY2krnLb7qihMvQNdWvfdg/view?usp=share_link This is different from situations where a large number of training samples can be generated or much longer training time $T_{training}$ for model tuning and testing time $T_{prediction}$ for prediction are available. Technically the number of training data is itself a proxy for data generation time $T_{Data}$. In this context, note that $T_{prediction}$ for our model is at least **10 times lower** than the DC3 approach or Zumzum & Baker's approach because our model does not involve a projection step that requires solving power flows or access to a power flow solver during either training or testing. When compared to models that have fast testing times, our results show up to an order of magnitude improvement in feasibility gaps, without compromising the optimality of predictions. E.g.: the optimality gap reduces to 0.089 by proposed method from 1.284 best point prediction method for case118 in Table 2. Second, our proxy, due to the use of BNN, is able to generate non-trivial confidence bounds on the accuracy of prediction despite a limited number of 1000 validation data. As shown in Figure 4, prior DNN based models when using the same validation dataset, and Hoeffding's or Empirical Bernstein inequality, do not produce meaningful performance bounds. 2. We emphasize that our contributions extend beyond sandwich learning BNNs to include simple BNNs for constrained optimization and Bernstein inequality-based error guarantees. The results demonstrate that BNNs are well-suited for data- and time-constrained environments. For smaller power grids, the sandwich model effectively reduces power balance violations, achieving its intended goal. However, for larger systems, limited training time prevents the model from fully utilizing available information. Nonetheless, to maintain practicality, we uniformly restrict training to 10 minutes for all examples. 3. Yes, the DNN training used 512 samples and was capped at 600 seconds. The goal was to compare BNN and DNN performance under low data and limited training time conditions. While larger datasets could be constructed, the time required increases significantly with system size, as noted in [Peng-Pascal paper]. More critically, this added data creation time is a major drawback in power system operations, where datasets are often unavailable or difficult to obtain, especially for N-1 line outage scenarios. This setting aligns closely with the motivation of the proposed work. 4. Yes, the DNN model for 2000-bus has too many trainable parameters and is failing because of the very limited data and training time. Also, we agree that with more data, and training time, DNNs will start to perform better and at some point exceed the BNN performance (since DNNs can be trained more efficiently compared to BNN). However, we want to again highlight that our motivation is not to say that BNNs are better than DNNs in general. Our limited submission assertion is that-- under data and time constrained settings BNN based models outperform the DNN models for ACOPF problem, as described in motivation before. Questions: 1. We used 512 labeled data points for DNN training across all cases, as outlined in Section 5. 2. While this work focuses on BNNs we agree that sandwich learning has potential applicability for DNNs, as it could reduce reliance on labeled datasets. However, we wish to emphasize that using a DNN alone can only provide Hoeffding’s and Empirical Bernstein bounds on the expected error, as discussed in Section 4 of the paper. An intriguing direction, without losing the Bayesian nature of predictions, would be to use a DNN exclusively for the unsupervised part of training. However, this approach would require substantial effort in designing sequential priors for the supervised BNN stages that follow the unsupervised DNN stages. We plan to explore this direction in future work. 3. It is standard practice to train NNs with two hidden layers for the ACOPF problem. Almost all existing works adopt a two-hidden-layer architecture. An intuitive explanation for this stems from the fact that power flow equations are quadratic in nature, so two or three nonlinear transformations should suffice to capture the input-output relationship. While it cannot be ruled out that deeper networks might offer better performance, they are much harder to train within the constrained training time setting of the proposed method due to the significantly larger number of parameters.