Training Non-Differentiable Networks via Optimal Transport

原文アブストラクト

We optimize losses that jump: spiking thresholds, quantized layers, and discrete routing put jumps in the forward pass, where backpropagation does not apply. Finite differences fail: at a derivative-estimating radius, 99.5% of probe pairs on a quantized network leave the loss bit-identical, against 1.6% on a smooth control. At a jump, Clarke and conservative stationarity are undefined, and any radius-$h$ smoothing steepens like $1/h$, so no vanishing radius defines a derivative. Fixed-resolution stationarity survives. PolyStep attains it from forward passes alone, ranking probe points on a rotated polytope and stepping along a softmax-weighted average of their directions, the $\lambda=0$ endpoint of a KL-penalized transport program. A second impossibility fixes the frame: on a plateau wider than the probes reach, every rule that reweights one cost row freezes, softmax included, and only the column constraint escapes. In expectation, the cost-weighted average of the probed directions is exactly a gradient step on a smoothed loss, for every bounded measurable objective, with no probe asked to clear the jump set. Over $T$ steps it reaches subspace stationarity at rate $O(T^{-(1/2-\gamma)})$ up to a bias floor with an interior optimal probe radius, upgrading to Goldstein stationarity where the loss is Lipschitz. At matched optimizer steps, PolyStep leads all 36 comparisons of six architectures against six gradient-free baselines. It reaches 93.0% on hard LIF spiking networks against 79.6% for the best tuned baseline, and satisfies 92.6% of clauses on million-variable MAX-SAT against an 87.5% random floor. Matching total evaluations instead reverses argmax attention and one MAX-SAT size, and levels hard MoE. A step costs one forward pass per subspace dimension, leaving PolyStep at chance at 4.2M parameters from scratch; where gradients exist Adam is faster and more accurate.