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

1. Question about whether ambient-space (e.g., pixel-space, point-space) learning is a promising direction.

   • Admittedly, latent generative modeling with two-stage training paradigm has been gaining popularity recently. However, we want to point out that this does not diminish the value of exploring generative models in ambient space. As a matter of fact, ambient/pixel space models do surpass latent space models on ImageNet-256, Tab. 3 (eg. VDM++ U-ViT vs DiT), there’s no empirical reason to believe in the superiority of learning in latent space instead of pixel space.
   • We kindly disagree that building domain-specific data compressors for each data type is a non issue. To design different VAEs for different data domains, efforts are required to curate architectures, training recipes and datasets. The efforts put on building a compressor for one data domain do not usually transfer to other domains.
   • Finally, as the reviewer pointed out “[... Generative models evolve so fast. I think making it end-to-end is not impossible ...]” Our contribution in this submission is exactly exploring an approach to build an end-to-end generative framework in a data-agnostic way. We aim to build a simple and flexible generative model that allows to tackle multiple data domains in a single training stage.

2. Building generative models with coordinate-value pairs may essentially restrict its application scenarios and conditional generation capabilities.

   • We want to kindly point out that the coordinate-value pairs are not constraint for broader applications. On the contrary, modeling data as coordinate-value pairs actually enables ASFT to be directly applied in different data domains. ASFT actually models an explicit neural field based on coordinate-value pairs. Since all outputs are predicted conditionally independent with each other given the learnable latent $z_{f_t}$. This means that given the latent $z_{f_t}$ we can continuously query the representation like implicit fields do. As an illustrative example, we can generate infinitely dense pointclouds that represent the continuous surface in 3D space (see Fig. 4 in the paper where we sample pointclouds with 100K points, also see Fig. 10 where sample pointclouds with 128k points from models trained on Objaverse). **ASFT does not suffer from generating implicit fields, as a matter of fact, it’s designed to predict explicit neural fields that can be continously evaluated. **
   • We benchmark sparse 3D point cloud generation mainly because (1) it’s unstructured data unlike images which provide a testbed for more diverse data domains and (2) point cloud generation has well-established benchmarks to compare our results with existing work.

3. Besides, the proposed method is only implemented with class label conditioned generation, which is quite out-of-date. Its conditional generation capabilities (e.g., text-guided, image-guided) are questionable.

   • We note that class-conditioned image generation on ImageNet is a widely used benchmark for generative modeling of images and it is considered the standard benchmark for evaluation. Therefore, we include benchmark our performance on ImageNet128 and ImageNet256. ASFT can directly integrate other conditional information (as latent diffusion models do). In light of your valuable suggestion, we are training an image-to-point-cloud generation model on Objaverse to showcase the conditional generation capabilities. In particular, the conditioning (i.e., image) is integrated into ASFT through cross-attention. In each block, the latent vector $z_{f_t}$ cross attends to image features from DINOv2. During training, the image conditioning is dropped randomly with 10% probability. Therefore, our model can also benefit from popular classifier-free guidance (CFG) to enhance the match between samples and conditions.

1. Q: The novelty of the proposed method in this paper is questionable, as it appears to be relatively straightforward.

   • The novelty of our work is to build a domain-agnostic generative model on coordinate-value maps (sometimes referred to as “fields”) that can be efficiently trained in large scale settings (eg. ImageNet-256). **Note that up to date, ASFT is the only generative model in function space achieving these results. **
   • Previous domain-agnostic generative models of fields/maps, like Functa [1], GEM [2] and DPF [3], have investigated generative models in field/function space. However, these works only tackle low-dimensional problems, like image generation of resolution 32 or 64 pixels (on smaller datasets like CelebA or LSUN-Church). In our work, we show that domain-agnostic generative models on fields/maps can achieve good performance on large datasets like ImageNet-256, which previous models fail to achieve.
   • We want to highlight that we build ASFT with simplicity in mind and aim to build a unified generative model that can be seamlessly applied to different data domains. We are pleased that the reviewer believes that our method is straightforward, we have put significant effort to present and discuss our methodology in a way thats easy to understand.

2. Q: The experimental evaluation could be strengthened by employing a more diverse range of image and point cloud datasets to demonstrate the generalizability of the proposed approach.

   • We evaluated our model on 3 image datasets (FFHQ, LSUN Church and ImageNet), as well, as ShapeNet for 3D shape generation. We compare with more than 20 different baselines across all these datasets and we use all the metrics reported in previous approaches. We believe this already represents a comprehensive comparison. In addition, to demonstrate the wide applicability of ASFT, we have added results on point cloud generation of Objaverse. We train an image-to-point-cloud ASFT model on Objaverse, containing more than 800k 3D objects of wide variety. Compared with SOTA 3D generative models like CLAY, our ASFT demonstrates competitive performance on image-to-point-cloud generation. These results further demonstrate the generalizability of proposed ASFT on diverse image and point cloud generative tasks.

     Model	ULIP-I (↑)	P-FID (↓)
     Shap-E	0.1307	-
     Michelangelo	0.1899	-
     CLAY	0.2066	0.9946
     ASFT (ours)	0.2976	0.3638

   • We also want to emphasize that on ShapeNet, ASFT achieves better performance than the SOTA latent diffusion model LION [1] on standard evaluation metrics, including MMD, COV, and 1-NNA. On image generation datasets FFHQ-256 and LSUN-Church-256, our model outperforms previous function-space generative models. On ImageNet-256, ASFT achieves comparable performance with models on ambient space. Admittedly, there is a performance gap compared with models trained in latent space from a pre-trained VAE. However, we want to point out that the VAE model to compute the latents is trained on a much larger dataset than ImageNet whereas ASFT is trained on ImageNet only. The widely used SD-VAE (https://huggingface.co/stabilityai/sd-vae-ft-mse) for latent space generative modeling, is trained on OpenImages (containing ~9M images) and then finetuned on subset of LAION (containing over 238k images) for a total of ~9.23M images. Whereas ImageNet contains ~1.28M images. This means that ASFT is trained on 13% of the data used to train DiT and SiT. We have updated Tab. 7 in the appendix to reflect this fact, which we believe explains the performance gap between ASFT and latent space generative models.

3. Q: The visualizations in Figures 1 and 2 are unclear and do not provide a satisfactory explanation of the proposed method.

   • We have updated Figures 1 and 2 as shown to better illustrate the pipeline of proposed ASFT.

1. Q: The overall image results are not very convincing - 128x128 images are not a solid proof for superiority of the model. Meanwhile, results on 256x256 are worse than many baselines (even if they are pretrained and domain-specific).

   • We want to highlight that comparing to baselines that are also trained on ambient space like Simple Diffusion (U-Net), RIN and HDiT, we acheive comparable performance while formulating a domain-agnostic approach that is directly applicable to other domains like 3D generation. Other baselines with large model sizes containing 2B parameters (Simple Diff U-ViT and VDM++ U-ViT) achieve better performance than ASFT, which we attribute to their model size being 3 times bigger than ASFT-XL, see Tab. 7 in the Appendix. We have updated Tab. 3 to also reflect this.
   • There is a performance gap compared with models using latent space with pre-trained VAEs. However, we want to point out that these VAEs are trained on a much larger dataset than ImageNet whereas ASFT is only trained on ImageNet. In particular, the widely used SD-VAE (https://huggingface.co/stabilityai/sd-vae-ft-mse) used for latent space generative modeling, is first trained on OpenImages (containing ~9M images) and then finetuned on subset of LAION (containing over 238k images) for a total of ~9.23M images. Whereas ImageNet contains ~1.28M images. This means that ASFT is trained on 13% of the data used to train DiT and SiT. We have updated Tab. 3 and Tab. 7 in the appendix to reflect this.

2. Q: Why is the unified representation important? From the practical point of view, it doesn't make that much of a difference to have two domain-specific backbones instead of a shared one (at least for images and point clouds). If we sacrifice the unification, is there any better backbone choice, especially for images that would result in better results and/or scalability?

   • We thank the reviewer for bringing up this point. Our hypothesis is that the ultimate goal of generative modeling is to train models on every existing bit of information. Unfortunately, these information bits are distributed across different data domains (ie. there are bits of information in images that are not captured by text datasets). A unified generative framework that can leverage different data modalities seamlessly is therefore an important direction to pursue.
   • As the reviewer points out, theres a tradeoff to explore when designing generative models, we can sacrifice the unification (ie. and its simplicity across domains and ability to scale up the training data) to benefit from more efficient domain-specific approaches. In particular, one may benefit from optimizing architectures for each modality separately. For example, we could make use of domain-specific biases like the frequency spectrum of images following a power law [4] or the fact that 3D models are multi-view consistent [5]. However, we want to point out that a trend recent works demonstrate, a Transformer-based architectures that trivially benefits from scaling has achieved superior performance in many applications like image [6], video [6], or even graph structured data [7]. We believe ASFT represents a promising in this direction, where we leverage Transformer-based domain-agnostic generative models that can be effectively trained.

3. Q: The model is shown only for images and 3D point clouds. What about other modalities? Are there any additional challenges?

   • Our model can be applied to other modalities with ease. Once certain modality is formulated as a mapping from coordinate space to signal space, ASFT can be directly applied to this data modality. For example, to build generative model on non-Euclidean spaces (eg. graphs or Riemannian manifolds), we can use an intrinsic coordinate system based on eigen-decomposition to define coordinate-value maps, as suggested in [2]. We specifically opted for images and point clouds since images are structured and dense in 2D space while point clouds are sparse and unstructured representations in 3D. These two settings cover most of the use cases for other data domains.

1. Q: Potential for Lower Fidelity in Complex Domains

   • We want to note that when comparing to baselines trained in ambient space like Simple Diffusion (U-Net), RIN, HDiT, our proposed ASFT achieves comparable performance despite being domain-agnostic and readily applied to other domains like 3D point cloud generation (see our new results on Objaverse in Sect G of Appendix). Other baselines with large model sizes containing 2B parameters (Simple Diff U-ViT and VDM++ U-ViT) achieve better performance than ASFT, which we attribute to their model size being 3 times bigger than ASFT-XL.
   • Admittedly, there is a performance gap compared with models using latent space with pre-trained VAEs. However, we want to point out that these VAEs are trained on a much larger dataset than ImageNet whereas ASFT is only trained on ImageNet. In particular, the widely used SD-VAE (https://huggingface.co/stabilityai/sd-vae-ft-mse) used for latent space generative modeling, is first trained on OpenImages (containing ~9M images) and then finetuned on subset of LAION (containing over 238k images) for a total of ~9.23M images. Whereas ImageNet contains ~1.28M images. This means that ASFT is trained on 13% of the data used to train DiT and SiT. We have updated Tab. 3 and Tab. 7 in the appendix to reflect this.

2. Q: Dependence on Model Size for Best Results

   • We agree with reviewer that the model benefits from increasing model size. However, we want to highlight that this actually shows that ASFT benefits from scale as widely studied in many diffusion models. Latent diffusion models also rely on scaling up the model size to achieve better performance. In fact, our largest model ASFT-XL has comparable model size as DiT-XL and SiT-XL. Investigating efficient variants of current version can be a valuable direction and we’ll include this discussion in future works of updated manuscript.

3. Q: Challenges with High Dimensional Data

   • For high dimensional data, we can subsample the decoded coordinate-value pairs as shown in Figure 3b during training to decrease the FLOPs. As a future direction, more efficient architectures like SSM can be applied. In addition, note that because ASFT models coordinate-value maps it can generate samples in arbitrary dimensions. We have also included additional examples of samples of higher resolution (i.e., 2048$\times$2048) from ASFT trained on ImageNet-256 in Appendix H. It indicates that ASFT can handle high-resolution data in inference trivially.

4. Q: Despite it is a domain-agnostic model, are there any possible way to include domain-specific knowledges to further boost the quality?

   • We agree with the reviewers that integrating domain-specific knowledge can boost the performance of proposed ASFT. For instance, one might take into account the power law of the frequency spectrum of images when designing coordinate-value maps [1] or add regularization between our ASFT and other representation learning models [2]. However, we want to kindly highlight that the scope of our work is to demonstrate a way of building unified generative models that can be applied to different domains with little to no domain-specific tweaking. As shown in the paper, this domain-agnostic framework already achieves comparable performance on image and point cloud generation as curated models for each domain. We believe this indicates a promising direction of building unified and flexible generative model.

References:

[1] Rissanen, Severi, et. al. "Generative modelling with inverse heat dissipation." arXiv preprint arXiv:2206.13397 (2022).

[2] Yu, Sihyun, et al. "Representation alignment for generation: Training diffusion transformers is easier than you think." arXiv preprint arXiv:2410.06940 (2024).