SPACE (Subspace Partitioning for Accessible Controlled Encodings)

spacelearn is a utility library built on top of PyTorch that provides tools for training deep-learning models in the Latent-SPACE regime. The SPACE method helps create structured, interpretable, and easily decodable embedding spaces by explicitly guiding a model's internal latent representation using side information available during training.

Rather than treating the latent space as an unconstrained representation, SPACE partitions it into dedicated subspaces that encode specific physical quantities, intermediate computations, or other known properties of the data. This allows these quantities to be decoded directly from the latent representation while still preserving unconstrained latent capacity for the model to learn additional features.

Documentation

The Method

1. Embedding Quantities

Since SPACE is used to embed specific information into the latent space, the first step is defining the target quantities. These quantities may have any number of dimensions; however, the method generally works best when targets contain between roughly 10 and a few thousand elements.

Soft guidelines for target quantities:

• Quantities should contain enough information that they cannot easily be represented by only a few latent dimensions.

  • Bad example: a 2D rotation matrix

[ \mathcal{R}(\theta)= \left( \begin{matrix} \cos(\theta) & -\sin(\theta) \ \sin(\theta) & \cos(\theta) \end{matrix} \right) ]

• Good example: the spatial gradient of a 64 × 64 image

[ G_x= \left( \begin{matrix} -1 & 0 & 1 \ -2 & 0 & 2 \ -1 & 0 & 1 \end{matrix} \right), \qquad G_y= \left( \begin{matrix} -1 & -2 & -1 \ 0 & 0 & 0 \ 1 & 2 & 1 \end{matrix} \right) ]

[ G_{A,x}=G_xA, \qquad G_{A,y}=G_yA ]

• Quantities with only a few elements can often be encoded directly into a small number of latent dimensions instead of first projecting them into learned subspaces.
• Quantities should exhibit predictable variance across samples or be analyzed using relatively high variance-retention thresholds. For quantities with highly variable sample-to-sample structure, it is often beneficial to estimate subspaces and decoders using large and diverse datasets.
• Since SPACE relies heavily on singular value decompositions and eigendecompositions, the underlying data distributions should be reasonably well-conditioned.

2. Setup / Hyperparameter Tuning

Once the target quantities have been defined, SPACE provides utilities to estimate all major SPACE-specific hyperparameters:

• optimal_bins() / solve_dims() for determining the pooling resolution n
• k_per_q() / solve_dims() for determining subspace dimensions k
• minimal_latent_dim() for determining the latent dimension D

Step 1: Choose Pooling Resolution

Use optimal_bins() (or solve_dims()) to estimate the minimum pooling resolution required to retain a desired fraction of spatial variance.

You must choose a target variance-retention threshold for the pooled quantities. Depending on the structure of the quantity, this threshold may have a dramatic or only minor effect on the resulting pooling resolution.

Step 2: Choose Subspace Dimensions

After determining the pooling resolution, use k_per_q() (or solve_dims()) to estimate the minimum subspace dimension required for each quantity.

The resulting value of k depends on:

• The variance-retention threshold chosen for the subspace
• The pooling resolution n
• The structure and complexity of the quantity itself

Relationship between pooling resolution (n) and subspace dimensionality (k) for 2D thermodynamic quantities (original quantity shape: 64 × 64).

Step 3: Choose the Latent Dimension

Once n and k are known, use minimal_latent_dim() to estimate the required latent dimension D.

This computation reserves enough latent capacity for all learned subspaces while optionally leaving a configurable fraction of the latent space unconstrained.

We recommend reserving free latent capacity because:

• Some quantities require additional latent freedom to find stable subspace geometries
• Joint training setups often benefit from latent dimensions that are not explicitly assigned to any known quantity
• Models frequently learn useful hidden representations that are not captured by the selected physical targets

3. Training

Step 1: Initialize Subspaces and Decoders

Before beginning the main training loop, initialize the subspaces and decoders.

We recommend collecting a dataset similar in size to the one used when estimating n and k.

Subspaces and decoders can then be initialized using:

WVA = solve_subspaces(Z, pooled_targets, k=k)

or individually through:

• principal_directions()
• subspaces()
• fit_decoders()

Step 2: Train Using the SPACE Objective

Add the SPACE objective to your training loss in the same way as any other regularization term.

The simplest approach is to use:

loss = combined_loss(
    Z,
    W,
    A,
    pooled_targets,
    W_prev
)

Internally this combines:

• alignment_loss()
• disentanglement_loss()
• isotropy_loss()
• stability_loss()

Step 3: Periodically Recompute Subspaces

We recommend periodically recomputing:

• W (subspace projectors)
• A (subspace decoders)

using solve_subspaces().

Experiments have shown that an exponential decay schedule often produces the best subspace geometries. In practice, this means recomputing relatively frequently during the early stages of training and less frequently later on.

We generally do not recommend recomputing subspaces within an epoch unless extremely large datasets are being used.

Architectures and Setups That Work Well

1. Separate Encoder / Decoder U-Net

Our initial experiments used a U-Net-style architecture in which the encoder (including the bottleneck) was trained under the SPACE regime, while the decoder was trained separately on the final inference objective.

This setup showed promising results for physics-guided inference tasks where conventional U-Nets often struggle to learn physically meaningful latent representations.

2. Additional Architectures

This section is currently a work in progress and may be published separately in the future.