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

Authors: Yuyang Wang, Anurag Ranjan, Joshua M. Susskind, Miguel Ángel Bautista
Primary area: generative models
Keywords: Generative Model, Flow Matching, Domain Agnostic

Abstract:
Flow matching models have emerged as a powerful method for generative modeling on domains like images or videos, and even on unstructured data like 3D point clouds. These models are commonly trained in two stages: first, a data compressor (\ie a variational auto-encoder) is trained, and in a subsequent training stage a flow matching generative model is trained in the low-dimensional latent space of the data compressor. This two stage paradigm adds complexity to the overall training recipe and sets obstacles for unifying models across data domains, as specific data compressors are used for different data modalities. To this end, we introduce Ambient Space Flow Transformers (ASFT), a domain-agnostic approach to learn flow matching transformers in ambient space, sidestepping the requirement of training compressors and simplifying the training process. We introduce a conditionally independent point-wise training objective that enables ASFT to make predictions continuously in coordinate space. Our empirical results demonstrate that using general purpose transformer blocks, ASFT effectively handles different data modalities such as images and 3D point clouds, achieving strong performance in both domains and outperforming comparable approaches. ASFT is a promising step towards domain-agnostic flow matching generative models that can be trivially adopted in different data domains.

Decision: reject
Reviewer scores: 5, 6, 6, 3

Reviews and discussion:
<thread>
creation date:
1730629617693

thread_type:
official_review

<thread_object>
creator: reviewer_pe6r

<content>
summary:
This paper claims that building latent-space diffusion models have several shortcomings such as non-end-to-end optimization and the requirement of domain-specific compression (e.g., VAE) models, and thus designs an ambient-space flow transformer (ASFT) architecture for generative tasks.

strengths:
1. The analyses of the potential shortcomings of existing latent diffusion models are presented. Based on the motivations, this paper constructs a single-stage learning model, 
2. Learning over coordinate-value pairs facilitates data-agostic processing. In this way, an image is actually treated as a "colored point cloud".

weaknesses:
1. Although the authors analyzed some aspects of potential drawbacks of latent-space diffusion models, it is hard to be convinced that ambient-space (e.g., pixel-space, point-space) learning is a more promising direction: 
-- First, I don't think training domain-specific data compressors is a cumbersome practice that should be criticized. When some powerful VAEs are already trained and released, the community can directly use them. There aren't that many types of data modalities. Building compressors for each type of data (e.g., image, video, point cloud, mesh) is totally acceptable.
-- Second, the current mainstream practice is to separately train the data compressor and the subsequent latent diffusion model, but it does not mean that such training workflow cannot be made end-to-end. We can't say it is a drawback just because we haven't explored it. Generative models evolve so fast. I think making it end-to-end is not impossible.
-- Third, anyway, for now the great success of various latent diffusion models seems to demonstrate the superiority of learning in the latent space instead of the ambient space.

2. Building generative models with coordinate-value pairs may essentially restrict its application scenarios and conditional generation capabilities. For 3D generation, point cloud is apparently not the final choice. What we want is the continuous surface, together with textures. Existing 3D generative models either use meshes or implicit fields. However, the proposed method faces difficulties in generating such data. Besides, I notice that 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. 

3. The experimental settings are not very persuasive. For image generation, the model is trained on ImageNet. For point cloud generation, the model is trained on ShapeNet. The experimental results cannot demonstrate the potential of the proposed method for learning from larger-scale high-quality data, such as LAION and Objaverse datasets. It can be observed that the diversity and quality of the generated images obviously cannot catch up the current state-of-the-art latent image diffusion models. The generated point clouds are also noisy and lack details, especially when compared with state-of-the-art 3D native diffusion models (like CLAY).

questions:
1. Does the proposed generation framework support more practical conditioning mechanism (e.g., guided by text or single image)?
2. Can the proposed learning model be scaled up using larger-scale high-quality datasets such as LAION and Objaverse?
3. The authors are suggested to provide further explanations about how the latent variable z_{f_t} for context encoding is obtained.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730677761524

thread_type:
official_review

<thread_object>
creator: reviewer_n7ez

<content>
summary:
The paper introduces **Ambient Space Flow Transformers (ASFT)** as a novel approach to generative modeling that simplifies the training process by eliminating the need for latent space data compressors. ASFT works directly in the ambient space (i.e., the original data domain), aiming to be a domain-agnostic model applicable across different types of data, such as images and 3D point clouds.

strengths:
- **Simplicity and Efficiency**: ASFT's single-stage training process is simpler than traditional two-stage models, making it easier to implement and tune.
- **Versatility**: It is designed to work across multiple data types without modifications, unlike models that require domain-specific architectures.
- **Resolution Scalability**: The ability to generate high-resolution outputs from lower-resolution training data provides flexibility and potential computational savings.
- **Competitive Results**: Despite its simplicity, ASFT achieves strong performance metrics comparable to state-of-the-art models across images and point clouds.

weaknesses:
- **Potential for Lower Fidelity in Complex Domains**: ASFT may not match latent space models in specific metrics, such as Fréchet Inception Distance (FID), particularly when latent space models are pre-trained on extensive datasets.
- **Dependence on Model Size for Best Results**: The model's performance improves significantly with scale, which may lead to high computational costs for larger ASFT versions, particularly in comparison with models leveraging pre-trained compressors.
- **Challenges with High Dimensional Data**: For very high-resolution applications, ASFT’s point-wise approach might face optimization difficulties due to the increased complexity in decoding large numbers of coordinate-value pairs.

questions:
Despite it is a domain-agnostic model, are there any possible way to include domain-specific knowledges to further boost the quality?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730681630205

thread_type:
official_review

<thread_object>
creator: reviewer_tvcm

<content>
summary:
The paper presents Ambient Space Flow Transformer (ASFT), a flow matching generative model working with an implicit representation of the data (INR), instead of a latent one. A modified version of PerceiverIO used as a backbone allows both images and 3D point clouds to bo modeled without any design changes.

strengths:
1. The novelty is good - single stage generation via flow matching on function space without any pretrained encoder-decoder models.
2. On images, ASFT beats function space approaches on FID.
3. On 3D point clouds, the method is better or on par with other models.
4. Due to the independence of the input coordinates, the model allows sub-sampling of point clouds and super-resolution of images.
5. The writing is very good. The information flow is maintained and all of the ideas are clearly explained.

weaknesses:
1. The image results presented are in low resolution only - 128x128 and 256x256.
2. 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 though domain-specific and using pretrained models).

questions:
1. 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?
2. The model is shown only for images and 3D point clouds. What about other modalities? Are there any additional challenges?
3. Does the image model scale to higher resolutions?
4. Figure 2 may suggest both image and 3D coordinates are passed to the network at the same time. The authors should consider clearly separating them in the figure and updating the caption to highlight this.
5. In Figure 3 (b) color palette is hard to read - the differences should be more visible.
6. It would be interesting to see a comparison with standard image upscaling in Figure 4 (a).
7. In introduction, citations for VAE, VQVAE, VQGAN, transformers, PointNet, UNet are missing.
8. "UNet" and "U-Net" used - please pick one.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730705370147

thread_type:
official_review

<thread_object>
creator: reviewer_1rvi

<content>
summary:
This paper introduces an Ambient Space Flow Transformer (ASFT), a flow-matching generative model designed to operate directly in the ambient space. The core innovation lies in eliminating the practical complexities of training latent space generative models, such as the reliance on domain-specific compressors for different data domains or the tuning of data compressor hyperparameters (e.g., adversarial weights, KL terms). Moreover, experimental results on both image and point cloud domains demonstrate competitive performance.

strengths:
This paper effectively tackles the challenge of simplifying the training process for flow matching models. The proposed method is innovative and well-supported by experiments conducted on diverse datasets, including images and point clouds. The clear presentation of the problem, method, and results makes the paper easy to follow.

weaknesses:
The novelty of the proposed method in this paper is questionable, as it appears to be relatively straightforward. Additionally, the experimental results are not sufficiently compelling. The use of only the ShapeNet dataset for point cloud modality limits the generalizability of the findings. More diverse datasets should be employed to validate the effectiveness of the proposed approach. Furthermore, the improvement over existing methods is not substantial, and a more comprehensive comparison with state-of-the-art methods, using a wider range of metrics, is necessary. Lastly, the visualizations in Figures 1 and 2 are unclear and do not provide a satisfactory explanation of the proposed method.

questions:
The following suggestions could enhance the paper: 1) Figures 1 and 2 should be refined to provide a clearer visualization of the proposed method; 2) 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; 3) Incorporating additional novel ideas could further enhance the distinctiveness and advancement of the proposed method.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732063658218

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Official Comments to All Reviewers

comment:
We thank all reviewers for their constructive suggestions that help substantially improve the quality of our paper. We have updated our manuscript accordingly, with changes highlighted in red color. Please find below the major clarifications and updates in response to the review: 

1. The main contribution of our work is to build a domain-agnostic generative model on coordinate-value maps (also referred to as “fields”) that can be efficiently trained in large scale settings (eg. ImageNet-256 and Objaverse). We aim to build a unified and simplified generative model that can be seamlessly applied to different data domains. We also want to emphasize that ASFT achieves competitive performance on standard benchmarks for image and 3D point cloud generations.  On ShapeNet, ASFT achieves better performance than the SOTA latent diffusion model LION on standard evaluation metrics. 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 as highlighted in the updated Tab. 3. 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.

2. We added experiments of image-to-point-cloud generation on Objaverse to validate the capability of proposed ASFT on **(1) larger and more challenging 3D generative tasks, (2) different conditioning inputs (i.e., image conditioning)**. Objaverse is a large-scale dataset that contains more than 800k 3D objects of wide variety. Results on Objaverse are listed in the table below. We report ULIP-I which measures the alignment between conditioning image and generated point clouds, as well as P-FID which measures the distribution similarity between sampled and real objects. Compared with SOTA 3D generative models like CLAY, our ASFT demonstrates strong performance on image-to-point-cloud generation. Please find more details in Appendix G and example samples in Fig. 10 of updated manuscript. 
| Model | ULIP-I (↑) | P-FID (↓) |
| -------- | ------- | -------- | 
| Shap-E           | 0.1307 | - |
| Michelangelo | 0.1899 | - |
| CLAY              | 0.2066 | 0.9946 |
| ASFT (ours)   | **0.2976** | **0.3638** |

3. In Fig. 11, we added more resolution agnostic sampling results. We showcase that given an ASFT trained on ImageNet-256, it can trivially generate images at high resolution from $512\times512$ to $2048\times2048$. This demonstrates the flexibility of ASFT as well as its capability to handle high-resolution in inference. Besides, we want to also highlight the results in Tab. 9 where ASFT shows better quantitative results than standard upsample strategies (e.g., bilinear and bicubic) in generating high-resolution images.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732941348204

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Discussion period

comment:
Dear SAC, AC and reviewers,

We kindly reach out to you once again to inquire about the status of the discussion. As we approach the end of the extended discussion period we want to note that only 1 out of 4 reviewers have acknowledged receiving our rebuttal. We believe our rebuttal has addressed all the points raised by reviewers (including new large scale experiments on Objaverse) and are happy to address any further suggestions from reviewers.

Thanks for your time
</content>
</thread_object>

</thread>


Metareview:
metareview:
The aim of the paper is to propose a domain-agnostic generative model on coordinate-value maps that can be efficiently trained at scale.
It shows 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.

On the positive side, the writing is very good. The information flow is maintained, ideas are clearly explained and the presentation of the problem, method, and results makes the paper easy to follow. The proposed generative model is domain-agnostic, meaning it can be applied to image generation but can also be trained with minimal modifications for 3D point cloud generation. In the experiments, the paper shows a compelling performance at image generation, when compared to other function space approaches as well as 3D shape generation, when compared to older point-cloud generation methods.

On the negative side, I see four weaknesses that need to be addressed (listed in order of importance):

1) The paper should better work out the advantage of a domain agnostic architecture, which remains unclear. We have good domain-specific architectures for image generation and 3D generation, that perform much better compared to the proposed model. What is the advantage of the proposed domain-agnostic model? The paper mentions end-to-end optimization as the single core advantage, which might be true, but this is not worked out well in the experiments, e.g. maybe training times/ compute budgets between would be significantly different? Maybe one could actually train a domain-agnostic model that can do both image and 3D generation? Overall, I think this is a very critical point that should be discussed in more detail in the introduction and in the experiments.

2) The current results, are promising, not sufficiently convincing yet. At shape generation, the proposed model only outperforms older point-cloud based methods such as LION. The paper misses showing that it also significantly lacks behind SOTA 3D generation methods that are domain-specific, such as XCube [1] or MeshGPT [2], which I think should be included in the report. The authors argue that their lack of performance compared to SOTA image generation methods is explained by their limited training data. However, I agree with the reviewers who suggested the next step to be an attempt to scale the model and training data size to be better comparable to other methods. 
[1] XCube: Large-Scale 3D Generative Modeling using Sparse Voxel Hierarchies, CVPR 2024.
[2] MeshGPT: Generating Triangle Meshes with Decoder-Only Transformers, CVPR 2024. 

3) Due to the independence of the input coordinates, the paper claims a contribution allowing a very simple mechanism for sub-sampling of point clouds and super-resolution of images. However, the experiments are purely qualitative. Overall, this might be an interesting property of the architecture, but it is not convincingly presented. E.g. the images look blurry, and it remains unclear how they would compare even to the simplest upsampling methods.

In summary, the weaknesses in the presentation and experiments out-weigh the strengths of an otherwise interesting approach in a well-written paper.