Wave Diffusion LM (548M) — language modeling with a Schrödinger mixer instead of attention
A non-autoregressive masked diffusion language model whose tokens are complex wave packets, whose noising is phase decoherence, and whose sequence mixer is an attention-free split-step Schrödinger integrator (schrod2d) rather than self-attention. Trained on ~10B tokens of FineWeb-Edu.
What's novel here: to our knowledge this is the first playable wave-physics diffusion LM with a live wavefield visualization of the denoising process. It is not the first diffusion LLM (see LLaDA / Dream / Mercury) nor the first complex/Born-rule language model — the contribution is the combination, made trainable and interactive, plus a clean empirical finding about the readout (below).
TL;DR
- 548M params (schrod2d mixer · linear readout · MDLM objective), trained on 9.99B tokens of FineWeb-Edu, sequence length 512, one RTX-5090-week.
- Generates fluent, on-topic English prose — a "small playable model," not a knowledge oracle (it rambles and is not factually reliable).
- Attention-free,
O(L log L) sequence mixing via FFT-based wave propagation.
- Comes with an interactive demo (two decoding modes) + a wavefield visualization.
Reported numbers (honest, both metrics)
Measured on a held-out FineWeb-Edu split, at the trained sequence length (512):
| Metric | Value | What it is |
|---|
| Generative NELBO | 5.67 bits/token (±0.05) | The real marginal-likelihood upper bound. Two independent MD4 estimators agree to 0.005 bits (schedule-invariance holds), so the bound is valid and well-estimated. This is the load-bearing number. |
| Conditional masked-LM | ~3.13 bits/token | Fill-in-the-blank quality: predict 15%-masked tokens given the other 85% as clean context. A favorable conditional slice, not a likelihood — reported for transparency, not as the headline. |
(We deliberately report both. The conditional number is much lower because it is an easier task; the NELBO is the fair generative-modeling figure.)
The one clean research finding: the readout, not the wave dynamics
In matched comparisons against a width/depth/objective-matched transformer-MDLM baseline, the model's earlier gap traced entirely to the readout, not the wave physics. The original quantum Born-rule readout — logit(v) ∝ log |⟨e_v, ψ⟩|² — is magnitude-only, phase-collapsing, rank/PSD-constrained, and norm-invariant, so it discards information the predicted field carries. Replacing it with a plain linear head on [Re ψ; Im ψ] recovers that information.
Capacity caveat (stated up front): a full linear head has 2·m·V params vs the Born head's m·V, so the raw comparison is capacity-confounded. The load-bearing result is the parameter-matched bottleneck-linear head (sized to the Born head's m·V): it still beats log-Born by ~0.39 bits at matched parameters — so the win is the readout form, not extra capacity. Separately, the linear-head wave model matches the transformer baseline at equal params (there the wave arm is actually smaller: 547.9M vs 592.7M).
How it works (one paragraph)
Each token v is a learned unit phasor e(v) = exp(i·Φ[v]) ∈ ℂ^m — identity lives in phase. The forward (noising) process is phase decoherence: rotate each position's phases by θ(t)·ω, ω ~ N(0, I), until the phasor's mean spirals to zero — "fully masked" is the θ → ∞ limit. The denoiser predicts the clean field ψ̂; its mixer, schrod2d, integrates a dissipative Schrödinger equation along the sequence by split-step Fourier: a unitary dispersion step in frequency space (learned ω_c(k) = a₁k + a₂k² + a₃k³) plus a content-dependent potential + absorbing step in position space, O(L log L). A linear head reads ψ̂ out to token logits. Trained with the MDLM masked-diffusion objective; generation is iterative parallel denoising. A full study guide (math + physics) ships alongside.
Usage — two decoding modes
Both run on the tuned sampler (conf_floor=0.4, top_p=0.90, rep_penalty=3) and must sample at sequence length 512 (the learned dispersion is tied to the trained length):
- Continue a prompt → block-decode (semi-autoregressive; better prompt adherence)
- Infill / any-order edit → parallel (the native non-autoregressive strength)
See the demo Space for an interactive version with the wavefield visualization.
Limitations (read these)
- Not factually reliable. It's a 548M model; it produces plausible-sounding prose, not correct facts. It rambles and drifts.
- Fixed length 512.
schrod2d's dispersion is tied to the trained k-grid; it does not transfer to other sequence lengths. Sample at 512.
- A stylistic phrase quirk. It over-uses a few stock phrases (e.g. "side by side") learned from the corpus — a model trait, not a bug we've suppressed.
- Weak precise recall. Strong at distributed/statistical context (infill), weak at sharp associative lookup (MQAR-style copy).
- Research artifact, not a product. Intended to be played with and studied, not deployed.
Training data
FineWeb-Edu (HuggingFaceFW/fineweb-edu), tokenized with a 32k byte-level BPE. ~10B tokens. English.
Citation / prior art to credit
Diffusion LLMs (LLaDA, Dream, Mercury); MDLM (Sahoo et al., 2024); absorbing/discrete diffusion (D3PM, SEDD); Block Diffusion (Arriola et al., BD3-LMs, ICLR 2025); split-step Fourier method / nonlinear Schrödinger numerics; Born rule & POVMs (Nielsen & Chuang). This model combines a wave/Schrödinger mixer with masked diffusion and a linear readout.
