GRAN-TED: Generating Robust, Aligned, and Nuanced Text Embedding for Diffusion Models

AI summary

Overview

Research area: Computer vision — specifically text encoders for text-to-image (T2I) and text-to-video (T2V) diffusion models. Technical level: Intermediate (the core ideas are accessible; the training recipe assumes familiarity with LLMs, contrastive learning, and diffusion transformers). Scope: The paper introduces a text-only benchmark for evaluating text encoders and a two-stage recipe for building a stronger one called GRAN-TED.

What This Paper Is About

Diffusion models translate text prompts into images or video, and the text encoder—the component that turns a prompt into a numerical representation—largely determines whether the output actually matches what was asked for. The problem is that there has been no cheap, reliable way to tell whether one text encoder is better than another: existing methods either require training a full diffusion model for every candidate (extremely expensive) or rely on NLP/retrieval benchmarks that do not predict visual generation quality. This paper builds a fast text-only benchmark to rank encoders, then uses it to guide the development of a superior encoder.

Key Contributions

1. TED-6K, a text-only benchmark for text encoders. A dataset of 6,641 evaluation instances, each pairing a detailed source caption with one positive statement and multiple "hard" contradicting statements, spanning nine semantic dimensions (action, spatial relations, temporal relations, coreference, adjectives, adverbs, quantity, OCR, and basic events). Evaluating an encoder takes roughly 4 minutes versus about 50 hours to train a diffusion model from scratch—roughly a 750x speedup.
2. A unified context aggregator. A small two-layer attention module with a learnable context token, trained with a contrastive loss, that mimics how a Diffusion Transformer consumes text embeddings. This allows CLIP, T5, LLM, and MLLM encoders to be compared fairly at the representation level rather than through a single pooled vector or misleading question-answering accuracy.
3. The GRAN-TED encoder. A two-stage construction: first fine-tuning Qwen3-VL-8B-Instruct on a large, generation-oriented VQA and captioning dataset to align its latent space with visual concepts, then applying a learnable layer-wise weighting module that fuses hidden states across the model's depth.
4. Validation of the benchmark against real generation. A correlation study showing TED-6K scores strongly predict downstream T2I and T2V performance, plus demonstrations that GRAN-TED improves actual generation results.

Main Findings

• The benchmark predicts downstream quality. TED-6K scores correlate strongly with generation metrics: Pearson r of 0.9923 and 0.9778 on T2I tasks (GenAI-Bench, UnifiedReward on DrawBench), and 0.9739 and 0.9866 on T2V tasks (GenAI-Bench, VideoAlign on TA-hard), all with statistically significant p-values. This is the paper's central claim: you can rank encoders without training generators.
• Decoder-only models beat encoder-only models. UMT5-XXL (about 5B parameters) reaches only 51.41 with multi-layer aggregation, below the much smaller Qwen3-VL-4B-Instruct at 55.20. The gap is architectural, not just about scale.
• Multimodal training helps even for a text-only test. MLLMs consistently outperform their text-only LLM backbones (for example, Qwen3-VL-8B-Instruct at 56.81 versus Qwen3-8B-Instruct at 55.77 with Norm-Avg), suggesting multimodal pretraining pushes text representations toward visually grounded semantics.
• Instruction tuning has an unclear effect. Unlike the multimodal trend, instruction-tuned versus base versions show no consistent advantage.
• "Thinking" models are double-edged. Reasoning models improve single-layer features but hurt when multiple layers are aggregated, likely because their prefill hidden states shift during reasoning.
• Scaling only shows up with multi-layer fusion. Single-layer features show no clear scaling trend with model size; the Norm-Avg strategy does.
• Aggregating layers beats picking one. Multi-layer normalization-and-averaging generally outperforms any single layer, since different layers encode different syntactic and semantic granularities.
• GRAN-TED reaches state of the art on TED-6K. A score of 57.42, outperforming all peers and closing the gap with 32B-parameter models while being an 8B model.
• Gains translate to generation. Compared to the Qwen3-VL-8B Norm-Avg baseline, GRAN-TED adds +1.24 GenAI and +0.007 UnifiedReward for T2I, and +2.39 GenAI and +0.136 VideoAlign for T2V.
• Semantic strengths and weaknesses. GRAN-TED improves on temporal (+2.67), action (+2.32), coreference (+1.39), OCR (+1.14), quantity (+0.86), event (+0.88), and spatial (+0.79) dimensions, but regresses slightly on adjectives (-1.17) and adverbs (-1.98). Fine-grained spatial detail and OCR remain the weakest categories in absolute terms.
• Freezing the layer weights is necessary. Training learnable layer weights continuously actually degrades performance (75.94 vs. 76.17 for fixed Norm-Avg). The two-step strategy—train weights jointly for a while, then freeze them—recovers and surpasses the baseline at 77.01. The authors visualize the weights drifting non-monotonically even after 200k steps, confirming that a moving text condition destabilizes diffusion training.
• Question-answering is a bad proxy. Direct QA scoring saturates near 97.9% across all encoders, offering no discrimination. Similarly, mean pooling helps encoder-only models but badly hurts decoder-only LLMs (Qwen3-8B-Instruct drops from 53.62 to 39.18).
• The benchmark is stable and context-dependent. Repeated aggregator training varies scores by at most 0.02, and shuffling captions away from their statements drops accuracy to roughly random guessing (27-28% versus 25% for four options).

Methodology in Plain English

The work proceeds in two halves: build a ruler, then use it to build a better model.

Building the ruler (TED-6K). The authors gathered high-quality images and videos from existing diverse datasets such as DREAM1K, CAPability, V-STaR, and CapArena. For each, Gemini 2.5 Pro wrote a richly detailed caption. From each caption, they generated statements testing nine semantic dimensions—one positive statement that agrees with the caption, plus several deliberately contradicting "hard negatives" created by swapping attributes, reversing relations, or changing counts. Human annotators verified that positives matched the source and that negatives were genuinely confusable rather than obviously wrong.

Making comparisons fair. You cannot just feed a prompt through an encoder and look at the output vector, because diffusion models consume a sequence of token-level hidden states, not a single pooled vector. So the authors train a small "aggregator"—a couple of attention layers with one extra learnable token—that reads the encoder's hidden states and compresses them into one context embedding, exactly the shape a diffusion transformer would consume. It is trained with a contrastive objective on 500,000 caption pairs (two different captions of the same image, generated by Qwen3-VL-235A22B-Instruct), learning to pull together embeddings of the same image and push apart different ones. The encoder itself stays frozen; only this small adapter trains. To score an encoder on TED-6K, you embed the source caption and all candidate statements, then check whether the positive statement is the most similar one. No image generation is involved.

Building the better encoder. Two stages. First, take Qwen3-VL-8B-Instruct—chosen because the benchmark showed MLLMs to be the strongest backbone—and fine-tune it on a large collection of image/video captions and VQA pairs focused on generation-relevant properties like object attributes, spatial layout, and temporal order. Second, rather than using a fixed rule to merge the model's internal layers, learn a softmax-weighted sum over layer-normalized hidden states. To keep the text conditioning stable while the diffusion model trains, these weights are trained only for an initial phase and then frozen. The total number of extra parameters equals the number of transformer layers.

Verification. The authors pick a subset of encoders, train matching T2I and T2V models with each, and compute correlations between TED-6K scores and standard generation metrics to prove the benchmark is predictive, then ablate the training strategy and evaluation design.

Why This Matters

Text encoders were largely treated as a settled detail in diffusion model design, even though prompt-following failures—wrong object counts, swapped attributes, ignored spatial relations—trace back to them. This work makes encoder quality measurable and improvable, and it reframes encoder selection as a cheap upstream decision rather than a trial-and-error side effect of expensive end-to-end training. The 750x cost reduction turns a multi-day experiment into a coffee break, which changes how quickly the field can iterate on this component.

Real-world applications:

• Image and video generation products: Faster, better-aligned text encoders directly reduce the frequency of prompts that produce wrong counts, misbound attributes, or ignored relationships—the most common user complaints.
• Content creation and advertising: Marketing teams generating on-brand visuals and short video ads benefit from more faithful adherence to detailed, attribute-heavy briefs.
• Storyboarding and previsualization: Video generation that correctly handles temporal ordering and coreference ("the character who entered first now raises the item she picked up") is far more useful for narrative planning.
• Model selection and procurement: Teams choosing between open and proprietary encoders can benchmark candidates on their own domain in minutes instead of retraining generators.

Industry relevance: The work is a collaboration involving Peking University, UCAS, Xi'an Jiaotong University, and the Kling Team at Kuaishou Technology, indicating direct applicability to commercial video generation. Any organization building or fine-tuning diffusion pipelines—generative media platforms, design tool vendors, cloud AI providers—could adopt TED-6K as an internal evaluation gate before committing GPU budget.

Future Directions

• Fixing the fine-grained weaknesses. GRAN-TED regressed on adjectives and adverbs and remains weakest on spatial relationships and OCR. Targeted data and objectives for these dimensions are the obvious next target.
• Stage-dependent conditioning. The paper notes that early denoising focuses on global structure while later steps refine detail, yet freezes a single set of layer weights. A principled, stable way to vary text conditioning across denoising timesteps could yield further gains.
• Better ways to handle non-stationarity. The two-step strategy is a pragmatic patch. A more theoretically grounded approach—such as weight schedules or alternating optimization with convergence guarantees—could remove the guesswork about when to freeze.
• Understanding why multimodal pretraining helps. The advantage of MLLMs over LLM backbones on a purely text-only test is demonstrated but not explained mechanistically; a clearer account could guide future pretraining.
• Extending scope. Larger backbones, more languages, longer video narratives, and multilingual or culturally specific prompts are untested, as is whether the benchmark transfers to domains outside natural images and video.

Target Audience

Diffusion model and multimodal researchers who need a principled way to compare text encoders; ML engineers at generative media companies deciding which encoder to build on; and graduate students entering the T2I/T2V area who want a compact, well-instrumented case study in evaluation-driven model development. Readers should be comfortable with transformer architectures, contrastive learning, and the basic mechanics of diffusion models to get full value; the benchmark and aggregator design are accessible on their own.

Authors’ abstract

The text encoder is a critical component of text-to-image and text-to-video diffusion models, fundamentally determining the semantic fidelity of the generated content. However, its development has been hindered by two major challenges: the lack of an efficient evaluation framework that reliably predicts downstream generation performance, and the difficulty of effectively adapting pretrained language models for visual synthesis. To address these issues, we introduce GRAN-TED, a paradigm to Generate Robust, Aligned, and Nuanced Text Embeddings for Diffusion models. Our contribution is twofold. First, we propose TED-6K, a novel text-only benchmark that enables efficient and robust assessment of an encoder's representational quality without requiring costly end-to-end model training. We demonstrate that performance on TED-6K, standardized via a lightweight, unified adapter, strongly correlates with an encoder's effectiveness in downstream generation tasks. Notably, under our experimental setup, compared with training a diffusion model from scratch, evaluating with TED-6K is about \textbf{750$\times$ faster}. Second, guided by this validated framework, we develop a superior text encoder using a novel two-stage training paradigm. This process involves an initial fine-tuning stage on a Multimodal Large Language Model for better visual representation, followed by a layer-wise weighting method to extract more nuanced and potent text features. Our experiments show that the resulting GRAN-TED encoder not only achieves state-of-the-art performance on TED-6K but also leads to demonstrable performance gains in text-to-image and text-to-video generation. Our TED-6K dataset and evaluation code are available at the following link: https://anonymous.4open.science/r/GRAN-TED-4FCC/.