Single-loop stationarity-constrained hypergradient

(1) Integration point: implement this as a bilevel optimizer for differentiable hyperparameter optimization or meta-learning. Let \(x\) denote hyperparameters or outer-network parameters, \(y\) denote inner-network weights, \(f(x,y)\) denote validation loss, and \(g(x,y)\) denote training loss. Replace several inner SGD steps plus unrolling or implicit differentiation with one simultaneous update of \(x\), \(y\), and \(\lambda\) per minibatch.

(2) Pseudocode:

initialize x, y, lambda
choose lr_x, lr_y, lr_lam, mu, R, h
for each iteration:
    sample training batch Btr and validation batch Bva
    gy = grad_y g(x, y; Btr)
    gx_f = grad_x f(x, y; Bva)
    gy_f = grad_y f(x, y; Bva)
    gy_xplus = grad_y g(x + h*lambda, y; Btr)
    gy_yplus = grad_y g(x, y + h*lambda; Btr)
    vx = gx_f + (gy_xplus - gy)/h
    vy = gy_f + (gy_yplus - gy)/h
    vlambda = gy - mu*lambda
    x = x - lr_x*vx
    y = y - lr_y*vy
    lambda = project_ball(lambda + lr_lam*vlambda, R)

The two finite differences are the paper's \(\widehat H_{xy}\lambda\) and \(\widehat H_{yy}\lambda\). For reduced truncation bias, test the centered estimator \([\nabla_y g(z+h v)-\nabla_y g(z-h v)]/(2h)\), using the appropriate perturbation of \(x\) or \(y\).

(3) Compute directly from the mathematics the inner stationarity residual \(\|\nabla_y g\|_2\), dual residual \(\|\nabla_y g-\mu\lambda\|_2\), multiplier norm, projection frequency, and finite-difference estimates. Estimate empirically the finite-difference step by testing \(h\in\{10^{-2},10^{-3},10^{-4}\}\) times the parameter RMS. Tune \(\mu\), \(R\), and the three learning rates independently; use separate training and validation minibatches to avoid leakage.

(4) First experiment: use a two-layer MLP on MNIST or a five-way Omniglot meta-learning split. Let \(x\) be either 10–50 scalar hyperparameters or an outer initialization and let \(y\) be the classifier weights. Compare against first-order MAML, 5–20-step truncated unrolling, and implicit differentiation at equal gradient-evaluation and wall-clock budgets. Track validation loss, validation accuracy, \(\|\nabla_y g\|_2\), multiplier clipping rate, and memory. The method is successful if it reaches lower validation loss or accuracy at equal compute, maintains a decreasing stationarity residual, and avoids exploding multipliers. Ablate dual regularization, projection, one-sided versus centered differences, and exact autodiff Hessian-vector products.