🤖 machine learning

DiffATS: Diffusion in Aligned Tensor Space

This paper introduces DiffATS, a generative framework that trains diffusion models directly on compact, data-adaptive tensor primitives derived from Tucker decomposition and orthogonal Procrustes alignment, enabling efficient high-resolution spatiotemporal generation without relying on pretrained deep autoencoders.

Original authors: Jinhua Lyu, Tianmin Yu, Brian Kim, Lizhuo Zhou, Chanwook Park, Naichen Shi

Published 2026-05-12

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Jinhua Lyu, Tianmin Yu, Brian Kim, Lizhuo Zhou, Chanwook Park, Naichen Shi

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: Trying to Paint a Masterpiece with a Sledgehammer

Imagine you want to teach a robot to paint a complex, high-resolution video of a stormy ocean or a swirling galaxy. The data for these scenes is massive—like a giant 3D block of numbers (a "tensor").

Standard AI models try to learn this by looking at every single number in that giant block. It's like trying to learn how to paint a masterpiece by staring at every single grain of sand on a beach. It's slow, expensive, and requires a supercomputer the size of a house just to do the math.

To fix this, other methods try to "compress" the data first. They use a pre-trained "squeezing machine" (an autoencoder) to turn the giant block into a smaller, simpler version, teach the AI on that, and then un-squeeze it later. But here's the catch: for scientific data (like weather patterns or fluid dynamics), we often don't have a pre-trained squeezing machine. Building a new one for every specific problem is expensive and difficult.

The Paper's Solution: A Smart, Custom Filing System

The authors propose DiffATS. Instead of using a pre-trained "squeezing machine," they build a custom, mathematical filing system specifically for the data at hand. They call this system Aligned Tensor Primitives.

Think of it like organizing a messy library.

1. The "Tucker" Idea: Breaking it Down

The paper starts with a mathematical trick called Tucker Decomposition. Imagine you have a giant, complex 3D puzzle. Instead of keeping the whole thing, you break it down into:

A small "core" piece that holds the main shape.
Several "factor" sheets that tell you how to stretch and rotate that core to rebuild the puzzle.

This is great because it uses way fewer numbers than the original puzzle. However, there's a huge problem: The puzzle is ambiguous.

2. The "Ambiguity" Problem: The Rotating Puzzle

Imagine you have a square puzzle piece. You can rotate it 90 degrees, and it still fits the same spot. Or you can flip it. Mathematically, there are infinite ways to rotate these "factor sheets" while still building the exact same picture.

If you try to teach an AI to generate these puzzles, the AI gets confused. It sees the same picture represented in a million different, rotated ways. It's like trying to teach a child to recognize a dog, but sometimes the dog is shown standing, sometimes upside down, sometimes spinning. The AI wastes energy trying to learn all the rotations instead of learning what a dog actually looks like.

3. The "OP Alignment" Fix: The Anchor

This is the paper's secret sauce. They use a technique called Orthogonal Procrustes (OP) Alignment.

Imagine you have a "Master Key" (an Anchor) that represents the most typical orientation of your puzzle pieces. Before you show any puzzle piece to the AI, you force it to rotate until it perfectly matches the Master Key.

Before: The AI sees a puzzle piece pointing North, then East, then South.
After: The AI always sees the puzzle piece pointing North because you aligned it to the Master Key first.

By doing this, the AI no longer has to learn how to rotate things. It only learns the actual content. This makes the learning process much faster and more accurate.

How DiffATS Works (The Pipeline)

Pick an Anchor: The system looks at a bunch of training data and picks one "representative" orientation (the Medoid) to be the Master Key.
Align Everything: Every single piece of data is mathematically rotated to match that Master Key.
Train the AI: The AI learns to generate these "aligned" pieces. Because they are all facing the same way, the AI learns the patterns much better.
Reconstruct: When the AI finishes generating a new piece, the system simply un-rotates it (using the math rules) to build the final, high-resolution image or video.

The Results: Smaller, Faster, Better

The paper tested this on three types of data:

Images: High-resolution faces (CelebA-HQ).
Videos: Moving digits (Moving MNIST).
Science: Simulations of fluid flow and weather (PDEs).

The Claims:

Compression: DiffATS shrinks the data by 3.9x to 210x. That's like turning a 100GB movie into a 1GB file without losing the story.
No Pre-training: Unlike other methods, it doesn't need a pre-trained "squeezing machine." It builds its own filing system on the fly.
Better Quality: In every test, DiffATS produced sharper, more accurate images and videos than the other "no-autoencoder" methods. It beat the competition even when they were using the same amount of compressed space.

The Bottom Line

DiffATS is a new way to teach AI to generate complex data. Instead of forcing the AI to learn how to rotate and flip data (which is confusing and wasteful), it forces the data to line up in a straight row first. This makes the AI's job easier, the training faster, and the final results much sharper, all without needing expensive pre-trained tools.

Technical Summary: DiffATS: Diffusion in Aligned Tensor Space

1. Problem Statement

Direct diffusion modeling of high-resolution spatiotemporal fields (e.g., videos, weather data, fluid dynamics, and partial differential equation (PDE) solutions) is computationally prohibitive. Standard diffusion models (DMs) require repeated function evaluations in the full data dimension, leading to immense training and sampling costs. For instance, training a weather forecasting model on tensors of size $84 \times 720 \times 1440$ requires days on dozens of TPUs.

To mitigate this, Latent Diffusion Models (LDMs) typically employ Variational Autoencoders (VAEs) to compress data into a lower-dimensional latent space. However, for scientific data, off-the-shelf pretrained autoencoders are often unavailable, and training task-specific autoencoders is expensive and requires sophisticated optimization. While autoencoder-free primitives based on Discrete Cosine Transforms (DCT) or fixed Tucker factorizations exist, they rely on predefined bases that may not optimally approximate the raw data's inherent structure.

The paper poses the question: Can we construct data-dependent, parameter-efficient primitives for high-resolution spatiotemporal data by exploiting their inherent low-rank structure to improve generative modeling without pretrained autoencoders?

2. Methodology

The proposed framework, Diffusion in Aligned Tensor Space (DiffATS), addresses the computational and representational challenges through a novel construction of tensor primitives based on Tucker decomposition and Orthogonal Procrustes (OP) alignment.

2.1 The Degeneracy of Natural Tucker Primitives

The authors identify that using standard Tucker decomposition factors (a core tensor and mode-wise factor matrices) directly as primitives for diffusion training is suboptimal. The Tucker decomposition is non-unique: a single tensor can be represented by infinitely many combinations of core tensors and factor matrices related by orthogonal rotations. This "gauge ambiguity" causes the primitive space to be dispersed, breaking the continuity of the mapping and complicating the learning of the score function for diffusion models.

2.2 Orthogonal Procrustes (OP) Alignment

To resolve this degeneracy, DiffATS introduces Orthogonal Procrustes (OP) alignment.

Anchor Selection: The method selects "medoid" anchor matrices from the training data. These anchors represent the most central or representative subspaces within the empirical distribution of factor matrices.
Alignment: For each data sample, the factor matrices are aligned to the corresponding anchor matrix by finding the optimal orthogonal rotation that minimizes the Frobenius norm distance. This selects a unique canonical representative from each Grassmannian equivalence class.

2.3 Grassmannian Primitives

Based on OP alignment, the paper defines two families of primitives:

Matrix Grassmannian Primitives (MGPs): For low-rank matrices (e.g., 1-D PDEs or image patches). The primitive consists of the aligned singular vectors and the transformed singular values.
Tensor Grassmannian Primitives (TGPs): For low-multilinear-rank tensors (e.g., videos or 2-D PDEs). The primitive consists of the core tensor and the aligned factor matrices for selected modes.

These primitives are data-adaptive (derived from the specific dataset's structure) and parameter-efficient, compressing the original data by factors ranging from $3.9\times$ to $210\times$ .

2.4 Theoretical Guarantees

The authors prove that under specific "anchor-overlap" conditions (where the data's subspaces have non-trivial overlap with the anchor), the mapping from the original tensor space to the primitive space is a homeomorphism. This theoretical result certifies that:

The mapping is bijective (one-to-one and onto).
The mapping is continuous, preserving the local topological structure of the data distribution.
The representation is non-degenerate, ensuring that the diffusion model learns a faithful distribution of the primitives.

2.5 The DiffATS Pipeline

Preprocessing: High-dimensional tensors are decomposed via HOOI (Higher-Order Orthogonal Iteration) or SVD.
Alignment: Factor matrices are aligned to medoid anchors via OP.
Training: A diffusion model (using a DiT backbone) is trained directly on the aligned primitives (core tensors and aligned factors).
Inference: The model generates samples in the primitive space, which are then reconstructed to the original data space via explicit multilinear reconstruction (multiplying the core by the aligned factors).

3. Key Contributions

Novel Primitives: Introduction of MGPs and TGPs as data-adaptive, parameter-efficient representations that retain the expressiveness of Tucker decompositions while resolving rotational ambiguity through anchored OP alignment.
Theoretical Foundation: Proof that the proposed primitive maps are homeomorphisms under overlapping conditions, providing the first formal homeomorphism guarantee for a low-rank primitive family used in generative modeling.
Autoencoder-Free Framework: Proposal of DiffATS, an end-to-end framework that trains and samples directly in the aligned primitive space without relying on pretrained deep compression autoencoders.
Empirical Performance: Demonstration that DiffATS outperforms existing autoencoder-free baselines (DCTDiff, SDIFT, AvgPool) and competitive neural operators (FNO) across images, videos, and PDE solutions, achieving superior generation quality at comparable or higher compression ratios.

4. Experimental Results

The authors evaluated DiffATS on six benchmark datasets:

Images: CelebA-HQ ( $1024 \times 1024$ ).
Videos: Moving MNIST.
PDEs: 1-d Burgers, 1-d Reaction-Diffusion, 2-d Burgers, and 2-d Kármán vortex street.

Key Findings:

Generation Quality: DiffATS achieved the best Fréchet Inception Distance (FID) on CelebA-HQ (23.57 vs. 32.89 for AvgPool) and the best Fréchet Video Distance (FVD) on Moving MNIST (415.42 vs. 487.82 for AvgPool).
PDE Prediction: In conditional generation tasks (predicting future trajectories given initial conditions), DiffATS significantly outperformed all baselines. For example, on the 1-d Reaction-Diffusion dataset, it reduced the relative $\ell_2$ error by approximately $9\times$ compared to the best baseline (SDIFT).
Compression: The method achieved compression ratios from $3.9\times$ to $210\times$ while maintaining high fidelity.
Ablation Studies: Experiments confirmed that using data-dependent bases with OP alignment is critical; removing alignment or using shared bases significantly degraded performance.

5. Significance and Claims

The paper claims that DiffATS offers a robust alternative to autoencoder-based latent diffusion for scientific and high-resolution spatiotemporal data where pretrained autoencoders are unavailable or impractical to train. By leveraging the inherent low-rank structure of such data and resolving the non-uniqueness of tensor factorizations through geometric alignment, DiffATS enables efficient and high-quality generative modeling.

The authors position their work as a solution to the "degeneracy" problem in standard tensor factorization for generative tasks, providing a theoretically grounded method to create stable, low-dimensional representations. They emphasize that their approach does not aim to surpass state-of-the-art image synthesis models that rely on massive pretrained autoencoders (like Stable Diffusion), but rather targets the specific domain of spatiotemporal physical fields where such resources are lacking and low-rank structure is a dominant feature.

Future directions suggested by the authors include exploring multi-anchor OP alignment and dynamic rank selection to relax the assumptions required for the homeomorphism guarantees.