Training Non-Differentiable Networks via Optimal Transport
Abstract
Hard thresholds, quantization, and discrete routing can produce training losses with flat regions and jumps, where ordinary gradients vanish or are undefined. We introduce PolyStep, a forward-only optimizer that evaluates rotated polytope probes and moves parameter blocks along weighted averages of the probe directions. We derive the weights from one-sided entropic transport and use its uncoupled softmax solution in our primary experiments. Our analysis explains when variation among probe costs produces motion and when that motion decreases the loss. On a regular simplex, nonconstant costs always give a nonzero direction. For monotone ridge losses, the softmax update cannot increase the loss at any positive temperature; a perturbation bound gives sufficient conditions for descent near curved jumps. For bounded measurable losses, we randomize the probe radii and identify an exact smoothing whose gradient equals the expected linear cost-weighted direction up to scale. This identity yields a stationarity bound for an idealized fixed-temperature variant: under regularity and sampling assumptions stronger than those met by our trained configurations, the bound has an O(T−1/2) term and a persistent bias floor. We evaluate the practical method on networks with hard operations, discrete optimization, and policy search. On MNIST with hard-threshold spiking neurons, PolyStep reaches 93.0 ± 0.2%, compared with 79.6 ± 5.2% for the best-tuned gradient-free baseline at matched evaluations. These gains come with a query cost proportional to the search dimension per fresh step, which limits the number of updates available at a fixed budget.
Related
- Statistical Inference for Adversarial Training: Central Limit Theorems via Optimal Transport
- Coarse-Graining Optimal Transport for Rollout-Free Post-Training of Flow Matching
- Brenier Meets Adversarial Training: Optimal Transport Geometry for Robust Learning
- Differentiable Hybrid Modelling for Learning and Optimising Chemical Transport Processes from Experimental Data
- RapidDOC: Riccati-Preconditioned Implicit Differentiation for Differentiable Optimal Control under Operator Mismatch
BibTeX
@misc{le2026training,
title = {{Training Non-Differentiable Networks via Optimal Transport}},
author = {An T. Le},
howpublished = {arXiv},
year = {2026},
month = oct,
doi = {10.48550/arxiv.2605.01928},
eprint = {2605.01928},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2605.01928},
note = {CC BY 4.0},
}
RIS
TY - UNPB
TI - Training Non-Differentiable Networks via Optimal Transport
AU - An T. Le
PY - 2026
DA - 2026/10/08
PB - arXiv
DO - 10.48550/arxiv.2605.01928
UR - https://arxiv.org/abs/2605.01928
AB - Hard thresholds, quantization, and discrete routing can produce training losses with flat regions and jumps, where ordinary gradients vanish or are undefined. We introduce PolyStep, a forward-only optimizer that evaluates rotated polytope probes and moves parameter blocks along weighted averages of the probe directions. We derive the weights from one-sided entropic transport and use its uncoupled softmax solution in our primary experiments. Our analysis explains when variation among probe costs produces motion and when that motion decreases the loss. On a regular simplex, nonconstant costs always give a nonzero direction. For monotone ridge losses, the softmax update cannot increase the loss at any positive temperature; a perturbation bound gives sufficient conditions for descent near curved jumps. For bounded measurable losses, we randomize the probe radii and identify an exact smoothing whose gradient equals the expected linear cost-weighted direction up to scale. This identity yields a stationarity bound for an idealized fixed-temperature variant: under regularity and sampling assumptions stronger than those met by our trained configurations, the bound has an O(T^−1/2) term and a persistent bias floor. We evaluate the practical method on networks with hard operations, discrete optimization, and policy search. On MNIST with hard-threshold spiking neurons, PolyStep reaches 93.0 ± 0.2%, compared with 79.6 ± 5.2% for the best-tuned gradient-free baseline at matched evaluations. These gains come with a query cost proportional to the search dimension per fresh step, which limits the number of updates available at a fixed budget.
ER -