FLAT: Resampling Image and Text into 1D Flexible-Length Aligned Transmodal Tokens for Retrieval and Generation

AI summary

Overview

Research area: Multimodal representation learning, cross-modal retrieval, and conditional generation (text-to-image synthesis and image captioning) in computer vision and natural language processing.

Technical level: Advanced. The paper assumes familiarity with contrastive learning (CLIP-style InfoNCE), vision-language model backbones, LoRA adaptation, rectified-flow / flow-matching diffusion decoders, and nested dropout.

Scope: The paper introduces FLAT, a single representation pre-training framework that encodes both images and text into a shared continuous 1D token sequence of flexible length, then uses that same representation simultaneously as a retrieval embedding and as the conditioning signal for both a text-to-image decoder and an image-to-text decoder.

What This Paper Is About

Most multimodal systems learn visual or textual representations in one stage and then bolt on a separate generative model in a second stage, usually behind frozen embeddings. This decoupling limits how much the generative model can benefit from the representation. FLAT instead jointly optimizes one shared multimodal encoder together with a text-to-image (T2I) decoder and an image-to-text (I2T) decoder, producing a representation that is contrastively aligned, linearly interpolatable, and directly consumable by generators in both directions.

Key Contributions

1. A shared 1D continuous representation space for both modalities. FLAT appends 256 learnable register tokens to either an image or a text input, encodes the sequence with a shared VLM encoder, and projects the register hidden states into a representation of dimension $N \times d$ with $d = 64$. Both modalities share the same registers and the same latent projection.

2. Flexible-length sequences via nested dropout over prefix-$K$ tokens. A keep-length $K$ is sampled from a geometric ladder $\mathcal{K} = {1, 4, 16, 64, 256}$, and only the prefix $\mathbf{z}_{:K}$ is passed to the contrastive objective and decoders. The paper reports that direct truncation (dropping the tail) works better in their setup than padding tail positions with learned null embeddings as in FlexTok, and that sampling from a geometric ladder avoids wasting optimization steps on imperceptible length differences.

3. Joint training of alignment and bidirectional generation. A single objective combines an InfoNCE contrastive loss computed as a late-interaction score over matching register positions, an autoregressive I2T captioning loss, and a conditional flow-matching T2I loss, with contrastive, text-generation, and image-generation weights all set to 1.0.

4. Demonstration that the learned geometry supports zero-shot latent operations. The paper shows linear interpolation between two FLAT codes, latent-space arithmetic of the form $\mathbf{z}_{\text{edit}} = \mathbf{z}_1 - \mathbf{z}_2 + \mathbf{z}_3$, and composed retrieval via add/remove clause arithmetic, all without task-specific editing supervision.

Main Findings

• Pre-training alone produces a working generative and retrieval model. From the jointly pre-trained checkpoint with no task-specific adaptation, FLAT reaches a GenEval score of 71.1 (0.711) at $K = 256$, scaling from 0.317 at $K = 1$, 0.613 at $K = 4$, 0.703 at $K = 16$, and 0.701 at $K = 64$. Zero-shot MS-COCO Recall@5 is reported as 69.1 (I2T) and 64.6 (T2I).

• Task-specific fine-tuning brings performance in line with reported state-of-the-art baselines. Fine-tuning yields 83.1 GenEval on T2I generation, 40.5 BLEU-4 and 138.6 CIDEr on MS-COCO image captioning, and Recall@5 of 86.8 (I2T) / 75.8 (T2I) on MS-COCO and 98.3 (I2T) / 93.6 (T2I) on Flickr30K.

• Retrieval accuracy is nearly invariant to prefix truncation. In Table 1 on the MS-COCO Karpathy 5k split, I→T R@5 moves only from 86.44 at $K = 1$ to 86.82 at $K = 256$; T→I R@5 moves from 75.59 to 75.83. On Flickr30k Karpathy 1k, I→T R@5 ranges from 98.10 to 98.30 and T→I R@5 from 93.16 to 93.72 across the sweep. The paper describes this as every metric staying within roughly one percentage point across a $256\times$ reduction in width.

• A single 64-dimensional token carries strong discriminative and semantic content. Each token contains 64 dimensions, spanning a 64-dimensional embedding at $K = 1$ to a 16,384-dimensional embedding at $K = 256$. The paper also reports that on MS-COCO the first token captures key cross-modal discriminability and that even at $K = 1$ FLAT outperforms baselines such as CrossFlow and SCD-Net across CIDEr, METEOR, and SPICE.

• Frozen FLAT representations beat other generative latent spaces on ImageNet-1K linear probing. A single linear classifier on concatenated frozen tokens gives 73.3 top-1 at $K = 1$ (64 dims), 77.0 at $K = 4$ (256 dims), 81.2 at $K = 16$ (1024 dims), and 81.8 at both $K = 64$ (4096 dims) and $K = 256$ (16384 dims). Comparison points in the same table are SD-VAE latent (dim 4, 8.0), REPA (62.5), MAE-B (dim 768, 68.0), and DREAM (dim 1024, 72.7).

• Task complexity maps onto required prefix length for generation. For GenEval categorical analysis, a single token captures basic semantics such as single object and color, while two-object, position, and color-binding categories score near zero at $K = 1$ (for example 0.03 for color binding), recover steeply by $K = 4$ and $K = 16$ (reaching 0.67). Increasing $K$ also sharpens caption detail, for example "street" → "alley" and "a shelf" → "a shelf with baskets".

• Pre-trained and fine-tuned checkpoints differ largely in output style, not semantics. In zero-shot captioning, performance rises to 63.8 CIDEr and 16.3 SPICE at $K = 256$. At $K = 64$, 99.6% of generated sequences terminate properly with an EOS token and average 10.8 words, but 27.1% begin with quotation marks characteristic of image titles. The paper attributes part of the fine-tuning gains on COCO to output-format alignment rather than improved semantic understanding alone.

• Zero-shot retrieval degrades slightly as prefix length grows. Without retrieval fine-tuning, MS-COCO I→T R@5 falls from 69.10 at $K = 1$ to 60.04 at $K = 256$, and T→I R@5 falls from 64.59 to 55.87. On Flickr30k, I→T R@5 falls from 89.70 to 84.30 and T→I R@5 from 86.04 to 81.30. Task-specific fine-tuning substantially mitigates this sensitivity.

• The three losses are mutually reinforcing rather than conflicting. In the ablation over all seven non-empty loss combinations trained for 40k steps, the full objective gives retrieval 52.2 (I2T) / 56.3 (T2I), B@4 40.2, CIDEr 137.5, and GenEval 0.327. Removing contrastive collapses retrieval to 13.6 / 6.7; removing generation collapses GenEval to 0.003; generation-only collapses retrieval to 0.6 / 1.1. The paper reports a Shapley value decomposition over all eight coalitions, showing each task is driven primarily by its matched loss while other objectives supply complementary signal.

• FLAT closes much of the modality gap. In PCA visualizations on MS-COCO, dual-encoder baselines such as CLIP and SigLIP2 show clear spatial separation between image and text representations, whereas FLAT produces highly overlapping distributions. A single register token cuts CLIP's centroid distance by half, and the full 256 tokens reduce it to a quarter.

• Pre-training matters most for generation, less for in-domain retrieval. Compared against task-only controls with the same architecture, data, budget, and evaluation protocol, FLAT pre-training improves GenEval by 0.16–0.24 for $K \geq 4$ and reaches 0.83 at $K = 16$ and $K = 64$, while the task-only model stays near 0.60 across prefix lengths and is higher at $K = 1$. The pre-trained initialization gains 0.34 from $K = 1$ to $K = 16$. Captioning improves by 9.0–10.9 BLEU-4 and 31.9–41.2 CIDEr. For retrieval, I2T R@1 improves by 0.6–1.3 points, but the task-only model is 0.4–0.9 points higher for T2I R@1.

Methodology in Plain English

FLAT starts from a pre-trained vision-language model (Qwen3.5-2B) and adds 256 learnable "register" tokens to every input, whether that input is an image or a piece of text. The encoder processes input plus registers under the instruction "Represent the user's input," and the hidden states at the register positions are linearly projected into the representation.

To make the representation's length flexible, each training step picks one length $K$ from ${1, 4, 16, 64, 256}$ and only keeps the first $K$ tokens. The same $K$ is used across all ranks and for all losses in that step, so the contrastive and generative objectives always see the identical prefix. At inference, $K$ can be varied freely, letting one encoder pass produce coarse or fine outputs.

The representation then feeds two decoders. An autoregressive decoder takes the visual tokens and generates a caption. A rectified-flow transformer (initialized from SANA-1.6B) takes the textual tokens as cross-attention conditioning and denoises image latents. In parallel, the image and text representations are contrasted position-by-position using a late-interaction score averaged over the $K$ registers, trained with InfoNCE at temperature 0.07 and cross-device negatives.

Pre-training runs on 65.7M image–text pairs for 135k steps (Table 4 lists 67M image–text pairs, composed of 30M long-caption, 33M short-caption, 4M prompt-style, and 0.1M curated high-quality pairs), with only 11.4M trainable parameters on frozen base weights. Fine-tuning then updates only the component relevant to one task: the image decoder for T2I (8,000 steps on 120K high-quality pairs), the text decoder LoRA for I2T (5,000 steps on COCO Karpathy train+restval, 113,287 images), and the encoder LoRA, registers, and latent projection for retrieval (5,000 steps for COCO, 320 for Flickr30K). Evaluation uses GenEval across all 553 prompts with 4 images per prompt, 20 sampling steps, CFG 4.5, and no prompt rewriting; COCO Karpathy test for captioning with greedy decoding; COCO Karpathy 5K and Flickr30K Karpathy 1K for retrieval; CIRR with 1,000 queries against a 1,000-image gallery for composed retrieval; and 1,281,167 ImageNet-1K training images for 20 epochs with evaluation on the 50,000-image validation split for linear probing.

Why This Matters

Impact on research. FLAT challenges the standard two-stage recipe in which representations are learned first and frozen for generators later. It shows that contrastive alignment and bidirectional generation can be optimized in a single pre-training stage without the objectives cannibalizing each other, and that the resulting space is not just discriminative but geometrically well-behaved: interpolatable, arithmetically composable, and free of the persistent image-text modality gap seen in dual-encoder models. It also extends 1D resampled tokenization, previously used mainly for image reconstruction, to a shared space covering text as well.

Real-world applications.

• Cross-modal search and recommendation, where one compact 64-dimensional token

Authors’ abstract

Traditional multimodal representation learning and generation are two stages: a contrastive or self-supervised visual encoder is trained first, followed by a separate downstream generative model. This setup bottlenecks generative performance behind frozen embeddings. To bridge this gap, we revisit joint multimodal representation learning and generation to produce linearly interpolatable embeddings that are directly consumable by generative decoders. We present FLAT (Flexible-Length Aligned Transmodal representations), a representation pre-training framework that jointly optimizes a shared multimodal encoder alongside downstream text-to-image (T2I) and image-to-text (I2T) decoders. By combining contrastive alignment with bidirectional cross-modal generative objectives, FLAT ensures its representations function as both discriminative semantic descriptors and generative conditions. Architecturally, FLAT maps visual and textual inputs into a unified continuous 1D sequence space, applying nested dropout over prefix-K tokens to enable dynamic output lengths. A single pre-training stage allows FLAT to perform cross-modal retrieval and generation across variable prefix K, achieving a T2I GenEval score of 71.1. Task-specific fine-tuning aligns model performance with state-of-the-art baselines: 83.1 GenEval on T2I generation; 40.5 BLEU-4 and 138.6 CIDEr on MS-COCO image captioning; and Recall@5 scores of 86.8 (I2T) / 75.8 (T2I) on MS-COCO alongside 98.3 (I2T) / 93.6 (T2I) on Flickr30K. Finally, qualitative evaluations demonstrate that FLAT representations natively support linear interpolation, latent space arithmetic, and zero-shot composed retrieval.