Skip to content

Solution-Map Differentiation

SolMapLayer is used when a scalar loss depends on the optimal decisions

\[ \hat z^\star=\mathcal O(\hat y). \]

DiffAPQP differentiates the KKT optimality conditions at the forward solution. It does not backpropagate through the iterations of the selected optimization solver.

Chain rule through the layer

For \(\hat y=f(x,\theta)\) and a scalar loss \(\ell(\hat z^\star)\),

\[ \mathrm D_\theta\ell = \mathrm D_{\hat z^\star}\ell\, \mathrm D_{\hat y}\mathcal O(\hat y)\, \mathrm D_\theta f(x,\theta). \]

The forecaster and loss derivatives are handled by PyTorch. DiffAPQP computes the vector-Jacobian product through \(\mathcal O\) with an adjoint KKT solve, avoiding construction of the full Jacobian \(\mathrm D_{\hat y}\mathcal O(\hat y)\).

Retrieval and the upstream gradient

The public decisions are recovered from the canonical solution by

\[ \hat z^\star=\mathcal R(\tilde z^\star). \]

The retrieval map \(\mathcal R\) is a linear block selection/permutation operation recorded during canonicalization. If \(g_{\hat z}=\nabla_{\hat z^\star}\ell\) is the gradient received from PyTorch, the canonical primal gradient is

\[ g_{\tilde z} = \mathrm D\mathcal R(\tilde z^\star)^T g_{\hat z}. \]

This vector becomes the primal block of the adjoint right-hand side.

KKT conditions

For the canonical APQP

\[ \begin{aligned} \min_{\tilde z}\quad &\frac{1}{2}\tilde z^T P\tilde z +q(\hat y)^T\tilde z+c(\hat y)\\ \text{subject to}\quad &A\tilde z=b(\hat y),\\ &G\tilde z\le h(\hat y), \end{aligned} \]

the Lagrangian is

\[ \mathcal L = \frac{1}{2}\tilde z^T P\tilde z +q(\hat y)^T\tilde z+c(\hat y) +\tilde\nu^T(A\tilde z-b(\hat y)) +\tilde\lambda^T(G\tilde z-h(\hat y)). \]

At an optimal primal-dual point \((\tilde z^\star,\tilde\nu^\star,\tilde\lambda^\star)\), the KKT conditions are

\[ \begin{aligned} P\tilde z^\star+q(\hat y) +A^T\tilde\nu^\star +G^T\tilde\lambda^\star &=0,\\ A\tilde z^\star-b(\hat y)&=0,\\ \operatorname{diag}(\tilde\lambda^\star) \big(G\tilde z^\star-h(\hat y)\big)&=0,\\ G\tilde z^\star-h(\hat y)&\le0,\\ \tilde\lambda^\star&\ge0. \end{aligned} \]

DiffAPQP offers two adjoint formulations derived from these conditions.

Full-KKT differentiation

Differentiating stationarity, equality feasibility, and complementarity with respect to parameter block \(\hat y_i\) gives

\[ K_a \mathrm D_{\hat y_i} \begin{pmatrix} \tilde z^\star\\ \tilde\nu^\star\\ \tilde\lambda^\star \end{pmatrix} = D_{a,i}, \]

with

\[ K_a = \begin{pmatrix} P & A^T & G^T\\ A & 0 & 0\\ \operatorname{diag}(\tilde\lambda^\star)G &0 &\operatorname{diag}\big(G\tilde z^\star-h(\hat y)\big) \end{pmatrix}, \]

and

\[ D_{a,i} = \begin{pmatrix} -Q_i\\ B_i\\ \operatorname{diag}(\tilde\lambda^\star)H_i \end{pmatrix}. \]

The signs follow from the affine mappings \(q(\hat y)\), \(b(\hat y)\), and \(h(\hat y)\) defined on the APQP page.

Full-KKT adjoint

Instead of solving one sensitivity system for every parameter component, reverse mode solves one adjoint system. Let

\[ g_u = \begin{pmatrix} g_{\tilde z}\\ 0\\ 0 \end{pmatrix}. \]

The full-KKT adjoint is obtained from

\[ K_a^T s_a=g_u. \]

When an iterative least-squares solve is used, this is interpreted as

\[ s_a^\star = \arg\min_{s_a} \left\|K_a^T s_a-g_u\right\|_2^2. \]

The gradient with respect to parameter block \(\hat y_i\) is then

\[ \nabla_{\hat y_i}\ell = D_{a,i}^T s_a^\star. \]

The full matrix has dimension

\[ n_d+n_{\mathrm{eq}}+n_{\mathrm{in}}. \]

It is generally nonsymmetric because of the complementarity block and can be ill-conditioned when inequality slacks or duals are close to zero.

Reduced-KKT differentiation

The reduced formulation uses only inequalities active at the forward solution. For an exact solution, define

\[ \mathcal A = \left\{ j: G_j\tilde z^\star=h_j(\hat y) \right\}. \]

Locally, an inactive inequality has zero dual sensitivity contribution. After retaining the equality constraints and active inequality rows, the QP becomes

\[ \begin{aligned} \min_{\tilde z}\quad &\frac{1}{2}\tilde z^T P\tilde z +q(\hat y)^T\tilde z+c(\hat y)\\ \text{subject to}\quad &A\tilde z=b(\hat y),\\ &G_{\mathcal A}\tilde z=h_{\mathcal A}(\hat y). \end{aligned} \]

Its differentiated KKT system is

\[ K_s \mathrm D_{\hat y_i} \begin{pmatrix} \tilde z^\star\\ \tilde\nu^\star\\ \tilde\lambda_{\mathcal A}^\star \end{pmatrix} = D_{s,i}, \]

where

\[ K_s = \begin{pmatrix} P&A^T&G_{\mathcal A}^T\\ A&0&0\\ G_{\mathcal A}&0&0 \end{pmatrix}, \qquad D_{s,i} = \begin{pmatrix} -Q_i\\ B_i\\ H_{i,\mathcal A} \end{pmatrix}. \]

The reduced matrix is symmetric and has dimension

\[ n_d+n_{\mathrm{eq}}+|\mathcal A|. \]

The adjoint therefore satisfies

\[ K_s s_s=g_s, \qquad g_s= \begin{pmatrix} g_{\tilde z}\\ 0\\ 0 \end{pmatrix}, \]

and

\[ \nabla_{\hat y_i}\ell = D_{s,i}^T s_s. \]

When the local active set is correctly identified and the relevant regularity conditions hold, the reduced system represents the same local solution sensitivity as the full KKT conditions while removing inactive complementarity rows. Its benefit grows when \(|\mathcal A|\ll n_{\mathrm{in}}\).

Numerical active-set detection

Numerical solutions rarely satisfy an active constraint with exact equality. The implementation marks inequality \(j\) active when

\[ \left| (G\tilde z^\star)_j-h_j(\hat y) \right| \le \max\left( \epsilon_{\mathrm{abs}}, \epsilon_{\mathrm{rel}} \big(1+|h_j(\hat y)|\big) \right). \]

The absolute and relative thresholds can be passed as act_ineq_tol and act_ineq_tol_rel. When either is omitted, DiffAPQP infers it from matching forward-solver tolerances when available, with a lower bound of \(10^{-4}\). Exact precedence and recognized option names are documented under SolMapLayer adjoint configuration.

A tolerance that is too small can omit genuinely binding constraints because of forward-solver residuals. A tolerance that is too large can include nearly active or dependent rows and make \(K_s\) singular or poorly conditioned. See Numerical considerations.

Linear-system modes

The mathematical formulation and numerical linear solver are separate choices:

Formulation Default mode Rationale
Full KKT (reduced_kkt=False) lsqr Handles the nonsymmetric adjoint as a least-squares problem
Reduced KKT (reduced_kkt=True) minres Exploits the symmetric reduced matrix

spsolve is available as a direct sparse solve for either square system. lsqr is permitted for the reduced system but normally gives up its symmetric structure. minres is not accepted for the nonsymmetric full KKT system. Argument defaults and override rules are listed in the SolMapLayer reference.

Regularity and interpretation

The KKT derivative is a local sensitivity. A unique optimizer and stable gradient generally require suitable regularity, such as an isolated primal solution and linearly independent active constraints. At degenerate points:

  • \(K_a\) or \(K_s\) may be singular;
  • multiple optimal solutions may produce different valid sensitivities;
  • small changes in tolerance can change \(\mathcal A\);
  • QP/LP solutions can agree while their reported gradients differ.

These are properties of the optimization map, not merely implementation failures. LPs are particularly exposed because \(P=0\). The numerical page discusses regularization, solver tolerances, and validation.

Interface consequence

Only the named decisions in primal_dict are connected to autograd. Canonical dual arrays are returned for inspection, but losses built from those returned duals are not differentiated. See the complete SolMapLayer contract.