Training Non-Differentiable Networks via Optimal Transport

• URL: http://arxiv.org/abs/2605.01928v1
• Date: Sun, 03 May 2026 15:20:56 GMT
• Title: Training Non-Differentiable Networks via Optimal Transport
• Abstract summary: PolyStep is a gradient-free method that updates parameters using only forward passes.<n>We train genuinely non-differentiable models where existing gradient-free methods collapse to near-random accuracy.<n>On hard-LIF spiking networks we reach 93.4% test accuracy, outperforming all gradient-free baselines by over 60pp and closing to within 4.4pp of a surrogate-gradient Adam ceiling.
• Score: 1.5229257192293202
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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.

Related papers

• Closed-Form Spectral Regularization for Multi-Task Model Merging [96.82449201305234]
  Model merging combines several independently fine-tuned experts into a single multi-task model without any training data.<n>State-of-the-art merging methods formulate merging as a layer-wise interference problem.<n>We propose SWUDI, a closed-form method that combines a soft exponential filter, which matches the gradient-flow trajectory of iterative descent.
  arXiv  Detail & Related papers  (2026-06-05T14:00:47Z)
• Towards Scalable Persistence-Based Topological Optimization [44.16669776030478]
  Persistence-based topological optimization deforms a point cloud $X subset mathbbRd$ by minimizing objectives of the form $L(X) = ell(mathrmDgm(X))$, where $mathrmDgm(X)$ is a persistence diagram.<n>In practice, optimization is limited by two coupled issues: persistent homology is typically computed on subsamples, and the resulting topological gradients are highly sparse, with only a few anchor points receiving nonzero updates.
  arXiv  Detail & Related papers  (2026-05-09T15:47:20Z)
• PolyGLU: State-Conditional Activation Routing in Transformer Feed-Forward Networks [0.0]
  We introduce PolyGLU, a drop-in replacement for SwiGLU that enables each FFN neuron to dynamically route among K=4 activation functions.<n>On standard benchmarks, PolychromaticLM achieves 62-89% of Qwen3-0.6B-Base performance despite training on 3,600x fewer tokens.
  arXiv  Detail & Related papers  (2026-03-07T10:39:56Z)
• Layer-wise QUBO-Based Training of CNN Classifiers for Quantum Annealing [0.0]
  We propose an iterative framework based on Quadratic Un Binary Optimization (QUBO) for training the head of convolutional neural networks (CNNs)<n>A per-output decomposition splits the $C$-class problem into $C$ independent QUBOs, each with $(d+1)K$ binary variables, where $d$ is the feature dimension and $K$ is the bit precision.<n>We evaluate the method on six image-classification benchmarks (sklearn digits, MNIST, Fashion-MNIST, CIFAR-10, EMNIST, KMNIST)
  arXiv  Detail & Related papers  (2026-03-03T13:10:36Z)
• Closed-Form Last Layer Optimization [72.49151473937319]
  Under a squared loss, the optimal solution to the linear last layer weights is known in closed-form.<n>We show this is equivalent to alternating between gradient descent steps on the backbone and closed-form updates on the last layer.
  arXiv  Detail & Related papers  (2025-10-06T09:14:39Z)
• Decentralized Nonconvex Composite Federated Learning with Gradient Tracking and Momentum [78.27945336558987]
  Decentralized server (DFL) eliminates reliance on client-client architecture.<n>Non-smooth regularization is often incorporated into machine learning tasks.<n>We propose a novel novel DNCFL algorithm to solve these problems.
  arXiv  Detail & Related papers  (2025-04-17T08:32:25Z)
• Unified Gradient-Based Machine Unlearning with Remain Geometry Enhancement [29.675650285351768]
  Machine unlearning (MU) has emerged to enhance the privacy and trustworthiness of deep neural networks. Approximate MU is a practical method for large-scale models. We propose a fast-slow parameter update strategy to implicitly approximate the up-to-date salient unlearning direction.
  arXiv  Detail & Related papers  (2024-09-29T15:17:33Z)
• Nearly Minimax Optimal Regret for Learning Linear Mixture Stochastic Shortest Path [80.60592344361073]
  We study the Shortest Path (SSP) problem with a linear mixture transition kernel. An agent repeatedly interacts with a environment and seeks to reach certain goal state while minimizing the cumulative cost. Existing works often assume a strictly positive lower bound of the iteration cost function or an upper bound of the expected length for the optimal policy.
  arXiv  Detail & Related papers  (2024-02-14T07:52:00Z)
• Bridging Discrete and Backpropagation: Straight-Through and Beyond [62.46558842476455]
  We propose a novel approach to approximate the gradient of parameters involved in generating discrete latent variables. We propose ReinMax, which achieves second-order accuracy by integrating Heun's method, a second-order numerical method for solving ODEs.
  arXiv  Detail & Related papers  (2023-04-17T20:59:49Z)
• Orthogonalizing Convolutional Layers with the Cayley Transform [83.73855414030646]
  We propose and evaluate an alternative approach to parameterize convolutional layers that are constrained to be orthogonal. We show that our method indeed preserves orthogonality to a high degree even for large convolutions.
  arXiv  Detail & Related papers  (2021-04-14T23:54:55Z)
• GradInit: Learning to Initialize Neural Networks for Stable and Efficient Training [59.160154997555956]
  We present GradInit, an automated and architecture method for initializing neural networks. It is based on a simple agnostic; the variance of each network layer is adjusted so that a single step of SGD or Adam results in the smallest possible loss value. It also enables training the original Post-LN Transformer for machine translation without learning rate warmup.
  arXiv  Detail & Related papers  (2021-02-16T11:45:35Z)
• A Unified Analysis of First-Order Methods for Smooth Games via Integral Quadratic Constraints [10.578409461429626]
  In this work, we adapt the integral quadratic constraints theory to first-order methods for smooth and strongly-varying games and iteration. We provide emphfor the first time a global convergence rate for the negative momentum method(NM) with an complexity $mathcalO(kappa1.5)$, which matches its known lower bound. We show that it is impossible for an algorithm with one step of memory to achieve acceleration if it only queries the gradient once per batch.
  arXiv  Detail & Related papers  (2020-09-23T20:02:00Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.