Does the whole coordinate cloud look like an iid Gaussian sample?
the-puzzler Matteo
AAG · GENERATIVE MODELLING
Amortised Assignment Generation
Turn a dataset into Gaussian assignments, then learn to generate in one forward pass.
Summary
Generative modelling can be split into representation and reconstruction. AAG solves representation first by assigning every training sample a fixed Gaussian latent, then learns one-step generation as ordinary supervised reconstruction. It can converge quickly and samples in one forward pass, but depends on having enough data to cover the effective latent dimension.
RUNNING IN YOUR BROWSER
Play the Doom world model
Each click encodes the latest three frames, samples a fresh latent \(z\), and predicts the next frame for the selected action. The generated frame is then fed back into the context.
The model was trained on 10M frames of Doom. It runs locally after a one-time 25 MB download. No frames or controls leave the browser.
seed context
Preview ready. Load the model to enable controls.
RUNNING IN YOUR BROWSER
Sample conditional faces
Checkbox changes keep the same latent batch, so each condition modulates all eighteen faces. The New faces button samples fresh \(z\) values.
The demo starts with Male and Receding Hairline selected. Nine controls are exposed here.
Load the 14 MB model to generate faces.
Introduction
Methods such as GANs and denoising models can be viewed as solving two problems at once: deciding how samples should be represented by noise, and learning how to reconstruct them from that representation. If we could solve the assignment problem first, generation itself could become ordinary supervised learning.
Implicit Maximum Likelihood Estimation points in this direction. For each data example, IMLE searches sampled latent codes for one whose generated output is nearby, then trains the generator on that match. Reversing the usual regression direction ensures that every real example gets a generated neighbour, rather than allowing several possible outputs to collapse toward their mean. The difficulty is that repeatedly searching through noise during training does not scale gracefully.
Amortised Assignment Generation removes that search from generator training. It gives every example a fixed Gaussian coordinate ahead of time, choosing the assignment to disturb the geometry of a learned representation as little as possible. The generator is then trained directly on those persistent pairs.
Gaussianise the data. Supervise the assignment.
Theory
Suppose each example \(x_i\) is paired with a coordinate \(z_i\), and the set \(\{z_i\}\) is indistinguishable from a finite sample of \(\mathcal N(0,I)\). A generator trained on the fixed pairs \((z_i,x_i)\) can then be sampled with fresh \(z\sim\mathcal N(0,I)\). Matching every one-dimensional projection is enough to determine the joint distribution by Cramér–Wold, which motivates the sliced rank transports used here.
The assignment must also be smooth enough for the generator to interpolate between coordinates. In the conditional case there is one more requirement: \(z\perp c\). Otherwise, sampling an independent \(z\) with a chosen condition \(c\) asks the generator to handle combinations it never saw during training.
INTERACTIVE THEORY
Representation
The moon-half label c is attached to every point.
Generate
The main scaling bottleneck is not the cost of globally moving the latent dataset, but whether \(N\) is large enough relative to the effective sampled dimension for the decoder to interpolate between assignments.
\[\text{computational cost}\sim O(Nd+N\log N)\]
\[\text{statistical difficulty}\sim N^{-1/d_{\mathrm{eff}}}.\]
Statistical Difficulty
Around a fresh prior sample \(z\), suppose the prior density is \(p(z)\). In a small \(d_{\mathrm{eff}}\)-dimensional ball of radius \(r\), the probability mass is approximately
\[p(z)V_{d_{\mathrm{eff}}}r^{d_{\mathrm{eff}}},\]
where \(V_{d_{\mathrm{eff}}}\) is the volume of the unit ball in \(d_{\mathrm{eff}}\) dimensions. With \(N\) assigned samples, the expected number falling in that neighbourhood is therefore
\[N p(z)V_{d_{\mathrm{eff}}}r^{d_{\mathrm{eff}}}.\]
Setting this equal to \(k\), the typical distance to the \(k\)-th neighbour is
\[r_k(z)\approx\left(\frac{k}{N p(z)V_{d_{\mathrm{eff}}}}\right)^{1/d_{\mathrm{eff}}},\]
and therefore, ignoring constants,
\[r_k\propto N^{-1/d_{\mathrm{eff}}}\quad\Longleftrightarrow\quad\text{coverage}\propto N^{1/d_{\mathrm{eff}}}.\]
This is the relevant curse of dimensionality: halving the characteristic gap requires roughly \(2^{d_{\mathrm{eff}}}\) times more samples.
Under conditioning, replace \(N\) and \(d_{\mathrm{eff}}\) with the locally relevant sample count and residual stochastic dimension:
\[\text{conditional coverage}\propto N_{\mathrm{eff}}^{1/d_{\mathrm{eff},z}}.\]
Conditioning reduces \(N_{\mathrm{eff}}\), but if it removes enough stochastic degrees of freedom, the residual problem can still become substantially easier.
Computational Cost
For \(N\) samples in \(d\) latent dimensions, a global projection step requires a matrix-vector product over all latent codes, sorting one scalar projection per sample, and applying the resulting rank update:
\[O(Nd)+O(N\log N)\]
per accepted direction.
These operations are highly GPU-friendly. Candidate directions can be evaluated in parallel as a matrix multiplication, while the expensive full-dataset operation is only a projection and scalar sort. Even millions of samples with \(d\sim64\text{-}256\) require only hundreds of megabytes to a few gigabytes of latent storage in FP16, so the full persistent assignment can remain practical on modern accelerators.
Method
// construct the persistent assignment
z ← PCA-whiten(E(x))
repeat
a ← worst random projection of z // score on a subset, apply to all
z ← rank-transport(z, a, Gaussian orders)
for each of c conditional steps // if conditioning
S ← examples with a similar condition
// exact group, or k-NN in condition space
a ← worst projection of z[S]
z[S] ← partial-rank-transport(z[S], a)
optionally, interleaved
every 2 steps:
slab-cleanup(z)
// Gaussianise a tangent coordinate where |n·z - b| < ε
every 20 steps:
radial-calibration(z)
// rank-correct ||z|| toward the χ_d law
until
global: transport objective is inside its N(0,I) noise floor
conditional: independence ratio I has flattened near 1
// amortise the fixed mapping
train G(z ⊕ condition) → x // direct to pixels, MSE + LPIPSSorting a projection and matching its ranks to Gaussian order statistics is exact one-dimensional optimal transport. The surrounding choices, including greedy direction search, neighbourhood sampling and the balance of global and conditional steps, are heuristics.
The optional cleanup passes catch structure that global slices can miss. Offset-slab cleanup targets local tail spikes by Gaussianising a tangent coordinate inside a thin slab. Radial calibration corrects \(\lVert z\rVert\) toward the exact \(\chi_d\) law while leaving directions untouched. This matters in high dimensions because shell error can survive every one-dimensional projection test.
Once assignment is finished, generation is simply \(G(z)\) in a single forward pass, with \(z\sim\mathcal N(0,I)\).
Useful Metrics
I use three diagnostics for an assignment:
Do condition-local slices look like random slices of the same size?
Can nearby prior coordinates locally predict their assigned targets?
In the experiments so far, G and I predicted the better generator every time. R detects deliberately scrambled assignments, but mis-ranked both real comparisons with held-out FID, so I treat it as a description rather than a selection metric. Generator training loss is also misleading: it rewards assignments that moved the least.
Results
In practice, an autoencoder is first trained on the raw data to make assignment tractable. For images, this can use MSE and LPIPS. The decoder is immediately discarded, while the encoder is used once to embed the dataset and is then discarded too. These encoded representations undergo the assignment process. Finally, a new generator is trained on persistent pairs that map the assigned Gaussian coordinates back to the original samples.
Unconditional
CelebA unconditional training curves
Conditional
FM Comparison
I also trained flow matching models with similar parameter counts to compare their performance against AAG.
My general impression is that AAG converges much more quickly than flow matching and reaches a similar FID using only one forward pass. It may, however, be more sensitive to the amount of data. On CIFAR-10 it plateaus at a higher FID, while its monotonic improvement as more data is added suggests that the available dataset is still small relative to the effective sampled dimensionality.
Conclusion
AAG is based on the idea that generative modelling is two problems: representation and reconstruction. If one can solve representation first, reconstruction becomes trivial supervised learning. This is evidenced by the fast convergence of the FID, especially in the case of CelebA. The limitation of treating the dataset as a set of particles to be Gaussianised is that, if the dataset is sparse or has a low number of samples relative to its effective dimensionality, the resulting Gaussian cloud has large holes. For flow matching, this is less problematic because denoiser training covers the entire latent distribution, so it can learn to divert trajectories towards support. AAG, being a fixed supervised learner, cannot learn a function over space that is empty to begin with.
When AAG works, however, it clearly shows signs of competent generation, is capable of conditional generation, and achieves equivalent FID to flow matching with fewer parameters, training steps and less inference compute.
Going forward, I think AAG could be a standalone method as well as being combined with GANs and iterative denoising. By creating a fixed mapping over the latent space, mode coverage is guaranteed because every real sample has a unique supervised latent code. This could stabilise GAN training by preventing mode collapse, and could accelerate flow matching training.