Kink-Flux Regularization for Max Routers
Implementation & benchmark of arXiv:2607.09987 — Reflected Optimal Stopping with a Max-Type Payoff: Measure-Valued Stopping Gains and Killed Resolvent Representation
Source paper: Reflected Optimal Stopping with a Max-Type Payoff: Measure-Valued Stopping Gains and Killed Resolvent Representation arXiv:2607.09987 ⓘ · analyzed Aug 30, 2026
AI-generated research hypothesis, automatically tested. Not peer-reviewed.
Idea description
Use the paper's singular stopping-gain term to explicitly measure how much learned feature covariance crosses a max or routing boundary. Penalize excessive covariance in the normal direction to the switching surface, rather than pretending that the max operation has an ordinary Hessian.
Formulas
Mathematical statement
For the payoff G(x_1,x_2)=x_1\vee \alpha x_2, the switching surface is \Delta=\{x_1=\alpha x_2\}, with normal vector n=(1,-\alpha). If a=\sigma\sigma^ op is the local diffusion covariance, the paper identifies the singular diagonal contribution to the stopping gain as \Gamma^{\Delta}(dx)=-\frac{n^ op a(x)n}{2\sqrt{1+\alpha^2}}\sigma_{\Delta}(dx), where \sigma_{\Delta} is one-dimensional surface measure on \Delta. The quantity n^ op a n is the variance transported across the kink. For a neural max router with two logits z=(z_1,z_2), use u=z_1-\alpha z_2=n^ op z and estimate the surface delta by \delta_\varepsilon(u)=\exp(-u^2/(2\varepsilon^2))/(\sqrt{2\pi}\varepsilon). The adapted regularizer is R_{\mathrm{kink}}=\mathbb{E}[\delta_\varepsilon(u)n^ op A(z)n/(2\sqrt{1+\alpha^2})], where A(z) is a local covariance estimate of router-logit perturbations. This approximates the total singular kink flux without requiring a Hessian of the nonsmooth max.
Implementation notes
(1) Integrate this into the auxiliary loss of a maxout layer, hard two-expert MoE router, or max-pooling block. Let z_b be the two competing logits for batch item b, and let A_b estimate the covariance of their local perturbations. A cheap choice is A_b=J_bJ_b^ op, where J_b is the 2-by-d Jacobian of z_b with respect to the preceding representation and the input covariance is the identity; alternatively estimate A_b from two or four dropout or augmentation perturbations of the same example. (2) Pseudocode: compute z=model.router(h); u=z[:,0]-alpha*z[:,1]; n=[1,-alpha]; estimate A for each item; delta=exp(-u*u/(2*eps*eps))/(sqrt(2*pi)*eps); kink=(delta*(n^T A n)/(2*sqrt(1+alpha^2))).mean(); total_loss=task_loss+lambda_kink*kink; backpropagate total_loss. Detach A initially if Jacobian computation is expensive, then enable full gradients after warmup. Anneal eps from 0.2 times the batch standard deviation of u to 0.02 times that standard deviation. (3) The geometric coefficient and Gaussian surface-delta approximation come from the paper; A, epsilon, and lambda are empirical design choices. Apply the regularizer only near the switching surface through delta, not globally to all router variance. (4) First experiment: train a two-expert MLP and a maxout MLP on CIFAR-10 or a synthetic mixture-of-Gaussians classification task, comparing the baseline hard router against lambda_kink values in {0,1e-4,1e-3,1e-2}. Track validation accuracy, router flip rate under small input perturbations, gradient norm spikes, and calibration. Success means lower flip rate and fewer gradient spikes at equal accuracy, or improved accuracy under noise perturbations with less than 5% extra training FLOPs.
Verification
This idea has not been verified yet.
Verification happens in two stages: Stage 1 — a mechanism check on a toy system confirms the claimed mathematical phenomenon reproduces; Stage 2 — a benchmark implements the idea on a real (small) neural network task and compares it against a tuned baseline over 8 paired seeds with a permutation test.
Artifacts
Artifacts unavailable.