Training Non-Differentiable Networks via Optimal Transport

An Le

arXiv · PDF · Open in the Atlas

Code that ran

We lifted 34 functions out of this paper's own repositories and ran 26 of them in a sandbox. "Ran" means the function executed on a synthesized input and returned a value. It is not a reproduction of the paper's results.

Repositories linked to this paper

• anindex/polystep

Some links come from the archived Papers with Code dataset (CC BY-SA 4.0): attribution and licence.

Abstract

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 do not rescue it: 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 across it, so no vanishing radius defines a derivative. Fixed-resolution stationarity survives. PolyStep attains it from forward passes alone: it ranks probe points on a randomly rotated polytope and steps along a softmax-weighted average of their directions, the λ = 0 endpoint of a KLpenalized transport program. A second impossibility fixes that frame: on a plateau wider than the probes reach, every rule that reweights a single cost row freezes, softmax included, and only the transport 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 avoid the discontinuity set; the softmax step the runs use matches it up to a measured remainder. Over T steps the iterates reach subspace stationarity at rate O(T -(1/2-γ) ) up to a bias floor with an interior optimal probe radius, and where the loss is Lipschitz this upgrades to Goldstein stationarity at the sharp radius √ P δ out . Eight hypotheses carry the rate, with per-experiment coverage stated. At matched optimizer steps PolyStep leads every comparison in a six-architecture by six-baseline sweep on validation accuracy. At matched evaluations it reaches 93.0 ± 0.2% on spiking networks with hard leaky-integrate-and-fire thresholds against 79.6±5.2% for the best tuned gradientfree baseline, and satisfies 92.6% of clauses on million-variable MAX-SAT, against an 87.5% random floor that OpenAI-ES decays to within 0.3 points of. Matching total evaluations instead reverses argmax attention and one MAX-SAT size, and levels hard MoE. Two limits: a step costs about one forward pass per search-subspace dimension, leaving PolyStep at chance at 4.2M parameters from scratch, and where gradients exist Adam is faster and more accurate. Code: https://github.com/anindex/polystep.

For agents

The same record, over MCP at https://syntology.ai/mcp:

get_harvested_code_for_paper("2605.01928")
get_code_for_paper("2605.01928")
have("2605.01928")

Connect an agent — have() is free.