Article

Inference perks

Inference modes of a segregated functional autoencoder

Daedalus keeps its encoder and decoder apart from the processor that learns an operator. This separation makes caching possible in training, and it also enables two modes of inference. A rollout can remain in the latent space for any number of steps, and a point cloud far too large for the device can be encoded and decoded one chunk at a time. This article describes both and their measured cost.

Based on Section 4 and Appendix D.4 of the paper

Separation of encoding, processing and decoding

An operator in Daedalus composes three maps. The encoder turns every partition of a field into one embedding, the processor maps the embeddings of the input to those of the output, and the decoder turns embeddings back into values at any points. The encoder and the decoder are pre-trained once, on synthetic functions, and stay frozen for every task. Only the processor learns the physics of a task.

Encoderone partition in, one embedding out, on its own
Processorembeddings to embeddings; the only part that couples partitions
Decoderone point out, from its own coordinate and its partitions' embeddings

Two properties of this split underlie both modes. First, the processor's output is itself a set of embeddings of the next state, so it can be fed back to the processor directly. Second, the encoder and the decoder act on every partition independently, so they do not need the whole cloud at once. The first property yields latent rollouts, and the second chunked inference.

Latent rollouts

A time-dependent operator advances a state by one step, and a forecast applies it again and again, which is called a rollout. A model that works on fields has to go back to the field between steps: decode the prediction at every point, partition the new field, and encode it again before the next step. Decoding and re-encoding touches every point of the cloud at every step, and on a dense cloud it costs far more than the step itself.

When the processor predicts the next state directly, its output is a state embedding and needs no decoding or re-encoding. A latent rollout encodes the initial state once, lets the processor step in the latent space for as many steps as required, and decodes only the final state, or any state that is wanted.

Steps
Schematic of two rollouts that forecast the same state after the same number of steps. Above, the field is decoded and re-encoded around every step, so every step pays for the encoder and the decoder. Below, the state is encoded once, the processor iterates in the latent space, and only the final state is decoded.

Each step of the latent loop costs only the processor, so the single encoding and decoding are amortized over the length of the rollout. Its memory does not grow with the length either, since only the current state is kept. The mode requires the processor to predict the state itself rather than its change, because with a nonlinear decoder the sum of two embeddings does not decode to the sum of two fields. With residual or derivative steps a rollout returns to the field at every step.

Measured cost of latent rollouts

The paper's benchmark times rollouts of 1 to 4096 steps on one RTX 4090, on a coarse problem (16,384 points, 1024 partitions) and a dense one (262,144 points, 4096 partitions). The chart compares the two rollouts of Daedalus, the latent rollout and the rollout that decodes and re-encodes at every step.

–
faster at the longest rollout
Latency per sample of a whole rollout against its length, on log axes. A line ends where its setting runs out of memory.

A latent rollout costs nearly the same at short lengths, where the single encoding and decoding dominate, and grows only once the processor steps accumulate. At 4096 steps on the dense problem it is 6.6 times faster than the re-encoding rollout at a batch of one, 10 times faster at a batch of four, and 8.3 times faster at a batch of 32, and on the coarse problem 3.2 to 3.8 times faster. The Operator learning article compares these rollouts with those of other architectures.

Chunked inference

The memory of inference grows linearly with the points, because the encoder and the decoder keep intermediate features for every pair of a point and a partition it belongs to. On one GPU this bounds the size of a field. The encoder and the decoder, however, act on every partition independently. Only the processor couples partitions, and it operates on embeddings, fewer than the points by the size of a partition.

Chunked inference exploits this. The points are cut into chunks of equal size along a space-filling curve, and every chunk is partitioned on its own. The chunks are encoded one after the other, the processor runs on the embeddings of all chunks together, exactly as it would without chunks, and the chunks are decoded one after the other. Only the inputs and the output are held for all points at once; every intermediate array exists for one chunk at a time.

Chunks
Schematic of one inference in chunks. The gauge on the right shows the device memory, with the inputs and the output at the bottom and, above them, the work of the chunk being encoded or decoded. The dashed outline is what the whole cloud would need at once, more than the device holds with a single chunk.

Two details matter. The assignments of points to partitions are not stored between the encoding and the decoding, since for a huge cloud they would be the largest array of all. The decoder recomputes them chunk by chunk from the stored clustering, which gives exactly the assignments of the encoding. The mode also changes no weight, so a model trained without chunks runs with them unchanged.

The accuracy cost is small. A partition never crosses a chunk boundary, and a run of a space-filling curve can have an irregular shape, with a thin tail or a detached piece, where a partition may be longer and thinner than those the decoder was trained on. The more partitions a chunk holds, the rarer such partitions are. With the operator-learning backbone, on a field of about a million points in 4096 partitions, the median relative L2 error is 0.208% in one piece and 0.215% in 64 chunks.

Measured cost of chunked inference

The benchmark grows the problem with the partitions, at a fixed number of points per partition, and runs the full operator on one RTX 4090 with 24 GB until it runs out of memory. Each bar shows the largest field one GPU maps with a given number of chunks.

The largest number of points per function that fits in one inference, per number of chunks, at the chosen number of points per partition (PPP). A cross marks a bar whose next size ran out of memory.
537 M
Points at once with 64 chunks, 2.1 million partitions, in 86 s within 15.4 GB
16×
larger fields than in one piece, which stops at 33.6 million points
17×
less memory at 33.6 million points and 64 PPP, 1.1 GB in 64 chunks against 18.8 GB
0
weights changed; the mode is enabled at inference
Peak memory or latency per sample against the points per function, one line per number of chunks, at a batch of one, on log axes. PPP is the number of points per partition.

For small fields chunking adds latency, since the chunks run sequentially; at 65 thousand points and 64 points per partition, 64 chunks take 142 ms where one piece takes 7 ms. For large fields the relation reverses. Clustering many partitions at once costs more than clustering them in groups, so at 64 points per partition 64 chunks are faster than one piece from about 4 million points, and at 33.6 million points they take 3.8 s instead of 24.7 s. With 16 points per partition the same points form many more partitions, and chunking acts mainly as a speed-up. At 33.6 million points an inference takes 5.9 s in 64 chunks against 1229 s in one piece. Beyond 64 chunks the memory decreases little further, since the arrays kept for every point then dominate.

Summary

  • Neither mode requires training. Both follow from the structure of the model. A latent rollout feeds the processor's output back to the processor, and chunked inference reorders the work of the frozen encoder and decoder.
  • Latent rollouts reduce the cost of long forecasts. A step costs the processor alone, and the field is decoded only where it is needed.
  • Chunks raise the size limit of a GPU. One consumer GPU maps half a billion points per function in one inference.