OptimizationAdjoint for constrained optimization solution sensitivities

[jClugstor]

Checklist

• Appropriate tests were added
• Any code changes were done in a way that does not break public API
• All documentation related to code changes were updated
• The new code follows the
  contributor guidelines, in particular the SciML Style Guide and
  COLPRAC.
• Any new documentation only uses public API

Additional context

Add any other context about the problem here.

[jClugstor]

Summary

Adds a new OptimizationAdjoint sensitivity algorithm that computes $dG/dp$ for a downstream loss G evaluated at the optimum $x^*(p)$ of a (possibly constrained) OptimizationProblem, without re-solving the optimization. Handles equality constraints, two-sided inequality constraints (lcons/ucons), and variable box bounds (lb/ub). Sits alongside the existing UnconstrainedOptimizationAdjoint.

opt_sol = solve(prob, NLopt.LD_SLSQP())                # solve once
dgdu!(out, u, _, _, _) = (out .= dG_dx)                # cotangent of loss w.r.t. x*
dp = adjoint_sensitivities(opt_sol, nothing;
                           sensealg = OptimizationAdjoint(),
                           dgdu = dgdu!)

[jClugstor]

The Math

A simple way to say this is that this is basically the same thing as SteadyStateAdjoint, but just doing the implicit differentiation on the Nonlinear system that is made up of the gradient of the lagrangian

$$ \mathcal{L}(x,y,z,p) = f(x,p) + y^\top g(x,p) + z^\top h(x,p) $$

and the constraint equations all set to zero.

First of all though, we do need all of the KKT variables/ dual variables / lagrange multipliers, the variables that are multiplied by the constraint functions in the lagrangian that satisfy the KKT conditions. We get the optimal $x^*$ from the optimization forward pass, but not every solver will return the dual variables, so we need to recover them from the stationarity condition:

$$ [J_g; J_{h_I}]^\top \mu = -\nabla f $$

This is typically overdetermined (more x's than active constraints), so we solve via LinearSolve.QRFactorization — least-squares, which under LICQ recovers the exact multipliers. LICQ says that at the optimal point, the gradients of the active constraints are all linearly independent.

QR decomposition finds the least squares solution, so it's finding

$$ \min ||J^\top \mu^* - (-\nabla f)||^2 $$

but if we found a viable local minimum, the KKT theorem states that $\mu$ exists , so the residual there will be zero. Plus, LICQ means that the Jacobian has full column rank, so the $\mu$ is unique. So a QR solve finds the unique solution.

If any returned z_I is negative (KKT violation), the corresponding constraint was wrongly classified as active; drop it and redo the multiplier solve. The sign-check + redo handles boundary degeneracy that proximity-based active-set detection alone gets wrong.

Now for the sensitivity, the idea is that a solution of the constrained optimization problem satisfies the Karush-Kuhn Tucker (KKT) conditions:

$$\nabla_x f(x,p) + \sum_{i=1}^{m_e} y_i \nabla_x g_i(x,p) + \sum_{i=1}^{m_i}z_i \nabla_x h_i(x,p) = 0$$

$$g(x,p) = 0, \quad h(x,p) \le 0, \quad z \ge 0, \quad z^\top h(x,p) = 0$$

Taking only the set of inequality constraints that are actually active, $h_I$, the KKT conditions define a nonlinear system of equations:

$$F(w, p) = \begin{bmatrix} \nabla \mathcal{L}(w, p) \ g(w,p) \ h_\mathcal{I}(w,p) \end{bmatrix} = 0$$ where $w$ is $$w := (x, y, z_\mathcal{I})$$

The solution of the optimization problem satisfies this nonlinear system of equations.
Differentiate by $p$, using implicit differentiation:

$$ \frac{\partial F}{\partial w}\frac{\partial w}{\partial{p}} + \frac{\partial F}{\partial p } = 0 $$

Our end goal is to find $\frac{\partial w}{\partial p}$ , which is the sensitivity of the solution of the optimization problem $w$, with respect to the parameters.

Using the above,

$$\frac{\partial w}{\partial p} = -\frac{\partial F}{\partial w}^{-1}\frac{\partial F}{\partial p}$$

Now consider that we have some scalar cost function that will be a function of the solution of the optimization problem, $C(w)$.

$$\frac{dC}{dp} = \frac{\partial C}{\partial w}\frac{\partial w}{\partial p } + \frac{\partial C}{\partial p} = \frac{\partial C}{\partial w}\left(- \frac{\partial F}{\partial w}^{-1} \frac{\partial F}{\partial p} \right) + \frac{\partial C}{\partial p}$$

$$\frac{dC}{dp} = \frac{\partial C}{\partial p} - \left(\frac{\partial C}{\partial w} \frac{\partial F}{\partial w}^{-1} \right)\frac{\partial F}{\partial p} $$

But this term $\frac{\partial C}{\partial w} \frac{\partial F}{\partial w}^{-1}$ can be represented as a linear solve,

$$\frac{\partial F}{\partial w} \lambda = \frac{\partial C}{\partial w}$$

so

$$\frac{dC}{dp} = \frac{\partial C}{\partial p} - \lambda \frac{\partial F}{\partial p}$$

that $\lambda \frac{\partial F}{\partial p}$ is what vecjacobian! does.

[jClugstor]

@ChrisRackauckas this is finally ready to be looked at
@SebastianM-C I'm pinging you too since this concerns Optimization.jl

[ChrisRackauckas]

All forward?

[jClugstor]

What do you mean? The finite difference grad should be using FDT, right?

[ChrisRackauckas]

There's no reverse mode gradient methods here?

[jClugstor]

Oh right, for context this only gets used for the gradient of the lagrangian. The actual VJP is performed by vecjacobian!.

If we want to be able to use reverse mode for the lagrangian grad I think the best way to do it would be to just use DifferentiationInterface and just add a field in the OptimizationAdjoint for the ADType. We only need to differentiate with respect to one variable so it should work fine.

That probably won't work for all combinations of AD backend for lagrangian gradient and autojacvec of course because it will be AD of AD, but should work for some combinations I think.

[ChrisRackauckas]

Why not just require that the optimization has be initialized with a gradient function? Optimization.jl has a whole system for that already, and you just get an error that you need to do the standard Optimization AD setup if you want to do this.

[jClugstor]

We could have Optimization.jl build the gradient of the lagrangian, lag_g? The problem I can see with that is that we have to differentiate lag_g with respect to p. So because of DI preparation it might be difficult to figure out how to push Duals through it, in the case of autojacvec = true.

[jClugstor]

@ChrisRackauckas the problem with using the grad created for the OptimizationFunction is that it uses a DI preparation, which bakes in the types. So there's no way to put Duals in to it for ForwardDiff autojacvec without using some kind of diffcache, or having the grad reprepare every time.

[ChrisRackauckas]

It's probably better just to require there is a hess and grad function from the standard system, rather than doubling it up.

[jClugstor]

Oh I see, those fields get populated anyways depending on the solver. Makes sense.

[jClugstor]

First order constrained solvers don't populate the required second order pieces that we need, so for those you would always have to provide either lag_h or hess and cons_h. Is that ok?

[SebastianM-C]

Can we add that via DI? What solvers do you need to support this?

[jClugstor]

The main issue here is that in instantiate_function, it looks at which pieces the algorithm requires, and if it doesn't require that piece, it just doesn't build it, or store it in the cache, even if the user directly provided it.
e.g. here https://github.com/SciML/Optimization.jl/blob/ae548dc4213083be50afc9e35dd7e06eab3e0c12/lib/OptimizationBase/src/OptimizationDIExt.jl#L99

if the algorithm doesn't require the hessian even a user provided hessian won't be stored in the cache, so it can't be used in the adjoint.

I think the choices are:

• Have Optimization.jl always build all of the derivatives and keep them in the OptimizationCache. This would allow sensitivity analysis for constrained problems that used single order constrained solvers even if the user didn't provide second order information.

• Just keep the second order information in the cache if it's provided by a user, even if the solver doesn't use the information. This would mean that for first order solvers users would have to provide second order information to find the sensitivities.

• Make an OptimizationBase extension in SciMLSensitivity, check if the cache has second order information, if it doesn't have it call OptimizationBase.instantiate_function to build it.

[jClugstor]

What solvers do you need to support this?

I would think as many as possible. The calculation of the adjoint doesn't use the optimization solver at all, it only needs the function, first order, and second order information to find the derivative.

[jClugstor]

Ok, so this works, but I basically have a custom gradient wrapper for the inner gradient of the lagrangian. I just take the objective function and constraints and create the lagrangian, and use AD to find the gradient, then pass that on to vecjacobian! later.

If we want to reuse the grad in the OptimizationFunction, we would need to change it so that it's able to accept dual numbers, for when vecjacobian! uses ForwardDiff (also for a part of the KKT matrix). Something like this SciML/Optimization.jl#1229 would let us pass Duals through.