Training Non-Differentiable Networks via Optimal Transport
cs.LG, cs.NE, cs.RO, math.OC
Submitted: 2026-05-03
Updated: 2026-09-21
Comments: 95 pages, 29 tables, 8 figures. Accepted at Transactions on Machine Learning Research. Code: https://github.com/anindex/polystep
Code: https://github.com/anindex/polystep
License: http://creativecommons.org/licenses/by/4.0/
The gist: Hard thresholds, quantization, and discrete routing can produce training losses with flat regions and jumps, where ordinary gradients vanish or are undefined.
Terminology
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 plus or minus 0.2%, compared with 79.6 plus or minus 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.
Sources
- Convergence of a class of gradient-free optimisation schemes when the objective function is noisy, irregular, or both
- Gaussian Loss Smoothing Enables Certified Training with Tight Convex Relaxations
- Gradients without Backpropagation
- Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation
- Low-rank surrogate modeling and stochastic zero-order optimization for training of neural networks with black-box layers
- Gradient-Free Training of Quantized Neural Networks
- Node Perturbation Can Effectively Train Multi-Layer Neural Networks
- Mean-Field Model for Two-Layer Neural Networks Trained with Consensus-Based Optimization
- Nonsmooth Optimization with Zeroth Order Comparison Feedback
- A Tutorial on Bayesian Optimization
- The CMA Evolution Strategy: A Tutorial
- The Forward-Forward Algorithm: Some Preliminary Investigations
- Beyond Discreteness: Sample Complexity Analysis of Straight-Through Estimator for 1-bit Quantization
- On the Complexity of Deterministic Nonsmooth and Nonconvex Optimization
- Deterministic Nonsmooth Nonconvex Optimization
- SpikingGamma: Surrogate-Gradient Free and Temporally Precise Online Training of Spiking Neural Networks with Smoothed Delays
- LOTION: Smoothing the Optimization Landscape for Quantized Training
- A Gaussian smoothing-based zeroth-order method for Goldstein second-order stationarity
- Signal-Adaptive Trust Regions for Gradient-Free Optimization of Recurrent Spiking Neural Networks
- NoProp: Training Neural Networks without Full Back-propagation or Full Forward-propagation
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: Local-Consistency Optimisation for Bayesian Neural Networks
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks