Coordinate In and Value Out: Training Flow Transformers in Ambient Space

We interpret our empirical data distribution $q$ to be composed of maps $f\sim q(f)$. These maps take coordinates ${\bm{x}}$ as input to values ${\bm{y}}$ as output. For images, maps are defined from 2D pixel coordinates ${\bm{x}}\in\mathbb{R}^{2}$ to corresponding RGB values ${\bm{y}}\in\mathbb{R}^{3}$, thus $f:\mathbb{R}^{2}\rightarrow\mathbb{R}^{3}$, where each image is a different map. For 3D point clouds, $f$ can be interpreted as a deformation that maps coordinates from a fixed base configuration in 3D space to a deformation value also in 3D space, $f:\mathbb{R}^{3}\rightarrow\mathbb{R}^{3}$, as in the image case, each 3D point cloud corresponds to a different deformation map $f$. For ease of notation, we define coordinates ${\bm{x}}$ and values ${\bm{y}}$ of any given map $f$ as ${\bm{x}}_{f}$ and ${\bm{y}}_{f}$, respectively. Fig. 1(a) shows an example of such maps in the image domain.

In practice, analytical forms for these maps $f$ are unknown. In addition, different from previous approaches (Dupont et al., 2022a; Du et al., 2021a), we do not assume that parametric forms of these maps can be obtained, since that would involve a separate training stage fitting an MLP to each map (Dupont et al., 2022a; Bauer et al., 2023; Du et al., 2021a). As a result, we assume we are only given sets of corresponding coordinate and value pairs resulting from observing these maps at a particular sampling rate (e.g. at a particular resolution in the image case). Therefore, we need a to develop an end-to-end approach that can take these collections of coordinate-value sets as training data.

We consider generative models that learn to reverse a time-dependent forward process that turns data samples (i.e. maps $f$ in our case) $f\sim q(f)$ into noise $\epsilon\sim\mathcal{N}(0,{\mathbf{I}})$.

Both flow matching (Lipman et al., 2023) and stochastic interpolant (Ma et al., 2024) formulations build this forward process in Eq. 1 so that it interpolates exactly between data samples $f$ at time $t=0$ and $\epsilon$ at time $t=1$, with $t\in[0,1]$. In particular, $p_{1}(f)\sim\mathcal{N}(0,{\mathbf{I}})$ and $p_{0}(f)\approx q(f)$. In this case, the marginal probability distribution $p_{t}(f)$ of $f$ is equivalent to the distribution of the probability flow ODE with the following velocity field (Ma et al., 2024):

where the velocity field is given by the following conditional expectation,

Under this formulation, samples $f_{0}\sim p_{0}(f)$ are generated by solving the probability flow ODE in Eq. 2 backwards in time (e.g. . flowing from $t=1$ to $t=0$), where $p_{0}(f)\approx q(f)$. Note that both the flow matching (Lipman et al., 2023) and stochastic interpolant (Ma et al., 2024) formulations decouple the time-dependent process formulation from the specific choice of parameters $\alpha_{t}$ and $\sigma_{t}$, allowing for more flexibility. Throughout the presentation of our method we will assume a rectified flow (Liu et al., 2023; Lipman et al., 2023) or linear interpolant path (Ma et al., 2024) between noise and data, which define a straight path to connect data and noise: $f_{t}=(1-t)f_{0}+t\epsilon$. Note that our framework for learning flow matching models for coordinate-value sets can be used with any path definition. Compared with diffusion models (Ho et al., 2020), linear flow matching objectives result in better training stability and more modeling flexibility (Ma et al., 2024; Esser et al., 2024) which we observed in our early experiments.

We now turn to the task of formulating a flow matching training objective for data distributions of maps $f$. We recall that in practice we do not have access to an analytical or parametric form for these maps $f$, and we are only given sets of corresponding coordinate ${\bm{x}}_{f}$ and value ${\bm{y}}_{f}$ pairs resulting from observing the mapping at a particular rate. As a result, we need to formulate a training objective that can take these sets of coordinate-value as training data.

In order to achieve this, we first observe that the target velocity field ${\bm{u}}_{t}(f_{t})d_{t}$ can be decomposed across both the domain and co-domain of $f_{t}$, resulting in a point-wise velocity field ${\bm{u}}_{t}({\bm{x}}_{f_{t}},{\bm{y}}_{f_{t}})d_{t}$, defined for corresponding coordinate and value pairs of $f_{t}$. As an illustrative example in the image domain, this means that the target velocity field can be independently evaluated for any pixel coordinate ${\bm{x}}_{f_{t}}$ with corresponding value ${\bm{y}}_{f_{t}}$, so that ${\bm{u}}_{t}({\bm{x}}_{f_{t}},{\bm{y}}_{f_{t}})\in\mathbb{R}^{3}$. Note that one can always decompose target velocity fields in this way since the time-dependent forward process in Eq. 1 aggregates data and noise independently (e.g. point-wise) across the domain of $f$. Again, using the image domain as an example, the time-dependent forward process of a pixel at coordinate ${\bm{x}}_{f}$ is not dependent on other pixel positions or values.

Our goal now is to formulate a training objective to match this point-wise independent velocity field. We want our neural network ${\bm{v}}_{\theta}$ parametrizing the velocity field to be able to independently predict a velocity for any given coordinate and value pair ${\bm{x}}_{f_{t}}$ and ${\bm{y}}_{f_{t}}$. However, this point-wise independent prediction is futile without access to additional contextual conditioning information about the underlying function $f_{t}$ at time $t$. This is because even if the forward process is point-wise independent, real data exhibits strong dependencies across the domain $f$ that need to be captured by the model. For example, in the image domain, pixels are not independent from each other and natural images show strong both short and long spatial dependencies across pixels. In order to solve this, we introduce a latent variable ${\bm{z}}_{f_{t}}$ that encodes contextual information from a set of given coordinate and value pairs of $f_{t}$. This contextual latent variable allows us to formulate the learnt velocity field to be conditionally independent for coordinate-value pairs given ${\bm{z}}_{f_{t}}$. The final point-wise conditionally independent CFM loss, which we denote as CICFM loss is defined as:

where the target velocity field ${\bm{u}}_{t}({\bm{x}},{\bm{y}}|\epsilon)$ is defined as a rectified flow (Liu et al., 2023; Lipman et al., 2023) or linear interpolant path (Ma et al., 2024):

One of the core challenges of learning this type of generative models is obtaining a latent variable ${\bm{z}}_{f_{t}}$ that effectively captures intricate dependencies across the domain of the function, specially for high resolution stimuli like images. In particular, the architectural design decisions are extremely important to ensure that ${\bm{z}}_{f_{t}}$ does not become a bottleneck during training. In the following we review our proposed architecture.

We base our model on the general PerceiverIO design (Jaegle et al., 2022) which provides a flexible architecture to handle coordinate-value sets of large cardinality (i.e. large number of pixels in an image). Fig. 2 illustrates the architectural pipeline of ASFT. At a high level, our encoder network takes a set of coordinate-value pairs and encodes them to learnable latents through cross-attention. These latents are then updated through several self-attention blocks to provide the final latents ${\bm{z}}_{f_{t}}\in\mathbb{R}^{L\times D}$ . To decode the velocity field for a given coordinate-value pair we perform cross attention to ${\bm{z}}_{f_{t}}$, generating the final point-wise prediction for the velocity field ${\bm{v}}_{\theta}({\bm{x}}_{f_{t}},{\bm{y}}_{f_{t}},t|{\bm{z}}_{f_{t}})$.

The encoder of a vanilla PerceiverIO relies solely on cross-attention to the latents $z_{f_{t}}\in\mathbb{R}^{L\times D}$ to learn spatial connectivity patterns between input and output elements, which we found to introduce a strong bottleneck during training. To ameliorate this, we make two key modifications to boost the performance. Firstly, our encoder utilizes spatial aware latents where each latent is assigned a “pseudo” coordinate. Coordinate-value pairs are assigned to different latents based on their distances on coordinate space. During encoding, coordinate-value pairs interact with their assigned latents through cross-attention. In particular, the learnable latent $z_{f_{t}}$ cross-attends to input coordinate-value pairs of noisy data at a given timestep $t$. Latent vectors are spatial-aware, this means that each of the $L$ latents only attends to a set of neighboring coordinate-value pairs. Latent vectors are then updated using several self-attention blocks. These changes in the encoder allow the model to effectively utilize spatial information while also saving compute when encoding large coordinate-value sets on ambient space. In the decoder, a given coordinate-value pair cross attends to ${\bm{z}}_{f_{t}}$ as in the original PerceiverIO. However, we found that a multi-level decoding strategy, which not only cross attends to the latents in the final layer but also the latent from the intermediate self-attentions layer is helpful. In particular, a given coordinate-value pair cross attends to latents from subsequent encoder layers sequentially to progressively refine the prediction. Finally, following previous work (Peebles & Xie, 2023; Ma et al., 2024), we apply AdaLN-Zero blocks for conditioning both on timestep $t$ and class labels whenever needed (e.g. for ImageNet experiments). More architectural details can be found in App. A.

Given that ASFT is a generative model for maps we compare it with other generative models of the same type, namely approaches that operate in function spaces. Tab. 1 shows a comparison of different image domain specific as well as function space models (e.g. approaches that model infinite-dimensional signals). ASFT surpasses other generative models in function space on both FFHQ (Karras et al., 2019) and LSUN-Church (Yu et al., 2015) at resolution $256\times 256$. Compared with generative models designed specifically for images, ASFT also achieves comparable or better performance. When scaling up the model size, ASFT-L demonstrates better performance than all the baselines on FFHQ-256 and Church-256, indicating that ASFT can benefit from increasing model sizes.

We also evaluate the performance of ASFT on large scale and challenging settings, we train ASFT on ImageNet at both 128$\times$128 and 256$\times$256 resolutions. On ImageNet-128, shown in Tab. 2, ASFT achieves an FID of $2.73$, which is a a competitive performance in comparison to diffusion or flow-based generative baselines including ADM (Dhariwal & Nichol, 2021), CDM (Ho et al., 2021), and RIN (Jabri et al., 2023) which use domain-specific architectures for image generation. Besides, comparing to PolyINR (Singh et al., 2023) which also operates on function space, ASFT achieves competitive FID, while obtaining better IS, precision and recall. The experimental results demonstrate the capabilities of ASFT in generating realistic samples on large scale datasets.

We report results of ASFT for ImageNet-256 on Tab. 3. Note that ASFT is slightly outperformed by latent space models like DiT (Peebles & Xie, 2023) and SiT (Ma et al., 2024). We highlight that these baselines rely on a pre-trained VAE compressor that was trained on datasets that are much larger than ImageNet, while ASFT was trained only with ImageNet data. In addition, ASFT achieves better performance than many of the baselines trained only with ImageNet data including ADM (Dhariwal & Nichol, 2021), CDM (Ho et al., 2021) and Simple Diffusion (U-Net) (Hoogeboom et al., 2023) which all use CNN-based architectures specific for image generation. Note that this is consistent with the results show in Tab. 1, where ASFT outperforms all function space approaches. When comparing with approaches using transformers architectures we find that ASFT obtains performance comparable to RIN (Jabri et al., 2022) and HDiT (Crowson et al., 2024), with slightly worse FID and slightly better IS. However, ASFT is a domain-agnostic architecture that can be trivially applied to different data domains like 3D point clouds (see Sect. 4.3 and 4.4). For completeness, we also include a comparison with very large U-Net transformer hybrid models, Simple Diffusion (U-ViT 2B) and VDM++ (U-ViT 2B) which both use approx. $\times 2.72$ more parameters than ASFT-XL, unsurprisingly, these much bigger capacity models outperform ASFT (see App. A for a more detailed comparison including training settings). We highlight that the simplicity of implementing and training ASFT models in practice, and the trivial extension to different data domains (as shown in Sect. 4.3) are strong arguments favouring our model. Finally, comparing with e.g. PolyINR (Singh et al., 2023) which is also a function space generative model we also find comparable performance, with slight worse FID but better Precision and Recall. It is worth noting that (Singh et al., 2023) applies a pre-trained DeiT model as the discriminator (Singh et al., 2023). Whereas our ASFT makes no such assumption about the function or pre-trained models, enabling to trivially apply ASFT to other domains like 3D point clouds (see Sect. 4.3).

To demonstrate the scalability of ASFT we train models of different sizes including small (S), base (B), large (L), and extra-large (XL) on ImageNet-256. We show the performance of different model sizes using FID-50K in Fig. 3(a). We observe a clear improving trend when increasing the number of parameters as well as increasing training steps. This demonstrates that scaling the total training Gflops is important to improved generative results as in other ViT-based generative models (Peebles & Xie, 2023; Ma et al., 2024). Due to the flexibility of cross-attention decoder in ASFT, one can easily conduct random sub-sampling to reduce the number of decoded coordinate-value pairs during training which significantly saves computation. Fig. 3(b) shows how number of decoded coordinate-value pairs affects the model performance as well as Gflops in training. An image of resolution 256$\times$256 contains 65536 pixels in total which is the maximal number of coordinate-value pairs during training. As see in Fig. 3(b), a model decoding 4096 coordinate-value pairs saves more than 20% Gflops over one decoding 16384. This provides us with an effective training recipe, which saves computation by only decoding a subset of $12\%$ of the image pixels during training. Interestingly, we see a performance drop when densely decoding 16384 coordinate value pairs. We hypothesize this could be due to optimization challenges of decoding large numbers of pairs and leave further analysis for future work.

To show the domain-agnostic prowess of ASFT we also tackle 3D point cloud generation on ShapeNet (Chang et al., 2015). Note that our model does not require training separate VAEs for point clouds, tuning their corresponding hyper-parameters or designing domain specific networks. We simply adapt our architecture for the change in dimensionality of coordinate-value pairs (e.g. $f:\mathbb{R}^{2}\rightarrow\mathbb{R}^{3}$ for images to $f:\mathbb{R}^{3}\rightarrow\mathbb{R}^{3}$ for 3D point clouds.). Note that for 3D point clouds, the coordinates and values are equivalent. In this setting, we compare baselines including LION (Vahdat et al., 2022) which is a recent state-of-the-art approach that models 3D point clouds using a latent diffusion type of approach. Following Vahdat et al. (2022) we report MMD, COV and 1-NNA as metrics. To have a straightforward comparison with baselines, we train ASFT-B with to approximately match the number of parameters as LION (Vahdat et al., 2022) (110M for LION vs 108M for ASFT) on the same datasets (using per sample normalization as in Tab. 17 in Vahdat et al. (2022)). We show results for category specific models and for an unconditional model jointly trained on 55 ShapeNet categories in Tab. 4. ASFT-B obtains strong generation results on ShapeNet despite being a domain agnostic approach and outperforms LION in most datasets and metrics. Note that ASFT-B has comparable number of parameters and the same inference settings than LION so this is fair comparison. Finally, we also report results for a larger model ASFT-L (with $\times 2$ the parameter count as LION) to investigate how ASFT improves as with increasing model size. We observe that with increasing model size, ASFT typically achieves better performance than the base version. This further demonstrates scalability of our model on ambient space of different data domains.

We also experiment ASFT on large-scale 3D point generation conditioned on images. To this end, we train an image-to-point-cloud generation model on Objaverse (Deitke et al., 2023), which contains 800k 3D objects of great variety. In particular, conditional information (i.e., an image) is integrated to our model through cross-attention. For each object in Objaverse, we sample point cloud with $16k$ points. To get images for conditioning, each object is rendered with 40 degrees field of view, $448\times 448$ resolution, at 3.5 units on the opposite sides of $x$ and $z$ axes looking at the origin. We extract features via DINOv2 (Oquab et al., 2023) which is concatenated with Plucker ray embedding (Plucker, 2018) of each patch in DINOv2 feature. In each block, the learnable latent vector $z_{f_{t}}$ cross attends to image feature.

Tab. 5 lists the performance of ASFT in comparison of recent baselines on Objaverse. We report ULIP-I (Xue et al., 2024) and P-FID (Nichol et al., 2022) following CLAY (Zhang et al., 2024). PointNet++ (Qi et al., 2017a; b; Nichol et al., 2022) is employed to evaluate P-FID. ULIP-I is an analogy to CLIP for text-to-image generation. ULIP-I is measured as the cosine similarity between point-cloud features from ULIP-2 model (Xue et al., 2024) and image features from CLIP model (Radford et al., 2021). Numbers of baseline models are directly borrowed from CLAY (Zhang et al., 2024). We calculate the metrics of our ASFT on 10k sampled point clouds. In our case, P-FID is measured on point clouds with 4096 points following Shape-E (Jun & Nichol, 2023) while ULIP-I is measured on point clouds with 10k points following ULIP-2 (Xue et al., 2024). Note that since CLAY (Zhang et al., 2024) is not open-source, we do not have the access to the exact evaluation setting or the conditional images rendered from Objaverse. But all evaluation settings of ASFT are provided for reproduction purpose. As shown in Tab. 5, our ASFT achieves strong performance on large-scaled image-conditioned 3D generative tasks. Compared with CLAY (Zhang et al., 2024), which is a 2-stage latent diffusion model, ASFT demonstrates very strong performance on both ULIP-I and P-FID. Fig. 9 and 10 show additional examples of sampled point clouds and corresponding conditional images. As shown, ASFT learns to generate 3D objects with rich details that match the conditional images ultimately being able to generate a continuous surface.

An interesting property of ASFT is that it decodes each coordinate-value pair independently, allowing resolution to change during inference. At inference the user can define as many coordinate-value pairs as desired where the initial value of each pair at $t=1$ is drawn from a Gaussian distribution. We show qualitative results of resolution agnostic generation for both images and point clouds. Fig. 4(a) show images sampled at resolution $512^{2}$, $1024^{2}$, and $2048^{2}$ (together with their $256^{2}$ resolution counterparts generated from the same seed) from ASFT trained on ImageNet-256. Even though the model has not been trained with any samples at higher resolutions, it can still generate realistic images with high-frequency details. Fig. 4(b) shows point cloud with $32k$ and $128k$ points from ASFT trained on Objaverse with only $16k$ points points per sample (we visualize the generated $16k$ point could generated from the same seed). Similarly, ASFT generates dense and realistic point cloud in 3D without actually being trained on such high density points. These results show that ASFT is not trivially overfitting to the training set of points but rather learning a continuous density field in 3D space from which an infinite number of points could be sampled. Generally speaking, this also provides the potential to efficiently train flow matching generative models without the need to use large amounts of expensive high resolution data, which can be hard to collect in data domains other than images.