Training Non-Differentiable Networks via Optimal Transport

Abstract: Neural networks increasingly embed non-differentiable components (spiking neurons, quantized layers, discrete routing, blackbox simulators, etc.) where backpropagation is inapplicable and surrogate gradients introduce bias. We present PolyStep, a gradient-free optimizer that updates parameters using only forward passes. Each step evaluates the loss at structured polytope vertices in a compressed subspace, computes softmax-weighted assignments over the resulting cost matrix, and displaces particles toward low-cost vertices via barycentric projection. This update corresponds to the one-sided limit of a regularized optimal-transport problem, inheriting its geometric structure without Sinkhorn iterations. PolyStep trains genuinely non-differentiable models where existing gradient-free methods collapse to near-random accuracy. On hard-LIF spiking networks we reach 93.4% test accuracy, outperforming all gradient-free baselines by over 60~pp and closing to within 4.4~pp of a surrogate-gradient Adam ceiling. Across four additional non-differentiable architectures (int8 quantization, argmax attention, staircase activations, hard MoE routing) we lead every gradient-free competitor. On MAX-SAT scaling from 100 to 1M variables, we sustain above 92% clause satisfaction while evolution strategies drop 8--12~pp. On RL policy search, we match OpenAI-ES on classical control and retain performance under integer and binary quantization that collapses gradient-based methods. We prove convergence to conservative-stationary points at rate $O(\log T/\sqrt{T})$ on piecewise-smooth losses, upgraded to Clarke-stationary on the headline architectures and extended to the piecewise-constant regime via a hitting-time bound. These rates match the known zeroth-order query-complexity lower bounds that all forward-only methods inherit. Code is available at https://github.com/anindex/polystep.

• The paper introduces PolyStep, a gradient-free optimizer that leverages optimal transport to train networks with non-differentiable components.
• The paper presents theoretical convergence guarantees and empirical benchmarks demonstrating superior performance over existing gradient-free approaches across diverse tasks.
• The approach enables effective training of challenging architectures such as spiking, quantized, and combinatorial models while significantly reducing memory usage.

Training Non-Differentiable Neural Networks with PolyStep and Optimal Transport

Overview

The paper "Training Non-Differentiable Networks via Optimal Transport" (2605.01928) introduces PolyStep, a gradient-free optimizer capable of training neural networks containing genuinely non-differentiable components (e.g., spiking neurons, quantized layers, discrete routing, hard argmax attention, blackbox modules). PolyStep leverages optimal transport (OT) theory to guide parameter updates using only forward-pass loss evaluations, sidestepping the limitations of backpropagation and surrogate gradient methods. The approach utilizes structured polytope vertex perturbations in compressed parameter subspaces and computes softmax-weighted updates via barycentric projection, thereby inheriting the geometric structure of regularized OT optimizers without computationally expensive Sinkhorn iterations. PolyStep demonstrates both theoretical guarantees and strong empirical performance across a diverse spectrum of architectures and optimization regimes.

Motivation and Problem Statement

The increasing prevalence of non-differentiable modules in neural architectures poses significant challenges for standard training methods. Classical backpropagation is dependent on differentiability, while surrogate gradient and STE approaches introduce bias and require delicate tuning, often failing on architectures where small parameter perturbations yield no useful gradient signal. Existing gradient-free techniques (e.g., ES, SPSA, recent ZO methods like DeepZero, MeZO) require model smoothness to provide informative updates and degrade to random search when applied to hard-non-differentiable modules. PolyStep directly targets this gap by designing an optimizer that accepts arbitrary nn.Module programming, requires only forward losses, and is equipped to operate on models with discontinuous forward passes.

Algorithmic Framework

PolyStep's optimization is structured as follows:

• Parameters are projected into compressed subspaces via per-layer decompositions to mitigate dimensionality and computational burden. The optimizer operates on a population of "particles."
• For each particle, PolyStep samples step directions from a polytope (default: orthoplex/cross-polytope) under random rotations, generating candidate perturbations.
• Losses are evaluated at polytope vertices (probe points) using only forward passes. These losses form a particle-by-vertex cost matrix.
• A temperature-controlled softmax is applied to each cost row, yielding a transport plan that acts as a soft assignment from particles to low-cost vertices.
• Particles are updated via barycentric projection toward their assigned vertices, composing the new parameter vector.
• The step direction and assignment recover the exact one-sided limit of entropic OT (with empirical equivalence to softmax in subspace regimes), allowing for structured exploration and descent without gradient computation.

Subspace compression techniques (HybridSubspace, LinearSubspace, AdaptiveSubspace) and memory-efficient sparse projections are deployed for scalability, enabling PolyStep to operate in regimes up to millions of variables.

Theoretical Guarantees

PolyStep's optimization rule is supported by rigorous analysis:

• The paper proves convergence rates to conservative-stationary points ($O(\log T/\sqrt{T})$) for piecewise-smooth losses (i.e., losses differentiable almost everywhere except for a measure-zero discontinuity set). This covers architectures involving spikes, rounding, argmax, and staircase activations.
• For piecewise-constant objectives, PolyStep guarantees finite-time hitting of target level sets under mild conditions.
• The update rule is formally derived as the one-sided limit of regularized OT, with monotonic tightening via KL penalties and empirical validation for marginal-violation bounds.
• The theoretical analysis explains PolyStep's superiority over finite-difference or surrogate-gradient-based ZO methods on non-differentiable models, due to joint probe transversality and probability of escaping measure-zero non-smooth regions.

Empirical Results

PolyStep is systematically evaluated on five families of non-differentiable models and several benchmark tasks:

• Genuinely Non-Differentiable Networks:
  • Hard-LIF spiking networks: PolyStep achieves 93.4% test accuracy, outperforming all gradient-free baselines by over 60 percentage points and within 4.4 pp of surrogate-gradient-trained Adam.
  • INT8 quantized, staircase, hard argmax, and hard Mixture-of-Experts routing networks: PolyStep consistently outperforms ES, SPSA, and ZO methods, closing the gap to surrogate-gradient ceilings.
• Combinatorial Optimization:
  • MAX-SAT scaling from 100 to 1M variables: PolyStep maintains >92% clause satisfaction, while ES baselines drop 8–12 pp; performance remains robust across four orders of magnitude.
• RL Policy Search:
  • PolyStep matches OpenAI-ES on standard continuous control but retains performance under INT8 and binary quantization that causes PPO, DQN, and other gradient-based methods to collapse.
• Smooth Vision and Time-Series Benchmarks:
  • On MNIST, PolyStep achieves 96.0%, outperforming all gradient-free baselines and coming within 1.9 pp of Adam.
  • For ETTh1 oil temperature forecasting, PolyStep leads all other learned methods in MSE, surpassing Adam and ES.
• Memory Scaling:
  • PolyStep attains sub-linear memory on SNN recurrent evaluation ($\sim$30$\times$ reduction at $T=400$ compared to BPTT), with no activation checkpointing.

Ablation studies confirm that smooth softmax-weighted updates are essential; hard greedy moves or top-$k$ rules collapse to random accuracy. The full OT machinery delivers advantages in high-particle regimes and PolyStep's structured polytope probes are critical for functional signal extraction.

Strong Claims and Contradictory Evidence

• PolyStep outperforms all gradient-free baselines (CMA-ES, OpenAI-ES, SPSA, DeepZero, MeZO) on non-differentiable architectures by wide margins ($\geq$ 60 percentage points on SNN).
• PolyStep is robust to quantization, argmax routing, and staircase-induced discontinuities—all settings where gradient-based and finite-difference gradient-free methods fail or return zero gradients.
• On piecewise-constant combinatorial objectives (e.g., MAX-SAT), PolyStep remains competitive at million-variable scale, contrasting with the rapid degradation of ES methods.
• A sharp empirical limitation is observed: when applied to large transformer models trained from scratch (e.g., SST-2 with 4.2M parameters), PolyStep and other gradient-free optimizers collapse to near-random accuracy, strictly matching theoretical zeroth-order query complexity lower bounds.

Implications and Future Directions

Practical Implications

PolyStep fills a critical niche for neural architectures untrainable via standard methods due to forward non-differentiability, hardware constraints, or blackbox modules:

• Training is possible where neither surrogate gradients nor finite-difference estimators yield a valid signal.
• The approach is compatible with inference-only hardware, distributed setups, and hardware-in-the-loop configurations.
• PolyStep's forward-only computational model confers memory advantages particularly salient for SNNs and long-horizon, recurrent networks.

Theoretical Implications

• PolyStep demonstrates that optimal transport can serve as a principled framework for gradient-free descent, bridging geometric theory and practical neural optimization.
• The softmax/OT plan enables meaningful updates even on piecewise-smooth or constant objectives, expanding the class of trainable networks beyond those accessible with current ZO or surrogate methods.
• The convergence guarantees are robust for discontinuous loss landscapes, substantiating the update rule's correctness.

Speculation on Future AI Developments

• Forward-only optimization strategies like PolyStep can accelerate progress in neuromorphic computing, embedded inference, and combinatorial/blackbox optimization tasks, especially where gradient access is fundamentally restricted.
• Advances in subspace parametrization, OT solver acceleration, and distributed forward evaluation could further scale gradient-free model training toward deeper networks and NLP-sized architectures.
• Integration of OT-based descent with hybrid gradient-based recipes could enable selective optimization of non-differentiable submodules within differentiable networks.
• As inference-optimized hardware becomes prevalent, PolyStep's memory and compute profile may become increasingly relevant for real-world applications.

Conclusion

PolyStep constitutes a formally principled, practically validated gradient-free optimizer for genuinely non-differentiable neural architectures. Its OT-guided polytope probing enables structured, effective descent where all prior gradient-free and ZO baselines fail, and delivers strong empirical results on non-differentiable, combinatorial, RL, and even smooth tasks. The theoretical guarantees, empirical superiority, and hardware-aligned memory model identify PolyStep as a significant addition to the toolkit for hard-to-train models, closing longstanding gaps in neural network optimization methodologies.

Open-source code is available at https://github.com/anindex/polystep.