Dust: Zeroth‑Order Pretraining of Transformers Without Backpropagation

TL;DR

Dust is the first zeroth‑order method that rivals backpropagation on transformer pretraining. By perturbing activations token‑wise (a virtual population), it evaluates thousands of candidates in a single forward pass, achieving gradient estimates that converge to backprop’s as population grows and even surpass backprop in large‑compute regimes.


1. What Dust Does Differently

Activation‑Space Perturbation

  • Dust adds independent Gaussian noise to the output of each linear layer per token. Each token therefore acts as a separate population member, allowing a single forward pass to evaluate a population size on the order of the sequence length (thousands of members) rather than one per weight‑space ES member.
  • This “virtual population” avoids materializing weight copies; the perturbation lives in activation space, which is lower‑dimensional than the full weight matrix for a given layer.

Credit Assignment Without a Backward Pass

  • For a linear layer (y_t = Wx_t), Dust computes a reward (r_t) for each token based on the loss reduction at that token and, optionally, discounted future tokens (via attention). The reward‑weighted noise (\hat g_t = -\frac{1}{K\sigma}\sum_{i=1}^K r_t^{(i)} a_t^{(i)}) serves as the estimated output error.
  • The weight gradient estimate is then the outer product of (\hat g_t) with the cached input (x_t): [\hat G_W = \sum_t \hat g_t x_t^{\top}]
  • Attention internals are treated similarly, but their reward is derived from the change in attention output rather than token‑wise loss directly, reducing interference.

Reducing Interference

  1. Layer‑wise passes – Different layer types are jittered in separate forward passes, each with its own noise scale.
  2. Attention‑specific jitter – Only the attention output is recomputed for a jitter, keeping the rest of the network cached.
  3. Head‑only perturbation – The language‑model head is jittered on cached logits, evaluating a tiny slice of the vocabulary per draw.

These tricks keep the cost of a draw only marginally higher than a clean forward pass while allowing population sizes of (10^3)–(10^4) draws per update.


2. Empirical Findings

Pretraining Results (SGD Momentum)

Token budget Population (draws) Test loss (Dust) Test loss (Backprop) Test loss (EGGROLL)
100 k 64‑1 k below backprop – –
1 M 1 k‑16 k Approaches backprop, gap < 0.02 – –
10 M 4 k‑16 k Slightly above backprop, trend downward – –
20 M 16 k Extrapolation suggests Dust could undercut backprop (limit ≈ 4.43 vs backprop 4.63) – –

Dust consistently stays within a few hundredths of backprop’s loss, and at the largest populations the fitted power‑law limit lies below backprop.

Comparison to Weight‑Space ES (EGGROLL)

  • Even with 10³–10⁴× larger populations, EGGROLL lags behind Dust by 0.02–0.6 loss units across all token budgets.
  • Extrapolations indicate EGGROLL would need 10³–10⁴ times the population to match Dust’s smallest‑population performance.

Scaling with Model Size

Parameters Population 64 256 1 k 4 k 16 k
2 M 5.705 5.486 5.265 5.189 5.171
7 M 5.556 5.362 5.161 5.095 5.065
38 M 5.558 5.358 5.158 5.053 5.036
243 M 5.719 5.419 5.214 5.124 5.086

Larger models are more population‑efficient: a 243 M‑parameter model outperforms a 2 M model even at the smallest population and continues to improve where the tiny models plateau.

Gradient Alignment

Dust’s estimated gradients align with true backprop gradients increasingly with population. The cosine similarity follows: [\cos(K) = \frac{c_{\text{max}}}{\sqrt{1 + c/K}}]

  • Fits across layers achieve RMSE < 0.06.
  • At populations ≥ 16 k, cosine values exceed 0.9 for most layers, indicating near‑identical directionality.
  • By contrast, EGGROLL’s cosine stays < 0.05 even at 128 k draws, explaining its poor training performance.

3. Community Reactions on Hacker News

Commenter Key Insight / Critique
blt Warns that derivative‑free methods have historically failed to impact practice; argues gradients remain essential for smooth objectives and cites recent complexity‑gap proofs.
syntacticsalt Questions the “Bitter Lesson” framing, noting that zeroth‑order methods do not inherently solve non‑convexity and that the paper lacks a discussion of convexity assumptions.
eriwang915 Highlights the surprising result that a 243 M‑parameter Dust model beats a 120× smaller model across most population sizes.
polyomino Suggests a hybrid pipeline: fine‑tune a backprop‑trained checkpoint with Dust to explore whether later training stages benefit from the search‑based credit assignment.
soltanov Calls for wall‑clock time, energy, and memory measurements to assess practical viability; notes the current efficiency gap must close before Dust can replace backprop.
mkaic Points out that zeroth‑order methods may be valuable where gradients are unavailable (e.g., simulator‑based training) rather than as a universal replacement.
usernametaken29 Argues both Dust and backprop lie on the same Empirical Risk Minimisation frontier, but Dust could bypass Hessian‑conditioning limits that hinder backprop.
api Observes that Dust is less computationally efficient per update but highly parallelizable, which may be advantageous on massive hardware.

4. Why Dust Matters

  1. Proof of Concept for Search‑Based Credit Assignment – Demonstrates that a pure forward‑only algorithm can generate usable gradients for large‑scale language models, challenging the long‑standing belief that zeroth‑order methods cannot scale.
  2. Potential to Surpass Backprop in Compute‑Rich Regimes – At very large populations, Dust’s loss curve extrapolates below backprop, suggesting that with enough compute the algorithm could find better descent directions than the chain‑rule gradient.
  3. Enables Training of Non‑Differentiable Architectures – Since Dust does not require analytic gradients, it opens the door to models that incorporate discrete simulators, external programs, or recurrent loops that are hard to differentiate through.
  4. Architectural Insight: Overparameterization Benefits – Larger models exhibit better population efficiency, hinting that the geometry of high‑dimensional activation spaces may be more favorable for search‑based optimization than previously thought.

5. Open Questions and Future Directions

  • Efficiency Gap – Dust currently requires orders of magnitude more FLOPs than backprop. Engineering advances (e.g., specialized hardware, smarter population allocation) are needed to make it practical.
  • Hybrid Training Regimes – Combining backprop pretraining with Dust fine‑tuning could leverage the strengths of both approaches.
  • Higher‑Order Landscape Exploration – Understanding whether Dust implicitly captures curvature information (beyond first‑order) could explain its occasional superiority.
  • Broader Applicability – Applying Dust to tasks where gradients are unavailable (reinforcement learning with simulators, neuromorphic hardware) would test its generality.
  • Theoretical Foundations – Formal analysis of the population‑size vs. gradient‑alignment trade‑off, especially in non‑convex settings, would solidify the empirical observations.

6. Conclusion

Dust shows that zeroth‑order activation perturbations can train transformer language models at a level comparable to backpropagation, with larger models even becoming more sample‑efficient. While the method is far from replacing backprop in everyday practice due to its compute demands, it provides a compelling proof‑of‑concept that search‑based credit assignment can scale, opening new research avenues for non‑differentiable architectures and compute‑rich optimization regimes.

Sources

Related