4M-21: An Any-to-Any Vision Model for Tens of Tasks and Modalities

We thank reviewer 66Uu for the positive feedback. We address the main concerns and questions below:

Quality of generated images

Please see the PDF for detailed caption-conditioned generation metrics on COCO, as well as Section 2 of the common response for a discussion.

Detailed configurations of B, L, and XL models

Thank you for pointing this out. We've listed the model configurations in the table below. These configurations are derived from 4M and closely resemble those of T5 and UnifiedIO to ensure comparability with other encoder-decoder models in the literature.

Note that we use SwiGLU [15] in the feedforward, which employs three weight matrices instead of the typical two. To maintain equivalent parameter and operation counts, we've adjusted the feedforward dimension by scaling it by 2/3 (4 x 2/3 instead of 4x).

We thank reviewer PHF3 for the positive feedback. We address the main concerns and questions in the following response:

Encoder-Decoder vs Decoder-only.

Thank you for your question. Both encoder-decoder and decoder-only architectures are valid design choices, depending on the specific use cases and priorities, and this is somewhat that can be explored more solidly in future work. We explain the reasons for our choice of an encoder-decoder architecture below and discuss how a similar model might look like in a decoder-only setting.

Following 4M, we use an encoder-decoder architecture as it is directly compatible with masking approaches (e.g., T5, MAE), which is at the core of our method and enables any-to-any capabilities with a single training objective. In addition, after training, the encoder can be extracted and used as a ViT or a multimodal encoder.  Notably, the bidirectional self-attention in the encoder has been shown to have slight benefits over causal attention for representation learning & transfer tasks (e.g., Table 3d of AIM [9], Table 2 of T5 [10]).

If we were to design a 4M-like model using a decoder-only architecture, there are two main approaches to consider:

The first is a causal decoder (i.e. next token prediction without span or MAE masking). This approach is similar to multimodal LLMs operating on interleaved data (e.g. Gato [11], Fuyu [12], GiT [13], Chameleon [19]) and is the easiest to unify with LLMs. However, the 4M masking strategy would not be directly compatible with this approach. A naive strategy would be to keep everything unmasked and concatenate one modality after another, but this leads to much larger sequence lengths per example and redundancy between modalities. As many of the capabilities shown in our paper rely on cross-modal masking (e.g. the any-to-any capabilities and the ability to predict outputs from partial modalities), it is unclear whether those would be achievable with a causal decoder and what engineering challenges this would involve.

The second approach is using a prefix LM-like decoder (see Fig. 4 of T5 paper [10]), where unmasked inputs (i.e., encoder inputs in the 4M formulation) and masked inputs/targets (i.e., decoder inputs in the 4M formulation) are concatenated. The entire sequence is then given to a single decoder LLM. This approach allows preserving the masking strategy and training objective within a decoder-only architecture, but has seen less adoption than encoder-decoder approaches in the masking literature. However, it is more amenable to multi-turn or temporal inputs, as multiple sequences of unmasked inputs and masked targets can be concatenated one after the other.

In summary, we use an encoder-decoder architecture as it provides a straightforward way to achieve any-to-any capabilities through masking, and allows for downstream reuse of the trained encoder. While decoder-only approaches could potentially be adapted for similar purposes, we are unaware of work demonstrating this at the same scale (in terms of number of modalities) and we believe it to be a very impactful and exciting research direction.

As suggested, we will include additional references to decoder-only multimodal LLMs as well as parts of this discussion in our camera-ready version.

Do different feature maps count as different modalities? Comparison with approaches of GiT and Fuyu-8B.

This highlights an important point; thanks for the question. Three clarifications:

1. our primary goal in this paper is ‘modeling’ work, e.g. demonstrating the possibility of having one model that effectively does any-to-any prediction across many diverse modalities. Once that is established, the community can adopt the model on their own modalities that are justified based on the use case. Our goal is not committing a specific dictionary of modalities or ruling out others, but using a set that provides diversity and is relatively large to enable evaluation.

2. ‘what qualifies as a modality’ is an interesting nontrivial question. Many modalities, e.g., depth/3D, heat, etc., can be sensed using sensors, but also, they can be effectively inferred out of RGB images with a powerful task-specific prediction model. Being inferable with a task-specific neural network does not question if depth/3D qualifies as a modality of information. In general, in theory, any specific mode of information about an underlying scene qualifies as a modality. Which modalities are ‘useful’ is a different deeper question, often application specific, and out of our scope as the paper’s focus is making progress on the model.

3. adopting neural network feature maps as a modality is useful as, they can be used for chained conditioning [18], and also they provide a way for distilling popular single networks (e.g. DINO) into one multitask network.

Similar to GiT, we employ tokenization of several modalities via language, e.g. bounding boxes, metadata and color palette. However, this strategy is limited for tokenizing dense representations such as depth, normals, edges, instance segmentation masks, or feature maps. Thus, we employ different tokenization schemes to handle such modalities as depicted in Figure 3. We will add a discussion on this to the paper.

Challenge of negative transfer.

Thanks for the pointer, we will add a discussion. While we didn’t observe negative transfer when more modalities are included, investigating this direction further could be worthwhile.

Detail the training framework

Thanks. We will add a diagram to camera-ready to improve the clarity as suggested.

We thank reviewer b2ju for the constructive feedback. We address the main concerns and questions in the following response:

What are the performances of the trained tokenizers?

Please see the following items as the response: 1) Figure 1 in the rebuttal PDF for a visual comparison of multimodal tokenizer reconstructions (resolution 224x224), 2) the “tokenizer bound” reported in the main paper's Table 1 that represents the performance upper bound — it clearly shows that the prediction errors of all models are notably worse than the tokenization error indicating that the lossiness of discrete tokenization is not the bottleneck.

Generalization of tokenizers

We train tokenizers on CC12M to match our training data distribution, and find them to perform well within that distribution (see above answer). For more niche domains not captured in web-scale image distributions (medical, satellite, etc.), we recommend training domain-specific tokenizers. That said, we note that tokenizers trained on web-scale data perform significantly better in areas like human face and text reconstruction compared to popular tokenizers like VQ-GAN [1] that were trained on ImageNet-1k; a finding supported by MAGVIT-v2 [2].

Capturing high-frequency details

This is challenging for discrete tokenizers, but can be addressed by increasing token density and performing an additional token super-resolution step, similar to 4M and Muse. Furthermore, recent advances in VQ-free tokenization [2,3] show strong scaling trends in reconstruction and generation performance relative to the vocabulary size.

We would like to stress that discrete tokenization is one of the key design decisions that allows us to scale a unified model to such a large and diverse set of modalities, without multi-task loss-balancing, task-specific heads, or encountering training instabilities. We expect the above-mentioned advances in discrete tokenization to show direct benefits in multimodal generation and out-of-the-box performance.

Multimodal retrieval numbers

Please refer to the PDF Tab. 2 for additional evaluations on cross-modal retrievals. Our model has notable performance for different retrieval tasks (RGB-Text, RGB-Depth, RGB-Semantic) while being only trained on global embeddings extracted from the RGB images.

Regarding DINOv2, it is trained only on RGB images and is not a multimodal model, thus capable of only RGB-RGB retrieval. We also compared our model’s RGB retrieval performance with DINOv2/ImageBind using ImageNet kNN classification, (Tab. 1 of the main paper, also included in the PDF). The results in the table show that our model’s RGB retrieval capability successfully matches DINOv2 and ImageBind, and on top of that, it can perform multimodal retrieval which DINOv2 and other similar models cannot.

Conditional generation numbers

Please see the PDF for detailed caption-conditioned generation metrics on COCO, as well as Section 2 of the common response for a discussion.

How did you choose each vocabulary size?

We follow the best practices from community and tune the vocabulary size to obtain low reconstruction error while keeping the vocabulary size as small as possible. Please see Fig. 15 and Appendix J for an ablation.

Metadata decision process

Our choice of which metadata to include was guided by Omnidata [4] and SDXL [5]. Omnidata extracts several types of metadata from a multimodal dataset to give researchers a high-level set of controls to steer the data distribution towards high/low walkability, occlusion, etc. Our goal was to expose similar parameters, but for a multimodal generative model. We believe that such capabilities provide a path to answer questions on what the ideal multimodal data distributions are for different downstream tasks. Furthermore, we looked towards SDXL and more broadly ControlNet [6] as inspiration for metadata that can be used for image generation, e.g. original image size, colorfulness, etc. While many types of metadata can be included as text in a prompt, we observe strong condition adherence with metadata. In summary, our aim was to provide a broad set of control knobs by extracting various parameters from RGB images and the different pseudo labels. We note that this set can be easily extended towards scores like aesthetics, NSFW, watermark, etc.

Treating RGB tokens and pixels differently

Thank you, this is an important design choice in current multimodal models.

Discrete tokens enable iterative sampling, making them useful for generative tasks, which is why most token-based generative methods (e.g., Parti [7], MaskGiT [8]) use them. However, discretization leads to information loss, which is not ideal for visual perception tasks. On the other hand, using RGB pixels as input is more suitable for visual perception tasks. By avoiding the discrete bottleneck, there is no information loss during the tokenization step, and the projection layer can be more lightweight (in 4M's case, it is just a simple linear projection). Given these tradeoffs, we follow 4M by training on both and treating them as separate modalities, with RGB pixels as an input-only modality. This allows us to choose the most appropriate representation for a given task - discrete RGB for generation or RGB pixels for perception.

The added value is mostly due to engineering.

We find it hard to concretely respond to this comment as a weakness since, first, many undeniably key contributions in the community are basically “engineering”, and second, what would be “engineering” is too vague and broad. Engineering doesn’t mean unimpactful. All reviewers positively recognized that this paper required a significant amount of exploration and work, e.g., in terms of scaling, and the workload was viewed as a strength of the paper. Besides those, the paper made several contributions and introduced new capabilities across key axes, as summarized in L52-L70 in the main paper.