Un-0, from Unconventional AI, draws images by a route that is neither diffusion nor autoregression. It numerically integrates a field of coupled Kuramoto oscillators until they settle, then decodes their locked phases into pixels. The released CIFAR-10 checkpoint carries 19.4M parameters and takes 102 ms to draw a batch of 1,024 images on an A100. We got it to 3.2M parameters and 17.7 ms, with the quality metric sitting exactly where it started.
Most of that path is a verification story. The quality gate rejected a higher-capacity model, discarded a later checkpoint that scored worse than an earlier one, and charged a small measured tax for the fast precision. Three separate attempts to squeeze more out with a hand-written kernel or a graph capture mostly failed, and the failures are the useful part, because they mark where the real bottleneck actually is. Dead-ends included.
Un-0’s generator is that settling. A field of oscillators starts scattered, a learned coupling pulls them toward a shared rhythm, and the phases they lock into are read out as an image. Here is the mechanism at small scale, running live. Reset to re-scatter the phases and watch the order parameter climb from 0 toward 1.
Hold the ruler, beat everything under it
The target was not a smaller model that is almost as good. It was to hold Un-0’s own quality ruler exactly, clean-FID on CIFAR-10, and beat the size and latency underneath it. clean-FID is fixed as the gate before any optimization runs. A change that moves it does not count, however much it saves. That one rule is what turns the rest of this into a verification story as much as a speed story, because every lever has to clear the gate before its number is allowed to appear.
Where the weight actually is
Un-0 ties its 4,096 oscillators together with one learned, dense 4,096×4,096 coupling matrix, K. That single matrix is 16,777,216 parameters, 86.3% of the model’s 19,434,123. Everything else, the natural frequencies ω and the phase-to-pixel decoder, is a fixed backbone of exactly 2,656,907 parameters.
That backbone is a floor. It is byte-for-byte the same count in the dense model, in the compressed model, and in the higher-capacity variant tried later; only the coupling ever changes.
| Config | Coupling K | Backbone (ω + decoder) | Total |
|---|---|---|---|
| Dense (released) | 16,777,216 | 2,656,907 | 19,434,123 |
| Monarch depth-2 (shipped) | 524,288 | 2,656,907 | 3,181,195 |
| Monarch depth-3 + low-rank | 1,310,720 | 2,656,907 | 3,967,627 |
The backbone is identical to the last digit across all three: 19,434,123 - 16,777,216 = 3,181,195 - 524,288 = 3,967,627 - 1,310,720 = 2,656,907. Whatever differs between models is coupling, nothing else.
The coupling is also where the runtime concentrates, for the same structural reason it dominates size. K is applied at every one of the 10 Euler solver steps, and twice per step, so a single generated batch pays 20 applications of an O(n²) operator with n=4,096. One matrix is simultaneously 86% of the stored parameters and the dominant per-step compute.
Why the coupling is applied 20 times per generated batch
Each sample integrates the oscillator ODE with 10 explicit Euler steps. Computing the velocity at each
step needs the coupling applied to two different vectors, sinθ and cosθ (the reason is the
angle-difference expansion in the next section), so every step is 2 coupling applications and a batch is
20. At O(n²) per application, that one matrix decides both numbers that matter: how big the model is on
disk, and how long a step takes on the GPU. Any change that shrinks it moves size and latency together,
which is why it is the first and largest target.
There is a second, unrelated inefficiency in how the model runs. Un-0 ships computing in fp32 on the GPU’s general-purpose cores. The tensor cores, the matmul units several times faster that sit on every modern NVIDIA GPU, go entirely unused by default. That is a dispatch choice, not a modeling one, and it can be changed without touching a weight.
So there are two different costs to attack, and they are independent. The tensor-core one is close to settled the moment it is named. The math is already matmuls, the fast units are already on the die, and flipping the dispatch retrains nothing. The coupling is the harder question. It is 86% of the parameters and the dominant per-step compute, so it is plainly the thing to shrink, but “shrink it” is not yet a method. The reflex is to take the trained 4,096×4,096 K and factor it, an SVD or a low-rank fit to the matrix already learned. That reflex smuggles in two assumptions, that a learned dense coupling sits near some low-rank matrix, and that the model uses K as a matrix at all. Neither is given. A dense coupling has no reason to be near low-rank, and if the dynamics ever read K element by element, or needed its transpose or its eigenvectors, any structured stand-in would change what the model draws. Before writing a line of factorization code, the question worth settling is the narrow one. Does Un-0 use K as a matrix, or only as an operator that sends a vector to K·v? That is answered by reading the equation, not by touching code.
Read the equation before touching code
The equation is what decides where compression is even possible. The Kuramoto velocity for oscillator i is:
dθᵢ/dt = ωᵢ + Σⱼ Kᵢⱼ·sin(θⱼ - θᵢ)Expanding the coupling with the angle-difference identity is the decisive step:
sin(θⱼ - θᵢ) = sinθⱼ·cosθᵢ - cosθⱼ·sinθᵢSubstituted back, the per-oscillator velocity becomes ωᵢ + cosθᵢ·(K·sinθ) - sinθᵢ·(K·cosθ). K appears
only as K·sinθ and K·cosθ, a linear operator applied to two vectors. It is never read element by
element and never needed as a matrix. So K does not have to be a dense matrix; it only has to act like
one on vectors. Any object with a fast K·v and far fewer parameters can stand in, and the physics is
unchanged.
That answers the question the size table raised. The plans that start by factoring the trained K can be dropped, because nothing has to stay close to the dense matrix when nothing ever touches the matrix. What has to be reproduced is the action, K·v, and that reframes the whole search. Instead of approximating one specific 16.8M-number matrix, the job is to pick any operator class with a cheap K·v and enough freedom to learn the right action, then train it from scratch. Monarch is one such class.
The velocity expansion, term by term
Write the coupling sum with the identity sin(θⱼ - θᵢ) = sinθⱼ·cosθᵢ - cosθⱼ·sinθᵢ. The i-indexed
factors cosθᵢ and sinθᵢ do not depend on the summation index j, so they pull out of the sum:
Σⱼ Kᵢⱼ·sin(θⱼ - θᵢ) = cosθᵢ·Σⱼ Kᵢⱼ·sinθⱼ - sinθᵢ·Σⱼ Kᵢⱼ·cosθⱼ. The two remaining sums are exactly
(K·sinθ)ᵢ and (K·cosθ)ᵢ, two matrix-vector products against the same K. That is the whole reason a
step costs two coupling applications and not one, and the whole reason K is only ever consumed as an
operator. A structured object that computes K·v quickly reproduces the dynamics up to its own
arithmetic.
The precision lever is orthogonal and needs no retraining. fp32-on-general-cores is a dispatch default,
not a property of the weights. Turning on TF32 and a mixed-precision autocast routes the same matmuls
onto the vendor tensor-core kernels that were idle, with torch.compile fusing the elementwise glue
around them. Same weights, same code path, different kernel.
The two levers are independent. One is a flag on unchanged weights, the other is a retrained operator. They attack different costs, so they stack. That independence is the bet of the whole project. Two low-risk levers that multiply beat one risky rewrite that might not land at all.
Lever 1: the tensor-core path, no retraining
The released weights already run correctly in mixed precision, so lever 1 is the cheapest thing to try and it goes first. Turning on TF32 plus a bf16 autocast (fp16 on the T4) is the entire change. On an A100 batch of 1,024 it takes the sample from 102.244 ms to 21.202 ms, 4.82× on weights nobody retrained. The T4 gets 3.36× (2,138 to 7,187 img/s). The A100 gains more because its tensor cores are proportionally faster relative to its general cores.
This lever is banked first because it carries no quality risk in principle, since the weights are identical. But in principle is not a measurement. bf16 changes the arithmetic, so the fast path has to clear clean-FID before it ships, and it does that in the gating section below, not here.
Why lever 1 is a dispatch change, not a hand-written kernel
Tensor cores are matmul units that take low-precision inputs and add each product into a running
accumulator. cuBLAS and cuDNN already ship tuned kernels for them; the model simply was not dispatching
to those kernels. fp32-on-general-cores is what PyTorch selects by default for safety. TF32 plus
autocast changes only which existing kernel is chosen, and torch.compile fuses the surrounding sin,
cos, and scale operations into it. Nothing is written by hand. That is exactly why the 4.82× is free: it
is not new compute, it is unused compute switched on.
Lever 2: a Monarch-structured coupling
The K·v-only insight is what lets the second lever exist at all. K is replaced by a Monarch operator:
a product of square block-diagonal factors with a fixed transpose between them. For n=4,096 each factor
is √n = 64 blocks of 64×64, so applying it is a block-diagonal matmul plus a permutation, O(n^1.5)
instead of O(n²). A depth-2 product is two 262,144-parameter factors, 524,288 coupling parameters in
all.
- Coupling: 16,777,216 → 524,288, exactly 32× fewer numbers.
- Whole model: 19,434,123 → 3,181,195, a factor of 6.109×, so 6.1× smaller.
The 32× and the 6.1× are different quantities and both are correct. 32× is the coupling alone; 6.1× is the whole model, held back by the 2,656,907-parameter backbone that does not compress. The coupling stops being the bulk of the model, dropping from 86.3% of the parameters to 16.5% of a much smaller total.
The whole change is two edits to un0/model.py, behind a coupling_type flag. The dense branch is left
byte-identical, so --coupling-type dense still reproduces Un-0 exactly.
How a Monarch factorization replaces a dense 4,096×4,096 matrix
A Monarch matrix writes a large dense operator as a product of two block-diagonal matrices with a fixed
permutation between them. For n = 4,096 that is two factors of 64 blocks of 64×64, so the coupling’s
parameter count drops from n² = 16,777,216 to 2 × 64 × 64² = 524,288, a 32× cut, and applying it costs
O(n^1.5) instead of O(n²). The structure works here only because the dynamics never need K as a matrix.
They need K·sinθ and K·cosθ, and a block-diagonal matmul plus a transpose produces exactly those.
The zero-initialized low-rank residual, and why depth-2 is trained from scratch
A purely block-local Monarch operator can only mix oscillators inside the same 64-block; a purely global
low-rank operator is rank-limited and cannot express fine local structure. Each has a blind spot the
other covers. So the operator carries an optional low-rank residual (x·V)·Uᵀ with U zero-initialized,
LoRA-style, and the model starts as pure Monarch and grows a global path only if training asks for it.
This is also the reason the compressed model is trained from scratch rather than fitted to the released dense K after the fact. Block-local and globally-rank-limited structure form a strictly larger function class when learned jointly than either does when bolted onto a fixed matrix that already committed to neither. Training the two together is a modeling choice with a clear rationale; retrofitting a trained dense matrix is not.
6.1× smaller is not 6.1× faster
Here is where the first easy assumption breaks. 6.1× fewer parameters is not 6.1× faster. Compression alone, fp32, with no precision change, takes the A100 batch-1024 sample from 102.244 ms to 64.900 ms. That is 1.575×, call it 1.6×, nowhere near 6.1×.
The reason is structural, not a tuning miss. Monarch’s saving is many small 64×64 block matmuls, and small matmuls are kernel-launch-bound: the wall-clock is dominated by the fixed overhead of launching each kernel, not by the arithmetic inside it. Cutting the arithmetic 32× does not cut the wall-clock 32× when the wall-clock was never arithmetic in the first place.
Launch-bound, not compute-bound: what the 1.6× is measuring
A dense 4,096×4,096 matmul is one large kernel that keeps the GPU’s units saturated, so its cost really is arithmetic. Splitting it into 64 independent 64×64 blocks replaces that with dozens of tiny kernels, each doing very little work and each paying a fixed launch and scheduling cost to start. Below a certain size that fixed cost dominates, so throughput is set by how many kernels you launch, not how many multiply-adds they do. Monarch removes almost all of the multiply-adds and almost none of the launches, which is why fp32 compression alone lands at 1.575× while the parameter count fell 32×. The fix is not fewer blocks. It is running the blocks on the tensor-core path so each launch retires far more work, which is lever 1 doing its job on lever 2’s output.
This is why the two levers are not alternatives to choose between. Compression removes the parameters and the n² arithmetic but leaves launch-bound small matmuls. The tensor-core path turns those into fast tensor-core work. Neither alone reaches 5–6×. Stacked, they clear it, and the model that runs both is also the small one.
Stacking them, and the gate that said no twice
The compressed model is trained from scratch, so parity is not inherited; it has to be earned on Un-0’s own metric. The gate is clean-FID over 50,000 class-balanced CIFAR-10 samples, at two seeds.
With the ruler validated, the shipped depth-2 model comes in at 8.93 / 8.99 (min / mean, fp32) against the dense 8.88, a gap that sits inside the roughly 0.1 seed-to-seed noise floor. The 6.1× compression holds on quality.
The checkpoint that ships is not the last one trained, and that distinction is the gate doing its job. The depth-2 sweep sat at FID around 10.1–10.3 through epochs 700–1000, dropped to 8.93 at epoch 1100, then regressed to 9.22 at epoch 1200. The gate picks the best converged checkpoint, epoch 1100, not the final one. More epochs was not more quality.
The depth-3 dead-end
More capacity was the obvious next move, and it failed the gate. A depth-3 Monarch plus a rank-64 low-rank residual, a third factor and a global channel, is 786,432 + 524,288 = 1,310,720 coupling parameters, 3,967,627 in all, still 4.9× smaller than dense. Trained to epoch 1200, it scored 9.00 / 9.04. That is worse than depth-2’s 8.93. The extra factor and the global path bought nothing the metric could see, so the smallest structured model is the sweet spot. depth-3 is reported as a variant and not shipped.
Why depth-3 is 9.04, not the 8.96 the training log flashed
A single generation during depth-3 training logged an 8.96, close to depth-2, and it would have been easy to quote. Averaged over both evaluation seeds the model is 9.04. The 8.96 was one lucky draw, not a checkpoint that beats depth-2. Reporting the two-seed number instead of the best single sample is the same discipline that shipped epoch 1100 over epoch 1200: the gate reads the converged, averaged result, never the flattering instant.
One more thing has to clear the gate before any speed number counts, the bf16 fast path itself. bf16 is the deploy precision, so it is measured on clean-FID, not waved through as merely precision. It costs +0.09 FID on the dense model and +0.05 on depth-2, a small, disclosed tax with no quality cliff. At the actual deploy operating point both models land at 8.98 on the best seed. That is the headline, 8.98 = 8.98, stated on the bf16 path that delivers the speed rather than on a flattering fp32 comparison.
| Model | Params | FID fp32 | FID bf16 deploy | bf16 tax |
|---|---|---|---|---|
| Dense (released) | 19.43M | 8.88 (8.90) | 8.98 (8.99) | +0.09 |
| Monarch depth-2 (ours) | 3.18M | 8.93 (8.99) | 8.98 (9.04) | +0.05 |
| Monarch depth-3 + low-rank | 3.97M | 9.00 (9.04) | n/a | n/a |
On the bf16 deploy path the shipped models are dead even at 8.98 on the best seed. depth-3 scored worse than depth-2 in fp32, so it was never carried to a bf16 run.
The custom-kernel work that mostly failed
With both levers banked at 102 ms → 17.7 ms, one question is still open. Is there a hand-written win left
under torch.compile, from a bespoke kernel or a graph capture? It is worth treating as genuinely open
rather than assuming either way, and the answer after profiling is mostly no. The shape of that no is the
finding.
The profile is what decides where effort can even help, so nothing gets written until it is read. On an A100 batch-1024 sample (bf16, 10 iterations) the wall-clock is 18.41 ms, of which 11.86 ms is GPU-busy. The other 6.5 ms, 35.6%, is host-side idle, the Python solver loop launching kernels one step at a time while the GPU waits between launches. Of the GPU-busy time, the coupling is still the largest single slice at 23.2%, even after a 32× parameter cut, because Monarch’s small blocks lower to a memory-bound GEMV rather than a tensor-core GEMM.
- GPU-busy64.4% · 11.86 ms
- host-idle35.6% · 6.5 ms
- Image-decoder GEMMs/convs30%
- Coupling GEMV (gemvx)23.2% · largest single op, post-compression
- bf16-cast + K-scale multiply12.6%
- Velocity recombine6%
- Monarch transpose1.2%
- Remainder27%
An independent cross-check confirms the coupling share. A standalone MonarchCoupling forward runs
0.4179 ms, and 0.4179 × 20 applications per sample is 8.36 ms, which matches the 23% GEMV slice. So the
profile splits the wall-clock into two targets that no single tool reaches at once. One is a compute
slice, the coupling GEMV, still the biggest op on the GPU and exactly the kind of thing a hand-written
kernel exists for. The other is host orchestration, a third of the wall-clock spent not computing at
all, which no per-step kernel can touch and only a change to the loop can. Each of the three attempts
below aims at one of these two, and each returns a result that says as much about where the bottleneck is
as about the kernel it tried.
A Triton coupling kernel: a real win, capped by Amdahl
The first attempt is a Triton kernel that fuses each depth-2 factor into one launch: a tensor-core block
matmul whose output index folds the transpose in, so the coupling is 2 launches with no materialized
GEMV and no separate transpose, against roughly 6 launches from the compiler. It is bit-exact in fp32
(max relative error 0.000e+00), 1.36e-3 in TF32, 4.06e-3 in bf16, and it runs 1.07× (fp32) to 1.34×
(bf16) faster than torch.compile on the coupling in isolation. A real win on that slice. But the
coupling is 23% of compute, so Amdahl’s law caps the end-to-end gain at about 1.1×, and that is a
projection from the slice, not a measured end-to-end run. It also does nothing for the 35% host idle.
Not where the time is, so not shipped.
The Amdahl arithmetic behind the ~1.1× ceiling
Amdahl’s law bounds an end-to-end speedup by the fraction of time the optimization actually touches. The
coupling is p = 0.232 of GPU-busy compute, and the kernel is s = 1.34× on that fraction. The overall
speedup is 1 / ((1 - p) + p/s) = 1 / (0.768 + 0.232/1.34) = 1.063×. The docs round that to about 1.1×.
Even an infinitely fast coupling kernel would cap at 1 / 0.768 = 1.30×, and that ceiling ignores the
35.6% host idle entirely, which sits outside the GPU-busy fraction the law is computed over. A 1.34×
kernel on a 23% slice is a real result and a small end-to-end one at the same time, and both statements
are true.
Widening the kernel hands the win back to the compiler
The next attempt widens the kernel to the whole per-step velocity: sin and cos, both factors, the
K-scale, and the ω + cos·ws - sin·wc recombine fused into one path, roughly 10 launches per step down
to 5. It is correct (fp32 max relative error 1.70e-7, bf16 1.16e-2) and slower than torch.compile:
0.83× in fp32, 0.69× in bf16. This is the key negative result. The moment the kernel widens past the
isolated GEMV, Inductor wins, because it computes sin and cos once and reuses them, folds the recombine
into matmul epilogues, and picks better tiles than a hand-written fixed 64×64. The 5a win exists only
because it attacks the one primitive the compiler lowers badly. Give the compiler back the surrounding
elementwise work and it beats the hand kernel.
Two graph-capture attempts, both fail
The 35% host idle is the real target, and CUDA graphs are the standard tool for it. Both ways of
reaching for them failed. torch.compile(mode="reduce-overhead"), which uses managed CUDA graphs, moved
17.82 ms to 18.61 ms, a 0.96× that helped nothing. A manual torch.cuda.CUDAGraph() capture of the full
10-step rollout failed outright with cudaErrorStreamCaptureInvalidated. The root cause is the same for
both: the model already runs torch.compile internally, and its compiled regions own their own
CUDA-graph trees. An outer manual capture cannot wrap them, and managed graphs cannot span the Python
solver loop that sits between the compiled regions.
Why both graph-capture attempts fail with one root cause
A CUDA graph records a fixed sequence of kernel launches once and replays it with almost no host
overhead, which is exactly the cure for a launch-bound Python loop. It requires the captured region to be
static and self-contained. Un-0’s sample() is neither from the outside: the 10 solver steps run in
Python, and inside each step the model calls its own torch.compile-generated regions, which already
build and manage CUDA-graph trees of their own. reduce-overhead mode adds managed graphs but can only
graph within a compiled region, not across the Python loop between them, so it nets 0.96×. A manual outer
capture tries to record across the loop and immediately hits the inner compiled graphs mid-stream, which
invalidates the capture (cudaErrorStreamCaptureInvalidated). Fixing it means replacing the Python
solver loop with one static, graph-safe rollout, a change to the integration loop rather than a kernel.
Past the one memory-bound GEMV the compiler handles badly, hand-written kernels give diminishing and then negative returns here. The ceiling on bespoke GPU code is about 1.1×, and the real remaining slack, the 35% host idle, is not a kernel problem at all. It is a solver-loop problem, and the solver loop is where the next lever lives.
How the quality gate held
The whole project holds together on one discipline. The gate is clean-FID, and negative results are reported, not buried. Five checks carry the credibility of every number above.
- The ruler is validated before it is trusted. Re-measuring the released dense checkpoint to 8.88–8.90, matching Un-0’s own 8.86 and 8.76 within noise, is what earns the right to compare anything against it.
- The gate is allowed to say no, and did. depth-2 passes at 8.93, inside the 0.1 noise floor. depth-3 fails at 9.00 and is dropped despite more capacity. The shipped depth-2 checkpoint is epoch 1100, not the regressed epoch 1200 at 9.22.
- The fast path pays its own bill. bf16 is not waved through as precision; it is measured at +0.09 on dense and +0.05 on depth-2, and the parity headline is stated at the bf16 deploy point, not a flattering fp32 one.
- The kernel investigation is negative by design. A 1.34× microbenchmark win is demoted to about 1.1× by Amdahl and not shipped. A wider kernel that loses to the compiler at 0.69× is reported as the key finding. Two failed graph captures are shown with their root cause.
- Correctness precedes speed. The coupling kernel is bit-exact in fp32 (0.000e+00) before its 1.34× is quoted; the velocity kernel’s fp32 error (1.70e-7) is reported next to its bf16 error (1.16e-2). No speed claim rests on an unverified kernel.
What we learned, and the next lever
A few things bound these numbers, and each one marks where the work goes next. Stating them precisely is what turns a limit into a direction.
The biggest lever still on the table is algorithmic, not a kernel. Un-0 runs 10 solver steps, and cost is close to linear in step count, since each step launches the same kernels and carries the same host idle. Cutting to 7 steps projects about 1.43×, and to 6 about 1.67×, a further 1.4–1.6× on top of everything above. It sits outside every measured number here on purpose, because fewer steps changes the generated samples, and that has to clear the same clean-FID gate the rest of the work cleared. That FID re-check is the one thing between the projection and a shipped number, and it is the next thing to run.
Where this goes
Two independent, low-risk, multiplicative levers carried the result at unchanged clean-FID.
| Lever | A100 batch 1024 | T4 batch 1024 |
|---|---|---|
| Tensor-core path only (no retraining) | 4.82× | 3.36× |
| Compression only (fp32) | ~1.6× | n/a |
| Full stack vs dense as shipped | 5.77× | 4.97× |
| Size | 6.1× smaller | same |
| clean-FID, bf16 deploy path | 8.98 = 8.98 | same |
Full stack: A100 102 → 17.7 ms (10,015 → 57,747 img/s); T4 2,138 → 10,631 img/s. Size 19.43M → 3.18M, coupling 32×. Every row holds clean-FID against the released model.
A precision flip that needs no retraining and a structural compression trained to a fixed quality bar together take the CIFAR-10 checkpoint 6.1× smaller and 5.77× faster, clean-FID held at 8.98. A bespoke kernel, tried three ways, did not beat that, and the reason is now precise. The compiler already owns everything except one memory-bound GEMV, and that one primitive is Amdahl-capped near 1.1×. So the next win is not a kernel, it is the solver loop. Cutting its 10 steps to 6 or 7 projects a further 1.4–1.6×, and a single clean-FID re-check is the one thing standing between that projection and a shipped number.