FLAT: Resampling Image and Text into 1D Flexible-Length Aligned Transmodal Tokens for Retrieval and Generation
Guangyu Sun, Shlok Kumar Mishra, Wentao Bao, Robert Zhenheng Yang, Xiao Wang, Xiyuan Wang, Yujunrong Ma, Chen Yuan, Max Xiangjun Fan, Jun Xiao, Jianpeng Cheng
cs.CV
2026-09-15
FLAT jointly trains a shared encoder with T2I and I2T decoders on nested 1D prefixes. Pre-training reaches 71.1 GenEval; task fine-tuning hits 83.1 GenEval, 138.6 CIDEr, and COCO R@5 of 86.8/75.8.
Contrastive encoders such as CLIP give a retrieval-friendly embedding. Generative models then freeze that embedding and train a decoder around it. Stable Diffusion, SDXL and Sana condition on frozen CLIP, T5 or LLM text features. LLaVA and BLIP-2 freeze a vision encoder and realign it to language. The generative ceiling sits at a representation that was never trained to be a generative condition.
FLAT trains one representation for both jobs. A shared encoder maps image and text into a single continuous 1D sequence that a retriever can score and that T2I and I2T decoders can condition on.
The encoder is Qwen3.5-2B with a frozen backbone and LoRA, prompted with "Represent the input". 256 learnable register tokens of 64 dimensions are appended; a linear map reads their hidden states as z. Image and text share the registers and the projection. Trainable encoder-side parameters are about 11.4M.
Nested dropout is applied to z. Each step samples a prefix length K from {1, 4, 16, 64, 256} and feeds only z[:K] to the contrastive loss and both decoders. The tail is truncated, not padded with FlexTok-style null embeddings. One K is broadcast across ranks so alignment and generation see the same prefix. A geometric ladder is used instead of every integer, so gradient is not spent on lengths that barely differ.
The image representation becomes soft tokens for an autoregressive captioner. The text representation conditions a SANA-1.6B rectified-flow transformer via cross-attention, denoising VAE latents. Contrastive training is late-interaction InfoNCE: cosine similarity averaged over the K aligned registers, both directions, temperature 0.07. The three losses are weighted equally.
Pre-training uses about 67M image-text pairs for 135k steps at global batch 1024 on 64 H200 GPUs. Task fine-tuning then splits: 8k steps of the image decoder on 120K clean pairs for T2I; 5k steps of the text LoRA for captioning; encoder LoRA plus registers for retrieval.
Prefix length tracks compositional difficulty. K=1 covers single-object and color; color binding scores 0.03. By K=16 the harder GenEval slices recover to 0.67 and overall GenEval reaches 0.83. Retrieval barely moves with width: COCO image-to-text R@1 goes from 63.00 at K=1 to 63.54 at K=256.
| Setting | Metric | Result |
| Pre-train, K=256 | GenEval | 71.1 |
| Fine-tune, K=16/64 | GenEval | 83.1 (MetaQuery-L 78, no rewrite) |
| Fine-tune, K=256 | COCO BLEU-4 / CIDEr | 40.5 / 138.6 (BLIP-2: 43.7 / 145.8) |
| Fine-tune | COCO R@5 I2T / T2I | 86.8 / 75.8 |
| Frozen features | ImageNet linear probe | 73.3 at K=1, 81.8 at K=64 (DREAM 72.7 at 1024-d) |
Drop the contrastive term and retrieval collapses (I2T R@1 52.2 to 13.6). Drop the image-generation term and GenEval falls from 0.327 to 0.003. A task-only control, same architecture and budget, stalls near 0.60 GenEval for K≥4; joint pre-training adds 0.16–0.24.
A single token cuts CLIP's image-text centroid distance from 0.824 to 0.372; 256 tokens reach 0.214. Linear interpolation and z1−z2+z3 arithmetic decode coherently on both heads. Zero-shot CIRR composed retrieval is strongest at K=1, Hit@1 44.0.
The prefix is a bandwidth knob: 64 dimensions retrieve, 16 tokens generate compositions. Continuous codes keep CLIP-style arithmetic while remaining a generative condition. This is an incremental systems result. The backbones are a 2B VLM and a 1.6B flow decoder; the contribution is the joint objective and the variable-length 1D interface, not a new generator.
There is no dedicated limitations section. The headline 83.1 GenEval and the retrieval numbers come from task-specific fine-tunes, not one checkpoint that tops all three tasks. Caption CIDEr 138.6 trails BLIP-2 (145.8) and CoCa (143.6). Retrieval is contrastively adapted on COCO and Flickr, so it is not a zero-shot CLIP score. Interpolation and arithmetic are mostly qualitative. At K=1 the task-only GenEval (0.60) beats the pre-trained init (0.49), so short-prefix generation does not come only from joint pre-training. Pre-training itself ran 135k steps on 64 H200 GPUs.