tjeongggg/wave-diffusion-548m

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.