File size: 14,253 Bytes
2d1810a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 | """
Infimal Convolution with Implicit Differentiation
This module implements the infimal convolution operation g(y) = inf_x [K(x,y) - f(x)]
as a differentiable PyTorch operation using implicit differentiation via the envelope theorem.
The key insight is that gradients with respect to the network parameters f_net can be
computed without backpropagating through the optimization process that finds the minimizer x*.
Instead, we use the envelope theorem which states that at the optimum x*:
d/dθ g(y,θ) = -d/dθ f(x*,θ)
This allows efficient gradient computation while maintaining correctness.
Technical Details:
- Forward pass: Solves the optimization problem to find x* = argmin_x [K(x,y) - f(x)]
- Backward pass: Applies envelope theorem to compute gradients wrt network parameters
- No gradient tracking through the optimization loop (uses functional_call with detached params)
- Supports multiple optimizers: LBFGS, Adam, and GD
- Early stopping based on relative change in objective value
- Optional L2 regularization on the optimization variable x
References:
- Envelope Theorem: https://en.wikipedia.org/wiki/Envelope_theorem
- Implicit Differentiation in Optimization: Amos & Kolter (2017), OptNet
Example:
>>> import torch
>>> import torch.nn as nn
>>> from tools.inf_convolution import InfConvolution
>>>
>>> # Define network f(x) = a^T x + b
>>> class LinearF(nn.Module):
... def __init__(self, d):
... super().__init__()
... self.a = nn.Parameter(torch.randn(d))
... self.b = nn.Parameter(torch.randn(()))
... def forward(self, x):
... return x @ self.a + self.b
>>>
>>> # Define kernel K(x,y) = 0.5||x-y||^2
>>> def K2(x, y):
... return 0.5 * ((x - y)**2).sum()
>>>
>>> # Compute infimal convolution
>>> fnet = LinearF(3)
>>> y = torch.randn(3)
>>> x0 = torch.zeros(3)
>>> g, converged = InfConvolution.apply(y, fnet, K2, x0, 100, 1.0, "lbfgs", 0.0, 1e-6, *list(fnet.parameters()))
>>> print(f"g(y) = {g.item():.4f}, converged = {converged}")
>>>
>>> # Compute gradients
>>> g.backward()
>>> print(f"Gradient wrt a: {fnet.a.grad}")
"""
import torch
from torch.autograd import Function
from torch.func import functional_call
class InfConvolution(Function):
"""
PyTorch autograd Function for computing infimal convolution with implicit differentiation.
Computes g(y) = inf_x [K(x,y) - f(x)] where:
- K(x,y) is a kernel function (typically a distance metric)
- f(x) is a neural network parameterized by θ
- The optimization is performed wrt x, not θ
The gradient computation uses the envelope theorem to avoid backpropagating through
the optimization loop, making it efficient and numerically stable.
"""
@staticmethod
def forward(ctx, y, f_net, K, x_init, solver_steps=30, lr=1e-1, optimizer="gd", lam=0.0, tol=1e-6, *params):
"""
Forward pass: Solve the optimization problem to compute g(y) = inf_x [K(x,y) - f(x)].
This method performs the following steps:
1. Initialize optimization variable x from x_init
2. Solve the minimization problem: x* = argmin_x [K(x,y) - f(x) + 0.5*lam*||x||^2]
3. Compute and return g(y) = K(x*,y) - f(x*)
The optimization is performed with manual gradient computation to avoid tracking
gradients through the network parameters during the forward pass. This is crucial
for efficiency and correctness of the implicit differentiation.
Args:
y (torch.Tensor): Input tensor of shape (..., d). The point at which to evaluate g(y).
f_net (nn.Module): Neural network representing f(x). Must be differentiable wrt x.
K (callable): Kernel function K(x, y) -> scalar tensor. Should be differentiable wrt x.
Common choice: K(x,y) = 0.5*||x-y||^2 (squared Euclidean distance).
x_init (torch.Tensor): Initial guess for the optimization variable x. Shape should
match the input dimension d.
solver_steps (int, optional): Maximum number of optimization steps. Default: 30.
Increase for harder optimization problems.
lr (float, optional): Learning rate for the optimizer. Default: 1e-1.
Typical values: 1.0 for LBFGS, 1e-2 to 1e-1 for Adam/GD.
optimizer (str, optional): Choice of optimizer. One of:
- "lbfgs": Limited-memory BFGS (recommended for smooth problems)
- "adam": Adam optimizer (good for non-smooth problems)
- "gd": Gradient descent / SGD (simplest, may need more steps)
Default: "gd"
lam (float, optional): L2 regularization weight on x. Adds 0.5*lam*||x||^2 to the
objective. Useful for improving conditioning. Default: 0.0 (no regularization).
tol (float, optional): Tolerance for early stopping. Optimization stops when
relative change in objective is less than tol. Default: 1e-6.
Formula: |obj_val - prev_obj| / (|prev_obj| + 1e-10) < tol
*params: Network parameters from f_net.parameters(). Needed for gradient tracking
in the backward pass. Automatically extracted if not provided.
Returns:
tuple: (g, converged) where:
- g (torch.Tensor): Scalar value of inf_x [K(x,y) - f(x)]
- converged (bool): Whether optimization converged before reaching solver_steps.
True if relative change fell below tol, False otherwise.
For LBFGS, always returns True (uses internal convergence criteria).
Notes:
- The forward pass does NOT accumulate gradients in f_net parameters.
- Uses torch.func.functional_call to evaluate f_net without parameter tracking.
- Manual gradient computation ensures efficiency and numerical stability.
- LBFGS typically converges faster but may need tuning for non-smooth problems.
- Adam is more robust to initialization and works well with default settings.
Raises:
RuntimeError: If optimization diverges or produces NaN/Inf values.
"""
# If no params passed, get them from f_net
if len(params) == 0:
params = tuple(f_net.parameters())
# Create detached parameter dict for functional_call so that the optimizer
# only updates `x_var` and not the network parameters.
params_dict = {name: p.detach() for name, p in f_net.named_parameters()}
x_var = x_init.clone().detach().requires_grad_(True)
# Solve argmin_x [K(x,y) - f(x) + 0.5*lam*||x||^2]
opt_name = optimizer.lower() if isinstance(optimizer, str) else "gd"
converged = False
with torch.enable_grad():
if opt_name == "lbfgs":
# LBFGS closure - compute objective and gradients
def closure():
if x_var.grad is not None:
x_var.grad.zero_()
# Compute objective using detached parameters (no param gradients)
k_val = K(x_var, y)
f_val = functional_call(f_net, params_dict, (x_var,))
obj = k_val - f_val + 0.5 * lam * x_var.pow(2).sum()
# Compute gradients wrt x only (params are detached)
obj.backward()
return obj
optim_obj = torch.optim.LBFGS(
[x_var],
lr=lr,
max_iter=solver_steps,
tolerance_grad=tol,
tolerance_change=tol,
line_search_fn="strong_wolfe"
)
optim_obj.step(closure)
# LBFGS handles its own convergence internally
converged = True # Assume converged if LBFGS completes
else:
if opt_name == "adam":
optim_obj = torch.optim.Adam([x_var], lr=lr)
else: # gd or default
optim_obj = torch.optim.SGD([x_var], lr=lr)
prev_obj = None
for step in range(solver_steps):
optim_obj.zero_grad()
# Compute objective using detached parameters (no param gradients)
k_val = K(x_var, y)
f_val = functional_call(f_net, params_dict, (x_var,))
obj = k_val - f_val + 0.5 * lam * x_var.pow(2).sum()
# Compute gradients wrt x only (params are detached)
obj.backward()
optim_obj.step()
# Check convergence based on relative change
obj_val = obj.item()
if prev_obj is not None:
rel_change = abs(obj_val - prev_obj) / (abs(prev_obj) + 1e-10)
if rel_change < tol:
converged = True
break
prev_obj = obj_val
# Mark as not converged if we completed all steps
if step == solver_steps - 1:
converged = False # Did not converge early
x_star = x_var.detach()
ctx.f_net = f_net
ctx.K = K
ctx.y = y
ctx.num_params = len(params)
ctx.save_for_backward(x_star, *params)
# Compute g value (no gradients needed for return value in forward)
with torch.no_grad():
g = K(x_star, y) - f_net(x_star)
return g, converged
@staticmethod
def backward(ctx, grad_output, grad_converged):
"""
Backward pass: Compute gradients using implicit differentiation via envelope theorem.
The envelope theorem states that for g(y,θ) = inf_x [K(x,y) - f(x,θ)], the gradient
with respect to parameters θ is:
dg/dθ = -df(x*,θ)/dθ
where x* = argmin_x [K(x,y) - f(x,θ)]. Notably, we do NOT need to compute dx*/dθ,
which would require backpropagating through the optimization loop. This makes the
computation efficient and numerically stable.
Mathematical Justification:
Let L(x,y,θ) = K(x,y) - f(x,θ). By definition:
g(y,θ) = min_x L(x,y,θ)
At the optimum x*, the first-order condition holds:
∂L/∂x|_{x=x*} = 0
Taking the total derivative:
dg/dθ = ∂L/∂θ|_{x=x*} + (∂L/∂x|_{x=x*}) · (dx*/dθ)
= ∂L/∂θ|_{x=x*} [since ∂L/∂x|_{x=x*} = 0]
= -∂f(x*,θ)/∂θ
Args:
ctx: Context object containing saved tensors and attributes from forward pass.
Stores: x_star (optimal x), params (network parameters), f_net, K, y.
grad_output (torch.Tensor): Gradient of loss wrt g (scalar). This is typically 1.0
when g is the final loss, or the upstream gradient in a larger computation.
grad_converged: Gradient wrt converged flag (always None, since it's a boolean).
Returns:
tuple: Gradients wrt all forward pass inputs, in the same order:
(grad_y, grad_f_net, grad_K, grad_x_init, grad_solver_steps, grad_lr,
grad_optimizer, grad_lam, grad_tol, *grad_params)
- grad_y: None (we don't compute gradients wrt input y)
- grad_f_net: None (not a tensor)
- grad_K: None (not a tensor)
- grad_x_init: None (initialization doesn't affect final output via optimization)
- grad_solver_steps: None (hyperparameter)
- grad_lr: None (hyperparameter)
- grad_optimizer: None (string)
- grad_lam: None (hyperparameter)
- grad_tol: None (hyperparameter)
- *grad_params: Tuple of gradients wrt each network parameter θ, computed as
-grad_output * ∂f(x*,θ)/∂θ
Implementation Details:
- Uses saved x_star from forward pass (no need to re-solve optimization)
- Computes gradients wrt parameters by evaluating network at x_star
- Multiplies by -grad_output (negative sign from envelope theorem)
- Uses saved parameters to ensure consistency with forward pass
- allow_unused=True handles networks with unused parameters gracefully
"""
f_net = ctx.f_net
saved_tensors = ctx.saved_tensors
x_star = saved_tensors[0]
params = saved_tensors[1:]
# Compute gradient wrt f_net parameters using envelope theorem
# d/dθ g(y,θ) = -d/dθ f(x*,θ) evaluated at the optimal x*
# grad_converged is None since converged is a boolean flag
with torch.enable_grad():
x_star_grad = x_star.detach().requires_grad_(True)
# Compute f(x*) using the SAVED parameters, not f_net.parameters()
# This ensures consistency with the forward pass computation
grad_theta = torch.autograd.grad(
outputs=f_net(x_star_grad),
inputs=params,
grad_outputs=-grad_output, # Negative sign from envelope theorem
retain_graph=False,
allow_unused=True,
)
# Return gradients in the same order as forward inputs
# (grad_y, grad_f_net, grad_K, grad_x_init, grad_solver_steps, grad_lr,
# grad_optimizer, grad_lam, grad_tol, *grad_params)
return (None, None, None, None, None, None, None, None, None) + grad_theta
|