Functional autoencoder

Functions as point clouds

The object the FAE works on is one discretized scalar function, a pair (x, u) of coordinates x ∈ ℝ^N×D in D spatial dimensions and values u ∈ ℝ^N, one scalar per point. The number of points and their locations are arbitrary and may differ from one function to the next. The encoder ℰ reduces such a pair to an embedding ζ, and the decoder 𝒟 returns the value 𝒟(ζ, y) of the represented function at any query coordinate y of the domain. Training minimizes the mismatch between the reconstructed and the observed values, so the task amounts to a near-lossless compression of a discretized scalar function.

Restricting the FAE to scalar functions keeps it independent of the number of channels. A velocity field, a multi-variable state or the three color channels of a photograph are reduced one channel at a time, and the representation is blind to the physical meaning of what it encodes. The input and the output coordinates are detached as well. The decoder is a conditional neural field attached to an embedding rather than a list of values tied to a mesh, so the cloud on which a field is observed and the cloud on which it is reconstructed need share neither their size nor their placement.

The pretext objective is the reconstruction of the uncorrupted function. The masked objectives common in vision instead install a bias toward the statistics of one data distribution, whereas an arbitrary function carries no semantics to infer, and a downstream operator requires a faithful and invertible compression.

Both networks work on one partition at a time, after it has been carried into its normalized frame. The encoder is applied to every partition of a field independently and returns one embedding per partition, and the decoder reads every query point from the embeddings of the partitions that contain it.

Two trained instances of this architecture appear throughout the paper and this site. The analysis backbone has a latent dimension of 128; every analysis of the paper and every figure and demo of this site uses it. The operator-learning backbone has a latent dimension of 8; every operator of the paper is trained on top of it, since its small embeddings make the operators far cheaper to train, at the price of more partitions for the same accuracy. The last section lists the two side by side.

The encoder

The encoder and the decoder are new architectures designed around a single partition. The encoder lifts every point of a partition, its normalized coordinate and value, by a feed-forward network, and aggregates the lifted points by a mean into an initial embedding. Every unit then reads the points again together with the current embedding, lifts the pairs, aggregates them by a mean and folds the result into the embedding. Aggregating by a mean leaves the encoder indifferent to the number and the order of the points, so one trained model accepts clouds of any size, and every unit sees the points themselves besides the summary of the previous unit.

Updates on the hypersphere

After every unit an RMSNorm returns the morphologic embedding to a fixed magnitude, the radius √d_ζ for a latent dimension d_ζ. Every embedding therefore lies on a hypersphere and represents a direction. Each unit refines the embedding through an update that is orthogonalized against the current embedding and applied under a learned scale shared by all units of the encoder. The component of an update along the embedding would only rescale it and be undone by the normalization, so its removal eliminates both a redundancy and a source of training instability.

The construction has two consequences for what follows. Every update keeps the new embedding within 90 degrees of the old one, so that a deep stack of updates rotates the embedding gradually instead of overwriting it. The latent space the frozen backbone provides to an operator is a sphere on which distances are angles, examined in the interpretability article.

Depth through weight sharing

Depth is partly decoupled from the parameter count. Each component holds a few distinct units and deepens its forward pass by re-applying them with shared weights, every unit in place and the whole sequence of units as a block, so the parameters applied in a forward pass far outnumber those stored. The encoder of the analysis backbone has a depth of 128 and its decoder a depth of 16. In the decoder the total depth sets a fixed residual scale of one over its square root, while the encoder learns one update strength shared by all its applications. The analysis backbone stores 4.83 million parameters and the operator-learning backbone 0.10 million.

The decoder as a neural field

The decoder evaluates every query point independently. For a query point and a partition that contains it, it reads the embedding of the partition together with the coordinate of the query in the frame of that partition, passes them through its units, and maps the output back through the value frame. A point that lies in several partitions averages their estimates. Nothing is computed across query points, and the value at a point depends on its own coordinate and on the embeddings of its partitions alone.

Two properties follow from this structure. A decoded field can be evaluated anywhere, at the encoded points, on a finer grid, or in a region where nothing was observed, and its spatial gradient is available exactly. Since the value at a point depends on no other point, the Jacobian of the decoded values with respect to the query coordinates is diagonal, and one forward-mode pass per coordinate axis, with the same unit tangent at every point, returns the derivative at every point without finite differences or a separate pass per point. The value frame lies inside the differentiated map, so the gradient comes out in physical units.

Because the decoder is a conditional neural field, the clouds on which a field is encoded and decoded are independent. An image can be encoded at one resolution and decoded on a finer grid, and every pixel of an image can be decoded when only a part of it was encoded. The memory of a decode is also independent of the output size, since the queries can be answered in blocks of any size with the result of a single pass.

Training on synthetic functions

The FAE is pre-trained without access to any PDE solution. Every batch is synthesized on the fly on the accelerator from analytic distributions on the unit disk, with no dataset on disk and no solver in the loop. Both backbones are trained on radial sines, sums of anisotropic elliptical wavefronts with random frequencies, drawn anew at every step.

Generation requires neither storage nor an input pipeline, so pre-training is limited by compute rather than by data. Since no function is presented twice, the model fits a distribution rather than a finite sample, and the transfer cannot stem from memorized solver output, as none is seen. As the distribution is fully controlled, being in or out of distribution becomes a statement about effective frequency rather than physical type.

The two trained backbones

Both backbones are trained on radial sines and differ in their size, as the table lists.