PolyStep: Gradient-Free Forward Optimizer

• PolyStep is a forward-only, gradient-free optimizer that navigates non-differentiable model tasks by evaluating loss at vertices of a local polytope around parameter blocks.
• It computes soft, cost-dependent weights using a row-wise softmax or OT plan to update parameters through barycentric averaging, requiring only batched forward passes.
• Empirical results show PolyStep excelling on tasks like hard-LIF spiking MNIST and reinforcement learning, with convergence guarantees for both piecewise-smooth and piecewise-constant losses.

PolyStep is a forward-only, gradient-free optimizer for training models whose forward pass is genuinely non-differentiable. It targets settings such as hard leaky integrate-and-fire spiking neurons, integer rounding, argmax routing, staircase activations, hard mixture-of-experts gates, black-box simulators, reinforcement-learning policy search, and piecewise-constant discrete objectives. Rather than estimating gradients, PolyStep evaluates the loss at vertices of a small polytope around many local parameter blocks, computes soft cost-dependent weights over those vertices, and moves each block to a barycentric average of low-cost vertices. The practical update uses only batched forward passes and a row-wise softmax, while the paper shows that this update is the one-sided limit of a regularized optimal-transport problem; the name “PolyStep” is explicitly tied to a polytope-based step with OT geometry (Le, 3 May 2026).

1. Problem setting and conceptual basis

PolyStep is designed for models in which non-differentiability is part of the forward computation itself, not merely an optimization inconvenience. The paper lists hard operators such as $\mathrm{sign}(\cdot)$, $\mathrm{round}(\cdot)$, $\mathrm{argmax}(\cdot)$, and $\mathbb{1}[u \ge u_{\mathrm{th}}]$ as representative examples. In such systems, exact backpropagation is either undefined or uninformative, and common replacements such as surrogate gradients or straight-through estimators optimize a smoothed or substituted backward model rather than the actual hard forward computation (Le, 3 May 2026).

The method is also positioned against finite-difference and evolution-strategy baselines. The paper argues that SPSA, DeepZero-style perturbation zeroth-order methods, and MeZO-style estimators fail on piecewise-constant losses because small perturbations usually remain in the same flat region, so the estimate is zero with high probability. Evolution strategies such as OpenAI-ES, CMA-ES, PEPG, and ARS are treated as perturbation-based search-distribution methods rather than structured local optimizers. PolyStep differs in that it does not form a perturb-and-regress gradient estimate. Instead, it evaluates a structured set of candidate directions, organizes the resulting losses into a cost matrix, computes soft assignments, and performs a barycentric displacement toward low-cost vertices (Le, 3 May 2026).

This yields a specific niche. PolyStep is presented as strongest when exact gradients are unavailable or unusable, and when surrogate-gradient fidelity to the true hard forward pass is unacceptable. It is not presented as a general replacement for first-order optimization on differentiable models; the paper explicitly notes that Adam remains faster, more accurate, and more scalable when gradients exist (Le, 3 May 2026).

2. Mathematical formulation

Let the model parameters be $\theta \in \mathbb{R}^d$. PolyStep either works in full space or in a compressed subspace of dimension $d_{\text{sub}}$. Parameters in the optimization space are reshaped into particles,

$$X = [x_1,\dots,x_P]^\top \in \mathbb{R}^{P \times d_p}, \qquad P = \left\lceil \frac{d_{\text{sub}}}{d_p}\right\rceil,$$

where each row $x_i \in \mathbb{R}^{d_p}$ is a particle. The default polytope is the orthoplex

$$\mathcal V = \{\pm e_j\}_{j=1}^{d_p}, \qquad V = 2d_p.$$

For each particle, PolyStep samples a rotation

$R_i \sim \mathrm{Uniform}(\mathrm{SO}(d_p)),$

and defines candidate vertices

$\mathrm{round}(\cdot)$0

Probe points are then placed along each vertex direction:

$\mathrm{round}(\cdot)$1

with jitter

$\mathrm{round}(\cdot)$2

The cost matrix averages probe losses,

$\mathrm{round}(\cdot)$3

so $\mathrm{round}(\cdot)$4 stores particle-by-vertex costs (Le, 3 May 2026).

The default source marginal is uniform,

$\mathrm{round}(\cdot)$5

and the practical update is a row-wise softmax,

$\mathrm{round}(\cdot)$6

Each particle is updated by barycentric projection,

$\mathrm{round}(\cdot)$7

This is the practical PolyStep algorithm used for the headline experiments (Le, 3 May 2026).

The OT interpretation is given by the balanced entropic transport problem

$\mathrm{round}(\cdot)$8

with uniform marginals $\mathrm{round}(\cdot)$9 and $\mathrm{argmax}(\cdot)$0, together with a one-sided KL-penalized interpolation

$\mathrm{argmax}(\cdot)$1

At $\mathrm{argmax}(\cdot)$2, the solution reduces to the row-wise softmax plan; at $\mathrm{argmax}(\cdot)$3, it converges to the full entropic OT plan. This is the formal basis for describing the practical rule as the one-sided limit of a regularized optimal-transport problem (Le, 3 May 2026).

3. Algorithmic procedure and implementation variants

A PolyStep iteration consists of reshaping parameters into particles, constructing a local rotated polytope around each particle, evaluating all probe losses in parallel, assembling the cost matrix, computing either row-wise softmax weights or a full OT plan, and updating particles by barycentric projection. The per-iteration forward-query cost is $\mathrm{argmax}(\cdot)$4. For the orthoplex with $\mathrm{argmax}(\cdot)$5 and $\mathrm{argmax}(\cdot)$6, this is $\mathrm{argmax}(\cdot)$7 forward evaluations; with $\mathrm{argmax}(\cdot)$8, it becomes $\mathrm{argmax}(\cdot)$9 (Le, 3 May 2026).

The paper describes several subspace variants. In full-space mode, $\mathbb{1}[u \ge u_{\mathrm{th}}]$0. In HybridSubspace mode, each layer $\mathbb{1}[u \ge u_{\mathrm{th}}]$1 uses a fixed projection

$\mathbb{1}[u \ge u_{\mathrm{th}}]$2

with reconstruction $\mathbb{1}[u \ge u_{\mathrm{th}}]$3. LinearSubspace and AdaptiveSubspace are also reported. Among these, HybridSubspace is described as best and most stable in the ablations (Le, 3 May 2026).

The default polytope is the orthoplex, but simplex and cube variants are also described. Particle dimensions $\mathbb{1}[u \ge u_{\mathrm{th}}]$4 are studied explicitly. The update rule can use either the softmax solver or full entropic OT with Sinkhorn iterations; the paper states that softmax and OT are identical in the common subspace regime, while OT becomes useful in the high-particle regime $\mathbb{1}[u \ge u_{\mathrm{th}}]$5 where the column marginal matters (Le, 3 May 2026).

Several “turbo” features are listed: biased rotation, warm-started duals for OT, Anderson acceleration, adaptive overrelaxation $\mathbb{1}[u \ge u_{\mathrm{th}}]$6, amortized OT or plan reuse, and dual momentum warm-start. These are presented as engineering accelerations rather than defining components of the method. The paper further notes practical details relevant to reproduction: theoretical convergence requires $\mathbb{1}[u \ge u_{\mathrm{th}}]$7, but reported experiments default to $\mathbb{1}[u \ge u_{\mathrm{th}}]$8 for deterministic reproducibility; new applications are recommended to use $\mathbb{1}[u \ge u_{\mathrm{th}}]$9; attention and recurrent layers require VmapSafe replacements; and SNN logits were multiplied by $\theta \in \mathbb{R}^d$0 before cross-entropy to increase probe cost contrast (Le, 3 May 2026).

4. Convergence theory and stationary notions

The theoretical analysis begins with a piecewise-smooth class of losses. A bounded measurable $\theta \in \mathbb{R}^d$1 is piecewise smooth if it is $\theta \in \mathbb{R}^d$2 on $\theta \in \mathbb{R}^d$3, locally Lipschitz on connected components of $\theta \in \mathbb{R}^d$4, and $\theta \in \mathbb{R}^d$5 is a closed semialgebraic or definable set of codimension at least one. The paper states that this covers hard LIF, rounding, argmax routing, and floor or staircase activations (Le, 3 May 2026).

Under orthoplex probing, $\theta \in \mathbb{R}^d$6, i.i.d. random rotations, the schedule

$\theta \in \mathbb{R}^d$7

and jitter $\theta \in \mathbb{R}^d$8 with $\theta \in \mathbb{R}^d$9, the smoothed surrogate

$d_{\text{sub}}$0

satisfies

$d_{\text{sub}}$1

Every limit point is conservative-stationary. For the headline architectures—hard-LIF, INT8 rounding, hard-MoE argmax routing, and staircase activations—the paper states that conservative-stationarity upgrades to Clarke-stationarity (Le, 3 May 2026).

The paper also studies the piecewise-constant regime. For $d_{\text{sub}}$2 with a positive-measure success set and a non-symmetric polytope such as the simplex under a uniform reachability condition, there exists $d_{\text{sub}}$3 such that the first hitting time

$d_{\text{sub}}$4

obeys

$d_{\text{sub}}$5

This gives a finite-time hitting guarantee for a piecewise-constant success region. The paper notes that this theorem does not directly apply to the orthoplex because its vertex sum is zero (Le, 3 May 2026).

A separate proposition analyzes schedule fragility. If a row of the cost matrix has a unique best vertex with margin $d_{\text{sub}}$6, then for sufficiently small $d_{\text{sub}}$7 the softmax becomes nearly one-hot and the step magnitude is pinned near $d_{\text{sub}}$8. The paper concludes that $d_{\text{sub}}$9 jointly is necessary for asymptotic stability, and that decaying $$X = [x_1,\dots,x_P]^\top \in \mathbb{R}^{P \times d_p}, \qquad P = \left\lceil \frac{d_{\text{sub}}}{d_p}\right\rceil,$$0 while holding $$X = [x_1,\dots,x_P]^\top \in \mathbb{R}^{P \times d_p}, \qquad P = \left\lceil \frac{d_{\text{sub}}}{d_p}\right\rceil,$$1 fixed is fragile. This theoretical point is reflected in the empirical schedule ablations (Le, 3 May 2026).

5. Empirical performance

The empirical evaluation spans supervised non-differentiable learning, combinatorial optimization, reinforcement learning, smooth benchmarks, and scaling or memory studies. The headline result is hard-LIF spiking MNIST, where PolyStep reaches $$X = [x_1,\dots,x_P]^\top \in \mathbb{R}^{P \times d_p}, \qquad P = \left\lceil \frac{d_{\text{sub}}}{d_p}\right\rceil,$$2 test accuracy, outperforming all gradient-free baselines by over $$X = [x_1,\dots,x_P]^\top \in \mathbb{R}^{P \times d_p}, \qquad P = \left\lceil \frac{d_{\text{sub}}}{d_p}\right\rceil,$$3 percentage points and closing to within $$X = [x_1,\dots,x_P]^\top \in \mathbb{R}^{P \times d_p}, \qquad P = \left\lceil \frac{d_{\text{sub}}}{d_p}\right\rceil,$$4 percentage points of a surrogate-gradient Adam ceiling (Le, 3 May 2026).

Setting	PolyStep	Comparator context
Hard-LIF spiking MNIST	$$X = [x_1,\dots,x_P]^\top \in \mathbb{R}^{P \times d_p}, \qquad P = \left\lceil \frac{d_{\text{sub}}}{d_p}\right\rceil,$$5	CMA-ES $$X = [x_1,\dots,x_P]^\top \in \mathbb{R}^{P \times d_p}, \qquad P = \left\lceil \frac{d_{\text{sub}}}{d_p}\right\rceil,$$6, OpenAI-ES $$X = [x_1,\dots,x_P]^\top \in \mathbb{R}^{P \times d_p}, \qquad P = \left\lceil \frac{d_{\text{sub}}}{d_p}\right\rceil,$$7, SPSA $$X = [x_1,\dots,x_P]^\top \in \mathbb{R}^{P \times d_p}, \qquad P = \left\lceil \frac{d_{\text{sub}}}{d_p}\right\rceil,$$8, surrogate Adam $$X = [x_1,\dots,x_P]^\top \in \mathbb{R}^{P \times d_p}, \qquad P = \left\lceil \frac{d_{\text{sub}}}{d_p}\right\rceil,$$9
INT8 quantization	$x_i \in \mathbb{R}^{d_p}$0	Adam surrogate $x_i \in \mathbb{R}^{d_p}$1
Argmax attention	$x_i \in \mathbb{R}^{d_p}$2	Adam $x_i \in \mathbb{R}^{d_p}$3
Staircase activations	$x_i \in \mathbb{R}^{d_p}$4	Adam $x_i \in \mathbb{R}^{d_p}$5
Hard MoE routing	$x_i \in \mathbb{R}^{d_p}$6	Leads all gradient-free competitors
Binary weights	$x_i \in \mathbb{R}^{d_p}$7	STE $x_i \in \mathbb{R}^{d_p}$8
Ternary weights	$x_i \in \mathbb{R}^{d_p}$9	STE $$\mathcal V = \{\pm e_j\}_{j=1}^{d_p}, \qquad V = 2d_p.$$0
MAX-SAT, $$\mathcal V = \{\pm e_j\}_{j=1}^{d_p}, \qquad V = 2d_p.$$1 variables	$$\mathcal V = \{\pm e_j\}_{j=1}^{d_p}, \qquad V = 2d_p.$$2 clause satisfaction	OpenAI-ES $$\mathcal V = \{\pm e_j\}_{j=1}^{d_p}, \qquad V = 2d_p.$$3
MNIST, smooth baseline	$$\mathcal V = \{\pm e_j\}_{j=1}^{d_p}, \qquad V = 2d_p.$$4	Adam $$\mathcal V = \{\pm e_j\}_{j=1}^{d_p}, \qquad V = 2d_p.$$5
ETTh1 forecasting	$$\mathcal V = \{\pm e_j\}_{j=1}^{d_p}, \qquad V = 2d_p.$$6 MSE	Adam $$\mathcal V = \{\pm e_j\}_{j=1}^{d_p}, \qquad V = 2d_p.$$7

Across MAX-SAT instances from $$\mathcal V = \{\pm e_j\}_{j=1}^{d_p}, \qquad V = 2d_p.$$8 to $$\mathcal V = \{\pm e_j\}_{j=1}^{d_p}, \qquad V = 2d_p.$$9 variables, PolyStep stays above $R_i \sim \mathrm{Uniform}(\mathrm{SO}(d_p)),$0 clause satisfaction, whereas the paper states that ES methods drop by $R_i \sim \mathrm{Uniform}(\mathrm{SO}(d_p)),$1–$R_i \sim \mathrm{Uniform}(\mathrm{SO}(d_p)),$2 percentage points. On the $R_i \sim \mathrm{Uniform}(\mathrm{SO}(d_p)),$3M-variable case, the run took $R_i \sim \mathrm{Uniform}(\mathrm{SO}(d_p)),$4 s on a single RTX 5090 with about $R_i \sim \mathrm{Uniform}(\mathrm{SO}(d_p)),$5 GB peak memory, enabled by delta evaluation over affected clauses (Le, 3 May 2026).

In reinforcement learning, PolyStep matches OpenAI-ES on CartPole-v1 and Acrobot-v1. On CartPole it attains $R_i \sim \mathrm{Uniform}(\mathrm{SO}(d_p)),$6 in Float32, INT8, and Binary, and on Acrobot it reports $R_i \sim \mathrm{Uniform}(\mathrm{SO}(d_p)),$7 in Float32, $R_i \sim \mathrm{Uniform}(\mathrm{SO}(d_p)),$8 in INT8, and $R_i \sim \mathrm{Uniform}(\mathrm{SO}(d_p)),$9 in Binary. The paper emphasizes that PolyStep retains performance under integer and binary quantization that collapses PPO and DQN; on CartPole, PPO falls to $\mathrm{round}(\cdot)$00 in INT8 and $\mathrm{round}(\cdot)$01 in Binary (Le, 3 May 2026).

The ablations are structurally important. Entropic OT and softmax-weighted PolyStep both achieve $\mathrm{round}(\cdot)$02 on subspace MNIST, while hard selection rules collapse: min-cost greedy reaches $\mathrm{round}(\cdot)$03, and top-$\mathrm{round}(\cdot)$04 mean reaches $\mathrm{round}(\cdot)$05. In full-space MNIST with $\mathrm{round}(\cdot)$06 and $\mathrm{round}(\cdot)$07, OT yields $\mathrm{round}(\cdot)$08 versus $\mathrm{round}(\cdot)$09 for softmax, illustrating when column-marginal enforcement matters. Smaller particle dimension performs better in subspace mode: $\mathrm{round}(\cdot)$10 gives $\mathrm{round}(\cdot)$11, compared with $\mathrm{round}(\cdot)$12 for $\mathrm{round}(\cdot)$13 and $\mathrm{round}(\cdot)$14 for $\mathrm{round}(\cdot)$15. On SNNVGG11Small at $\mathrm{round}(\cdot)$16 timesteps, PolyStep uses $\mathrm{round}(\cdot)$17 MB, versus $\mathrm{round}(\cdot)$18 MB for BPTT, a $\mathrm{round}(\cdot)$19 memory reduction (Le, 3 May 2026).

6. Scope, limitations, and disambiguation

PolyStep’s main practical limitation is the standard zeroth-order one: high query cost. Each step requires $\mathrm{round}(\cdot)$20 forward evaluations, and the paper explicitly ties this to known zeroth-order query-complexity lower bounds. It also states that PolyStep does not beat backpropagation on differentiable models, is sensitive to schedules and hyperparameters such as $\mathrm{round}(\cdot)$21, $\mathrm{round}(\cdot)$22, amortization, and subspace rank, and fails to solve large from-scratch smooth NLP models: on a $\mathrm{round}(\cdot)$23M-parameter SST-2 transformer trained from scratch, all gradient-free methods are near-random. In full-precision smooth continuous control, PPO far outperforms PolyStep on the Unitree G1 locomotion frontier experiment (Le, 3 May 2026).

The name should also be distinguished from several adjacent “step”-based methods. “Topological Optimization with Big Steps” introduces a topology-aware large-step update strategy for persistence-based optimization, but it is explicitly not a method named PolyStep (Nigmetov et al., 2022). “Step-TP” is a grounded, step-level dataset for LLM-guided tensor program optimization built around LEIR, atomic optimization strategies, and structured chain-of-thought supervision; it is a stepwise optimization dataset rather than the forward-only optimizer described here (Liu et al., 25 May 2026). In structural mechanics, the STEP and M-STEP literature concerns non-intrusive identification of polynomial reduced-order models for geometrically nonlinear structures, including the STiffness Evaluation Procedure and a modified STEP for 3D finite elements; this is a distinct STEP-family usage of “step” terminology rather than the OT-based PolyStep optimizer (Vizzaccaro et al., 2020).

Within its intended scope, PolyStep is therefore best understood as a forward-only optimizer for genuinely non-differentiable forward passes, with a practical softmax implementation, an OT-derived geometric interpretation, convergence guarantees on piecewise-smooth and piecewise-constant regimes, and strong empirical performance precisely where backpropagation and standard perturbation-based zeroth-order methods are least reliable (Le, 3 May 2026).