Box-Constrained Softmax

• Box-Constrained Softmax is a generalization of softmax that enforces strict lower and upper probability bounds.
• It solves a variational optimization problem using KKT conditions to yield an exact, closed-form solution with efficient computation.
• It is applied for fairness-aware classification and safety-critical systems, ensuring calibrated outputs with hard reliability guarantees.

Box-constrained softmax (BCSoftmax) is a generalization of the softmax function designed to explicitly enforce lower and upper bounds—termed box constraints—on the output probabilities of a model. Unlike the classical softmax, which provides only soft, parametric control via a temperature parameter, BCSoftmax ensures every output coordinate strictly adheres to user-specified constraints, making it suitable for applications requiring hard reliability guarantees, such as fairness-aware classification and safety-critical decision-making (Atarashi et al., 12 Jun 2025).

1. Mathematical Formulation and Variational Characterization

The standard softmax with temperature $\tau>0$ maps logits $x \in \mathbb{R}^K$ to the probability simplex $\Delta_K$ via

$$\mathrm{Softmax}_\tau(x)[i] = \frac{\exp(x_i/\tau)}{\sum_k \exp(x_k/\tau)}.$$

Softmax can be interpreted as the unique solution to the variational optimization problem:

$$\mathrm{Softmax}_\tau(x) = \arg\max_{y\in\Delta_K} \left\{ x^\top y - \tau \sum_{k} y_k \log y_k \right\}.$$

BCSoftmax extends this by imposing per-coordinate lower and upper bounds $a = (a_1,\ldots,a_K)$ and $b = (b_1,\ldots,b_K)$ with $0 \leq a_k \leq b_k \leq 1$ and $\sum_k a_k \leq 1 \leq \sum_k b_k$. The box-constrained softmax is defined as

$$\mathrm{BCSoftmax}_\tau(x; (a, b)) = \arg\max_{y \in \Delta_K,\, a \preceq y \preceq b} \left\{ x^\top y - \tau \sum_k y_k \log y_k \right\}.$$

An equivalent minimization (cross-entropy) form is

$x \in \mathbb{R}^K$0

where $x \in \mathbb{R}^K$1 is a transformed version of the input logits.

This formulation ensures that solutions always lie in the intersection of the simplex and the box $x \in \mathbb{R}^K$2, making it possible to enforce strict bounds on each output coordinate (Atarashi et al., 12 Jun 2025).

2. Exact Solution: KKT Conditions and Algorithmic Realization

Although the BCSoftmax evaluation is a convex program, the Karush–Kuhn–Tucker (KKT) conditions yield an exact closed-form characterization. Introducing Lagrange multipliers for the constraints, one establishes that, for each coordinate $x \in \mathbb{R}^K$3:

• $x \in \mathbb{R}^K$4 if the lower bound is active,
• $x \in \mathbb{R}^K$5 if the upper bound is active,
• $x \in \mathbb{R}^K$6 for free indices, where $x \in \mathbb{R}^K$7 is an appropriate normalization constant.

More precisely, there exists $x \in \mathbb{R}^K$8 such that

$x \in \mathbb{R}^K$9

with $\Delta_K$0.

The active set is determined by sorting the ratios $\Delta_K$1 (descending) and $\Delta_K$2 (ascending), then using a single pass or binary search over cumulative sums to find the threshold index $\Delta_K$3. Forward computation admits $\Delta_K$4 complexity, with $\Delta_K$5 possible via a quickselect strategy.

A concise summary of the evaluation procedure is as follows:

Step	Description	Complexity
Scaling & Sorting	Scale $\Delta_K$6, sort by $\Delta_K$7	$\Delta_K$8
Cumulative sum for feasibility	Precompute $\Delta_K$9	$$\mathrm{Softmax}_\tau(x)[i] = \frac{\exp(x_i/\tau)}{\sum_k \exp(x_k/\tau)}.$$0
Threshold search	Scan over $$\mathrm{Softmax}_\tau(x)[i] = \frac{\exp(x_i/\tau)}{\sum_k \exp(x_k/\tau)}.$$1; test for feasible $$\mathrm{Softmax}_\tau(x)[i] = \frac{\exp(x_i/\tau)}{\sum_k \exp(x_k/\tau)}.$$2	$$\mathrm{Softmax}_\tau(x)[i] = \frac{\exp(x_i/\tau)}{\sum_k \exp(x_k/\tau)}.$$3
Special case: UBSoftmax	Only upper bounds: can use $$\mathrm{Softmax}_\tau(x)[i] = \frac{\exp(x_i/\tau)}{\sum_k \exp(x_k/\tau)}.$$4 time	$$\mathrm{Softmax}_\tau(x)[i] = \frac{\exp(x_i/\tau)}{\sum_k \exp(x_k/\tau)}.$$5

This direct approach ensures tight box constraints are satisfied exactly and efficiently (Atarashi et al., 12 Jun 2025).

3. Gradient Structure and Differentiability

Away from the active set boundaries (i.e., for coordinates not saturated at box limits), BCSoftmax is differentiable. For $$\mathrm{Softmax}_\tau(x)[i] = \frac{\exp(x_i/\tau)}{\sum_k \exp(x_k/\tau)}.$$6 and solution $$\mathrm{Softmax}_\tau(x)[i] = \frac{\exp(x_i/\tau)}{\sum_k \exp(x_k/\tau)}.$$7, with indicator functions $$\mathrm{Softmax}_\tau(x)[i] = \frac{\exp(x_i/\tau)}{\sum_k \exp(x_k/\tau)}.$$8 and $$\mathrm{Softmax}_\tau(x)[i] = \frac{\exp(x_i/\tau)}{\sum_k \exp(x_k/\tau)}.$$9, the Jacobian with respect to $$\mathrm{Softmax}_\tau(x) = \arg\max_{y\in\Delta_K} \left\{ x^\top y - \tau \sum_{k} y_k \log y_k \right\}.$$0 is

$$\mathrm{Softmax}_\tau(x) = \arg\max_{y\in\Delta_K} \left\{ x^\top y - \tau \sum_{k} y_k \log y_k \right\}.$$1

where $$\mathrm{Softmax}_\tau(x) = \arg\max_{y\in\Delta_K} \left\{ x^\top y - \tau \sum_{k} y_k \log y_k \right\}.$$2 and $$\mathrm{Softmax}_\tau(x) = \arg\max_{y\in\Delta_K} \left\{ x^\top y - \tau \sum_{k} y_k \log y_k \right\}.$$3.

Similar low-rank plus diagonal characterizations exist for derivatives with respect to bounds $$\mathrm{Softmax}_\tau(x) = \arg\max_{y\in\Delta_K} \left\{ x^\top y - \tau \sum_{k} y_k \log y_k \right\}.$$4 and $$\mathrm{Softmax}_\tau(x) = \arg\max_{y\in\Delta_K} \left\{ x^\top y - \tau \sum_{k} y_k \log y_k \right\}.$$5, enabling both forward and backward passes to be efficiently realized in modern tensor frameworks. The O($$\mathrm{Softmax}_\tau(x) = \arg\max_{y\in\Delta_K} \left\{ x^\top y - \tau \sum_{k} y_k \log y_k \right\}.$$6) complexity applies to both computational and differential procedures (Atarashi et al., 12 Jun 2025).

4. Post-hoc Calibration via BCSoftmax

BCSoftmax underpins two principled post-hoc calibration methodologies:

1. Probability-Bounding (PB): Scalar functions $$\mathrm{Softmax}_\tau(x) = \arg\max_{y\in\Delta_K} \left\{ x^\top y - \tau \sum_{k} y_k \log y_k \right\}.$$7 and $$\mathrm{Softmax}_\tau(x) = \arg\max_{y\in\Delta_K} \left\{ x^\top y - \tau \sum_{k} y_k \log y_k \right\}.$$8, modeled by small neural networks or linear layers, induce the calibrated predictor

$$\mathrm{Softmax}_\tau(x) = \arg\max_{y\in\Delta_K} \left\{ x^\top y - \tau \sum_{k} y_k \log y_k \right\}.$$9

where $a = (a_1,\ldots,a_K)$0, $a = (a_1,\ldots,a_K)$1 are learned via

$a = (a_1,\ldots,a_K)$2

with parameters optimized over validation cross-entropy loss.

2. Logit-Bounding (LB): Leveraging the KKT structure, there exist $a = (a_1,\ldots,a_K)$3 such that

$a = (a_1,\ldots,a_K)$4

with $a = (a_1,\ldots,a_K)$5, $a = (a_1,\ldots,a_K)$6 parameterized and learned by analogous transforms. This reduces to applying softmax to logits after element-wise clipping.

Both approaches can reduce underconfidence and overconfidence, with top-1 accuracy preserved if $a = (a_1,\ldots,a_K)$7 (Atarashi et al., 12 Jun 2025).

5. Empirical Evaluation and Benchmarks

BCSoftmax-based calibration techniques, PB and LB, were evaluated across three datasets: TinyImageNet ($a = (a_1,\ldots,a_K)$8; train/val/test 90K/10K/10K), CIFAR-100 ($a = (a_1,\ldots,a_K)$9; 45K/5K/10K), and 20NewsGroups ($b = (b_1,\ldots,b_K)$0; 10.2K/1.1K/7.5K). Baseline models were ResNet-50 (TinyImageNet), DenseNet-12 (CIFAR-100), and GPCNN (20NewsGroups).

Calibration was quantified using empirical Expected Calibration Error (ECE) with $b = (b_1,\ldots,b_K)$1 bins. Compared to temperature scaling (TS), instance-based TS (IBTS), and Dirichlet calibration, BCSoftmax-based PB and LB achieved consistently lower ECE with negligible or no compromise in top-1 accuracy. For example:

Dataset	Method	ECE (↓)	Top-1 Accuracy (Δ)
TinyImageNet	TS	0.0162	—
TinyImageNet	PB-L	0.0139	Negligible
CIFAR-100	TS	0.0148	—
CIFAR-100	LB-C	0.0098	Negligible

Ablation studies confirmed that learning both upper and lower bounds, together with the temperature parameter, is essential for optimal calibration. Upper-only or lower-only constraints alone are suboptimal (Atarashi et al., 12 Jun 2025).

6. Implementation and Application Contexts

Implementing PB or LB methods involves appending a small head network to the pretrained classifier's penultimate layer, to parameterize functions generating the constraints. All parameters, including temperature, are optimized on a validation set, with constraint feasibility enforced via $b = (b_1,\ldots,b_K)$2 or softplus transformations.

BCSoftmax offers particular utility in applications demanding strict reliability, such as fairness-aware classification (to enforce equalized treatment), safety-critical systems (to cap overconfidence), and any downstream pipeline necessitating strictly bounded posterior probabilities. Its principled, exact, and computationally efficient enforcement of box constraints and compatibility with standard deep learning frameworks underlie its practical appeal for modern calibrated decision making (Atarashi et al., 12 Jun 2025).