The Native3D autoencoder, explained

What n3d/autoencoder.py does, part by part, and what it achieved. MAGGIE Stage 3, 2026-10-16. Code: /pscratch/sd/o/olawale/MAGGIE/E3SM/Native3D/n3d/autoencoder.py (338 lines, 4.07 M parameters).

In short

Why an autoencoder at all

A diffusion model learns by taking a state, adding noise to it, and training a network to remove the noise. It then generates new states by starting from pure noise and removing it step by step, 64 steps in our sampler. If the state were the raw 104.7-million-number field, every one of those steps would be a pass over 104.7 million numbers, and the network would have to learn the structure of the ocean from only 13,824 training months. That is too slow and too data-hungry.

So the prior does not work on the raw state. It works on the code: 3.7 million numbers arranged as a field on a coarse graph. The autoencoder is the translator between the two: it must compress enough to make the prior tractable, and reconstruct well enough that nothing the assimilation cares about is lost. The gate we set for it (Gate G2) was that the reconstruction of held-out months should explain at least 95 % of the variance of every variable.

The same thing as an analogy

Think of the ocean as a country of 465,044 towns (the mesh cells). Each town keeps a logbook with 335 entries: the temperature, salt and currents at each of its 64 floors from the surface down to the seabed, plus the sea level and the ice at the top, plus where the town is and how deep its seabed goes.

1. Every town writes a 128-word summary of its logbook (step 2). The same small writer is used for every town, so the summaries are comparable.
2. Every town reads its six neighbours' summaries and revises its own, twice (step 3). Now each summary also says what is going on within 90 km.
3. The country is divided into 29,000 districts of about 16 towns, drawn so that no district crosses land (step 4). Each district merges its towns' summaries into one 256-word district report.
4. Districts hold four rounds of regional meetings (step 5). In each round, a district reads the reports of its 32 nearest districts, paying more attention to some than to others, and revises its own. After four rounds a district report reflects the state of the ocean about 1,000 km around it.
5. The reports are condensed to 128 words each. The 29,000 condensed reports are the code (step 6).
6. To answer a question ("salt at this spot, 850 m down"), a clerk writes the question in a fixed form, then reads the report of the district the spot is in and the nine districts touching it, twice, weighing each by its position relative to the spot (steps 7 and 8). The clerk answers with a number and a confidence (step 9).

The whole machine is trained by asking 390,000 random questions about each month and comparing the clerk's answers with the real logbooks. Every part, writer, meetings and clerk, is adjusted together so the answers get better; nobody tells the district reports what to contain. That is what "learned" means here: the 128 words per district are whatever turned out to be most useful for answering questions.

The picture

In the code's own words: the encoder (steps 1–6) reads a whole state at once (no batch dimension; one state per GPU) and produces z of shape (29,000, 128); the decoder (steps 7–9) takes z and a set of queries and returns two numbers per query. The sections below follow the same nine steps, each with a plain statement of what happens, then the actual layer names and sizes.

Follow one point through the machine

Take a real place: 38° N, 60° W, in the Gulf Stream south of Nova Scotia, and ask for the salinity 850 m down in May 1980. Here is what happens to that one number.

For the full-state evaluation this is repeated for every one of the 25.1 million wet points and 11 variables, in chunks of 500,000 questions; it takes about a minute per month on one GPU.

1. Inputs: what a cell knows (cell_inputs)

What happens. Every column of water is written out as one fixed-length list of 335 numbers, so that one and the same small network can read every column.

The column is written into a dense (cells × 64 × 4) array through the index of wet points (wet_flat), which is how the cell-major wet-only storage of the data meets the per-cell view of the network. Stored z-scores are already of order one, which is why there is no normalisation layer here. "Z-score" throughout means: the value minus the climatological mean for that month and place, divided by the standard deviation at that depth, so a +2 at 4,000 m (where the variation is a few thousandths of a degree) is as unusual as a +2 at the surface (where it is a degree or more).

2. Column MLP: one token per cell (cell_in)

What happens. A two-layer network reads the 335 numbers of a column and writes a 128-number summary, the cell's token. It sees the whole column at once, so it can describe a thermocline or a cold deep layer as a shape rather than level by level.

h = MLP(335 → 256 → 128)(inputs)        # (465,044, 128)

A variant with 1-D convolutions over depth before this MLP (model_ae.column: conv, runs ae128c and ae128v) was tested and rejected: it learned more slowly in every depth band and cost 28 % more time. The option remains in the code with mlp as the default.

3. Message rounds: cells talk to their neighbours (msg, rel_msg, upd)

What happens. Each cell collects a message from each of its 5 or 6 neighbours, averages them, and revises its token. Doing this twice lets every cell know what the ocean looks like two cells away, about 90 km, before the detail is pooled away in the next step.

for round in 1, 2:
    nb  = W·h[neighbours] + R·rel          # (465,044, 6, 128): the 5–6 cells on cell
    agg = GELU(mean over valid neighbours of nb)
    h   = h + MLP([h | agg])               # (465,044, 2·128 → 128 → 128)

Two rounds of message passing on cells_on_cell (the mesh adjacency, 2.8 million directed edges). Each neighbour's token is passed through a linear map W and offset by a linear function of the neighbour's displacement in the cell's local east–north frame, in units of one cell spacing (REL_CELL = 0.005 rad ≈ 30 km), so the cell knows which message came from the east and which from the north. The messages are deliberately linear per edge: a hidden activation on each of the 2.8 M edges × 128 channels would not fit next to the rest of the state on a 40 GB GPU. The nonlinearity sits after the mean.

4. Pooling: regions become latent nodes (pool, members)

What happens. The 465,044 cell tokens are merged into 29,000 node tokens, one per patch of about 16 cells. This is where most of the compression happens: 16 cells become one row.

g = MLP([h[latent cell] | mean of members | max of members])   # (29,000, 3·128 → 256 → 256)

The latent graph (n3d/graph.py) picks 29,000 latent cells by farthest-point sampling (median spacing 124 km) and assigns every ocean cell to the nearest latent cell by a breadth-first search through water, so a region never crosses land: a cell on the Pacific side of Panama belongs to a Pacific node even if an Atlantic node is nearer as the crow flies. Regions hold 2 to 33 cells (median 16). They are stored as a padded (29,000 × 33) table; the pooling is a gather, a masked mean and a masked max over that table. The mean carries the typical state of the patch, the max carries its extremes (a front or an eddy that the mean would blur), and the latent cell's own token is kept as a third input so the node knows its centre as well as its region.

5. Latent attention: nodes see their rings (enc_blocks, NeighbourAttention)

What happens. Each node looks at its 32 nearest nodes and decides, separately for each of four "heads", how much to take from each of them; then revises its token. Four such blocks. After them a node reflects the ocean about 1,000 km around it, so basin-scale structure (a gyre, a meander, a warm pool) is in the code and not only local detail.

for block in 1..4:
    g = g + Attn(LN(g) → its 32 table tokens, bias = MLP(rel) per head)
    g = g + MLP(LN(g))                      # 256 → 1024 → 256
z = Linear(LN(g))                            # (29,000, 128)

"Attention" here means: the node computes a weight for each of its 32 candidates from the match between its own token and theirs, adds a bias that depends on where the candidate lies relative to the node (REL_LAT = 0.02 rad ≈ 125 km), turns the weights into fractions that sum to one, and takes that weighted mix of the candidates' tokens. The table of 32 is fixed per node: itself, its scale-1 neighbours, then scale-2 neighbours, padded with −1 where the ring is short. The NeighbourAttention class does this with one einsum over (node, key, head, channel); there is never a 29,000 × 29,000 matrix. The final LayerNorm and linear map produce the 128 latent channels.

In training a small Gaussian noise (z_noise 0.1) is added to z and a tiny L2 penalty (z_l2 10^-3) keeps its scale of order one; both exist for the prior's sake, so that the latent space is smooth (a slightly wrong code still decodes to a sensible ocean) and not arbitrarily scaled. The trained latents have a channel standard deviation of 2 to 3.

6. Query decoder: any point, any depth, any variable (decode)

What happens. A question (where, how deep, which variable) is turned into a token, that token reads the 10 code rows around the place, and a small head writes the answer: a value and how sure it is.

q   = MLP([Fourier(xyz) 36 | Fourier(log depth) 12 | static of the cell 8 | variable embedding 32])   # 88 → 128
tok = Linear([z | static of the latent cell])                                                          # 136 → 128
for block in 1, 2:
    q = q + Attn(LN(q) → the 10 latent tokens of the query's owner node, bias = MLP(rel))
    q = q + MLP(LN(q))
value, log-variance = Head(q)                                                                          # 128 → 128 → 2

A query is a tuple (unit vector on the sphere, depth in metres, cell, variable id). Its position enters as sine and cosine features at 6 frequencies each for x, y, z and for the log-scaled depth (tokens.query_features, 56 geometry numbers). Sines and cosines are used because a network finds it much easier to tell two nearby positions apart from a set of waves of different wavelengths than from three raw coordinates. The question also carries the static features of the cell it falls in and a learned 32-dimensional tag for which of the 11 variables is wanted. It then attends, in two cross-attention blocks, to the 10 latent tokens around it: the owner node of its cell and that node's scale-1 neighbours (k_dec 10, table_dec), with the same displacement bias as the encoder. The head returns a value and a log-variance. Ice variables are decoded at depth 0; ice velocity at vertex positions, with any adjacent cell as the query's cell.

Because the decoder is a function and not a grid, the same weights serve three jobs: reconstructing the full state (25 million queries, in chunks of 500,000), evaluating the model at observation positions during guidance, and comparing samples of the prior with data. The decoder's capacity is small on purpose (0.46 M of the 4.07 M parameters): most of the work is in the latent, and a small decoder keeps the full-state decode at about 7 s per variable.

7. Training queries: what is decoded each step (sample_queries)

What happens. Decoding all 104.7 million targets at every training step would be far too slow, so each step asks a random sample of 390,000 questions per state and checks the answers against the stored values.

The targets come straight from the stored state, so no interpolation is involved: the query depth is the mid-depth of the level the point sits in. Uniform sampling over wet points favours the upper ocean (the top 200 m hold about a third of all wet points and most of the variance); level balancing was the fix for the abyss: without it, the model saw a 5,000 m point once for every twenty surface points and learned the abyss poorly.

8. Loss: Gaussian likelihood plus a vertical-gradient term (loss)

What happens. The model is scored on each answer by how far the value is from the truth relative to its own error bar, plus a penalty for claiming a large error bar. This rewards answers that are close and error bars that are neither boastful nor timid. A second term scores the difference between vertically adjacent answers, which is what makes the model get the layering of the ocean right.

nll   = ½ · ((y − μ)² · exp(−logvar) + logvar)          mean over all queries, logvar clamped to [−7, 5]
vgrad = mean over pairs of ((μ_below − μ) − (y_below − y))²
loss  = nll + vgrad_weight · vgrad + z_l2 · mean(z²)

The first line is the negative log-likelihood of a Gaussian with the predicted mean and variance. It makes the decoder honest about its uncertainty: the variance cannot shrink without the error shrinking too, and the 1σ coverage in evaluation is 0.74–0.78 (0.68 would be perfect; the model is slightly wide). The vertical-gradient term (ae128g, weight 5) decodes the two points of each vertical pair and penalises the error of their difference. Since the mean-square of the stored vertical difference is only 0.11, a model that got each point right to 0.3 could still have the stratification wrong; this term asks for the right stratification directly. It is what removed the near-surface dip of the R² profile (the top 50 m, where the mixed layer sits on the thermocline).

forward = loss at the end of the file is what lets DistributedDataParallel wrap the module: each GPU calls forward on its own state and the gradients are averaged across the four GPUs of a node.

Why these numbers

Parameters and cost

One training step on one state takes 0.32 s on an A100 (forward 0.1 s, backward 0.2 s) and about 20–25 GB. Four states per step on four GPUs give 10–11.5 states/s; 100,000 steps are 10 hours, three or four interactive allocations. The first version of the model took 4.2 s per state: 94 % of it was the sort-based backward of fancy indexing (x[idx]) over 2.8 M cell edges and 3.9 M decoder keys. Replacing those gathers by F.embedding (gather at the top of the file), whose backward is a segment reduction, and replacing scatter pooling by the padded member table, brought it to 0.32 s.

Results: three runs, one winner

All runs: 100,000 steps, 4 states per step, learning rate 2·10^-4 with 1,000 warm-up steps, EMA weights evaluated. Full-state R² is computed by decoding every wet point of 8 held-out states (val = years 1990–94 starts, test = 1996–99), in stored units, so every level counts by its own variance.

How to read R². It is the fraction of the variance of the true field that the reconstruction explains: 1.0 would be a perfect copy, 0 would be no better than predicting the climatological mean everywhere. 0.978 for temperature means the reconstruction's errors have 2.2 % of the variance of the real anomalies, or an RMS error of about 15 % of the RMS anomaly. The gate, 0.95, corresponds to an RMS error of 22 %.

Gate G2 (held-out R² ≥ 0.95) is met by T, S, U, SSH and all ice fields on both held-out periods. V is at 0.936, held back only by the abyss (0.85 below 2 km). Test values sit close to validation values, so nothing is memorised from the training years. Precision is not the limit: decoding ae128g in fp32 instead of bf16 changes no R² by more than 0.0003.

Engineering notes worth knowing

Glossary

Source files: n3d/autoencoder.py, n3d/tokens.py, n3d/graph.py, train_ae.py, evaluate_ae.py; results in PLAN.md §6 and the walkthrough. Schematic drawn by docs/make_ae_simple.py.