Technical Summary: Stochastic Penalty-Barrier Methods for Constrained Machine Learning

Problem Statement

The paper addresses the challenge of Constrained Machine Learning (CML), specifically the optimization of a stochastic, non-convex, and non-smooth objective function subject to stochastic, non-convex, and non-smooth constraints. The problem is formulated as:

$$\min_{x \in \mathbb{R}^n} \mathbb{E}[f(x, \xi)] \quad \text{s.t.} \quad \mathbb{E}[g(x, \xi)] \le 0$$

where $\xi$ is a random variable, and $f$ and $g$ are locally Lipschitz functions. This setting naturally arises in deep learning applications such as fairness-aware training (enforcing demographic parity or equalized odds), physics-informed neural networks (PINNs) where physical laws act as constraints, and the integration of symbolic domain knowledge.

Existing methods often fail in this specific regime. Prior works typically assume deterministic constraints, convex objectives, or linear constraints. While methods like Stochastic Switching Subgradient (SSw) and Stochastic Smoothed Lagrangian (SSL-ALM) handle some stochastic non-convex cases, they struggle with general non-smooth, non-convex constraints or require impractical computations (e.g., Hessian storage in SQP methods). Furthermore, simple penalty regularization (e.g., $F_\rho(x) = f(x) + \rho \| \mathbb{E}[g] \|^2$) is often unstable for large penalty parameters or infeasible for small ones.

Methodology: Stochastic Penalty-Barrier Method (SPBM)

The authors propose SPBM, a primal-dual algorithm that extends the classical Penalty-Barrier Method (PBM) of Ben-Tal and Zibulevsky [20] to the stochastic, non-smooth, non-convex setting. The method unifies constraints using penalty-barrier functions and employs three key modifications to handle the stochastic and non-smooth nature of deep learning:

1. Exponential Moving Average (EMA) for Dual Updates:
   Since expectations are approximated via mini-batching, the dual variable updates are noisy. SPBM applies an exponential moving average to the dual updates ($\lambda$) to reduce variance and stabilize convergence. The update rule is:
   $$\lambda^{k+1} = \gamma \lambda^k + (1-\gamma) \lambda^k \odot \phi'(g(x^k) \oslash p^k)$$
   where $\gamma \in [0, 1)$ is a decay rate, $\phi$ is the penalty-barrier function, and $p$ represents penalty parameters.

2. Stabilized Penalty Parameter Schedule:
   In the original PBM, penalty parameters $p$ decrease towards zero, which can cause instability in the stochastic setting. SPBM introduces adaptive or linear penalty update functions that prevent $p$ from vanishing too quickly. The authors propose two strategies:

   • Linear: $p^{k+1} = \text{Clip}_{[0.1, 1]}[K_1 p^k]$ with $K_1 < 1$.
   • Adaptive: Adjusts $p$ based on the current satisfaction of the constraint (using $\phi'$), increasing penalties for satisfied constraints and decreasing them for violated ones, while maintaining bounds.

3. Moreau Envelope for Non-Smoothness:
   To handle the non-smoothness of the objective and constraints (common in neural networks with ReLU activations or specific loss functions), SPBM optimizes over the Moreau envelope of the Lagrangian. Instead of solving the exact non-smooth subproblem, the method performs a single step of Adam (OneStepAdam) to find an approximate minimizer of the Moreau envelope:
   $$\min_x \left( \bar{L}_{p^{k+1}}(x, \lambda^{k+1}) + \frac{\mu}{2} \|x - s^k\|^2 \right)$$
   where $s^k$ is a prox-center updated via exponential averaging. This effectively smooths the problem while preserving the minimizers.

The algorithm operates in a mini-batch setting, sampling both objective and constraint distributions simultaneously. It is designed to be a "drop-in" replacement for standard optimizers in frameworks like PyTorch.

Key Contributions

• Algorithmic Proposal: Introduction of SPBM, the first general method tailored for stochastic, non-convex, and non-smooth constrained optimization in deep learning.
• Theoretical Extension: Generalization of the deterministic PBM [20] to the stochastic setting, showing that SPBM can be viewed as a generalization of SSL-ALM [19] under specific parameter choices.
• Empirical Benchmarking: Comprehensive evaluation across 8 distinct experimental settings, including:
  • Fairness: Binary and multi-class classification on ACSIncome, Dutch Census, CIFAR-10, and CIFAR-100 datasets with demographic parity and equalized accuracy constraints.
  • Physics-Informed: Solving the Helmholtz and viscous Burgers equations using PINNs with boundary and initial value constraints.
• Open Source Implementation: Release of code that integrates seamlessly with PyTorch, requiring minimal code changes to specify constraints.

Experimental Results

The authors benchmark SPBM against Adam (unconstrained), SSw [18], and SSL-ALM [19].

• Performance: SPBM matches or outperforms existing baselines in terms of test loss and constraint satisfaction.
  • In fairness tasks (e.g., CIFAR-100 with 9,900 constraints), SPBM achieved significantly lower test loss than SSL-ALM while satisfying constraints.
  • In PINN experiments, SPBM demonstrated superior solution quality (lower test MSE) and better satisfaction of boundary conditions compared to Adam with regularization and SSL-ALM.
  • SSw generally underperformed, attributed to its limitation of handling only one constraint at a time, making it ill-suited for multi-constraint settings.
• Computational Efficiency: SPBM incurs only a linear runtime overhead compared to unconstrained Adam. For up to 10,000 constraints, the runtime was at most a factor of 3 slower than Adam (e.g., 5.86s vs 1.92s per epoch on CIFAR-100).
• Stability: Unlike standard penalty methods which require careful tuning of $\rho$ to avoid instability or infeasibility, SPBM was observed to be stable and converge to the feasible set across various hyperparameter settings.

Significance and Claims

The paper claims that SPBM provides a reliable and computationally feasible alternative for training neural networks under complex constraints. Its significance lies in:

1. Bridging the Gap: It addresses the "non-convex, non-smooth, stochastic" setting that is prevalent in modern deep learning but previously lacked a general, efficient solver.
2. Practicality: By maintaining linear computational scaling and offering a simple integration into existing ML workflows, it enables practitioners to enforce fairness or physical constraints without prohibitive costs.
3. Robustness: The method demonstrates robustness in high-dimensional constraint spaces (up to ~10,000 constraints) where other methods struggle.

The authors modestly note limitations: convergence guarantees in the fully stochastic non-convex non-smooth setting remain an open theoretical question, and the method introduces additional hyperparameters (e.g., penalty schedules) that may require tuning. The evaluation is currently limited to fairness and PINN applications.