💻 computer science

Twins: Learn to Predict Unified Representations with Focal Loss

The paper proposes "Twins," a unified continuous token space that concatenates ViT and VAE features, and addresses the resulting optimization imbalance in diffusion models by adapting a focal regression objective to significantly improve both image generation quality and multimodal understanding performance.

Original authors: Kaixiong Gong, Xin Cai, Bin Lin, Hao Wang, Yunlong Lin, Mingzhe Zheng, Bohao Li, Jian-Wei Zhang, Miles Yang, Zhao Zhong, Liefeng Bo, Xiangyu Yue

Published 2026-07-27

📖 3 min read☕ Coffee break read

CC BY 4.0

Original authors: Kaixiong Gong, Xin Cai, Bin Lin, Hao Wang, Yunlong Lin, Mingzhe Zheng, Bohao Li, Jian-Wei Zhang, Miles Yang, Zhao Zhong, Liefeng Bo, Xiangyu Yue

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to build a robot that can both see the world like a human and paint pictures like an artist. In the world of artificial intelligence, these two tasks usually require completely different toolkits. To "see" and understand an image (like knowing a picture shows a "golden retriever"), AI uses a brainy, high-level map that captures the big ideas but ignores tiny details like fur texture. To "paint" or generate a new image, the AI needs a different tool that holds onto every single pixel and fine detail, but often loses the big picture meaning. For a long time, scientists have been stuck trying to force these two different tools to work together, often having to choose between a robot that understands well but paints blurry pictures, or one that paints beautifully but doesn't know what it's drawing. This paper tackles that tricky balancing act, asking: can we create one single "language" for the robot that lets it do both perfectly at the same time?

The researchers behind this study, known as "Twins," decided to stop choosing between the two tools and instead glued them together. They took the "brainy" map (from a system called SigLIP) and the "detail-oriented" map (from a system called VAE) and stitched them side-by-side into one long strip of information. Think of it like giving your robot a pair of glasses where the left lens sees the big picture and the right lens sees the fine details, all in one view. However, when they tried to teach their AI model to use this new combined view, they hit a snag. The model was lazy; it loved the easy, big-picture part and completely ignored the hard, detailed part, resulting in generated images that were fuzzy and missing texture.

To fix this, the team discovered why the model was being lazy. They found that the "easy" part of the data was smooth and simple, while the "hard" part was full of complex, noisy details. The AI naturally preferred the smooth stuff, just like a student might skip the hard math problems to do the easy spelling ones first. To solve this, they invented a special training rule called "Focal Loss." Imagine a teacher who, instead of grading every question equally, gives extra credit points for getting the hardest questions right. This forced the AI to pay attention to the difficult, detailed parts of the image it was trying to learn.

The results were impressive. By using this new "Focal Loss" rule, the AI stopped ignoring the details. On a standard test of image quality, the new method improved the score by up to 10.57 points compared to the old way of training, without needing any extra tricks to guide it. The images it generated were sharp and clear, and the robot could still understand what it was drawing just as well as before. The paper suggests that this approach successfully breaks the old "impossible triangle" where you had to sacrifice one skill for another, proving that with the right training tricks, a single AI model can truly master both seeing and creating.

Technical Summary: Twins: Learn to Predict Unified Representations with Focal Loss

1. Problem Statement

Unified multimodal models aim to utilize a single model and a shared representation space to support both multimodal understanding and image generation. Existing approaches face a persistent trade-off, often described as an "impossible triangle" between Understanding, Reconstruction, and Generation:

Discrete Methods: Unify the interface via a shared codebook but often struggle with reconstruction fidelity or require complex tokenization.
Continuous Methods: Typically rely on two disparate representations: high-level semantic features (e.g., ViT) for understanding and low-level latents (e.g., VAE) for synthesis. This split creates mismatched latent spaces, forcing systems to perform extra decode–encode round trips to "understand" their own generations and breaking representational consistency.
Recent Attempts: Methods like UniFlow and UniLip attempt to unify representations by compressing high-dimensional CLIP features, but this dimensionality reduction compromises understanding capacity. Conversely, approaches like RAE that generate high-dimensional semantic embeddings (e.g., DINOv2) fail to reconstruct high-frequency details, resulting in poor image quality.

The core challenge is obtaining a single continuous representation that simultaneously supports strong semantic understanding, high-fidelity reconstruction, and generation-friendly predictability without increasing computational overhead (e.g., attention costs).

2. Methodology

2.1. The Twins Representation

The authors propose Twins, a unified continuous token space formed by channel-wise concatenating features from two pre-trained encoders on the same token grid:

ViT Component: Semantic-rich features from SigLIP2, optimized for discriminative semantic reasoning.
VAE Component: Detail-preserving latents from Flux.2 VAE, optimized for pixel-level reconstruction fidelity.

By concatenating along the channel dimension rather than the sequence dimension, Twins maintains the original sequence length. This ensures that the quadratic attention cost ( $O(L^2)$ ) does not increase with the token count, preserving computational efficiency.

2.2. The Optimization Imbalance

When training a Diffusion Transformer (DiT) to predict the unified Twins representation, the authors observed a severe optimization imbalance: the model fits the ViT (SigLIP) component well but struggles to match the VAE latent distribution, leading to blurred images and poor FID scores.

The paper attributes this imbalance to three sources of heterogeneity:

Spectral Bias: ViT features are dominated by low-frequency signals, while VAE features contain significant high-frequency energy (texture, noise). Neural networks naturally prioritize learning low-frequency functions (Spectral Bias theory), causing the DiT to fit ViT features first and treat VAE details as difficult noise.
Intrinsic Dimensionality (ID): While SigLIP has a higher physical dimension (768) than the VAE (128), its intrinsic dimension is significantly lower (~15 vs. ~35). The low-ID SigLIP manifold is easier to learn, whereas the high-ID VAE manifold presents a more complex optimization landscape.
Conditional Dependency: SigLIP features are highly structured and condition-aligned (deterministic given a class label), whereas VAE features retain significant condition-independent uncertainty. The model naturally prioritizes the predictable, condition-aligned objective.

2.3. Focal Loss for Flow Matching

To address this imbalance, the authors adapt a Focal Loss strategy for the Flow Matching objective. Instead of using standard Mean Squared Error (MSE) which treats all dimensions equally, they introduce a re-weighting scheme for the VAE channels:

The loss function assigns higher penalties to "hard" residuals (large errors) on the VAE dimensions.
The weighting factor is defined as $w_i = |v_i - v_\theta(z, t)_i|^{2\gamma}$ , where $\gamma$ is set to 0.5.
This forces the model to focus on the difficult-to-learn high-frequency VAE features, preventing the optimization from stagnating in a local optimum dominated by the easier ViT features.

3. Key Contributions

Unified Representation (Twins): A simple, efficient channel-wise concatenation of ViT and VAE features that supports both understanding and generation without increasing sequence length or attention costs.
Analysis of Imbalance: A systematic identification and analysis of the optimization imbalance in joint modeling, attributing it to spectral bias, intrinsic dimensionality, and conditional dependency.
Focal Loss Adaptation: The application of a focal reweighting strategy to Flow Matching, which effectively mitigates the imbalance by upweighting hard VAE dimensions.
Performance Gains: Demonstration that this approach yields substantial generation improvements over naive MSE loss while maintaining competitive multimodal understanding performance.

4. Experimental Results

4.1. Reconstruction

On the ImageNet-1K validation set, Twins achieves state-of-the-art reconstruction metrics among unified tokenizers:

PSNR: 31.46
SSIM: 0.90
rFID: 0.11
Twins significantly outperforms RAE (PSNR 18.83) and matches the reconstruction quality of dedicated generative autoencoders like Wan2.2 and SD-VAE 3, proving that the unified representation retains pixel-perfect consistency.

4.2. Multimodal Understanding

When integrated into a VLM pipeline (using Qwen2.5-7B), Twins performs competitively against specialized encoders:

It generally outperforms the strong SigLIP2 baseline.
The inclusion of low-level VAE features yields improvements in fine-grained tasks such as GQA (64.93 vs. 64.54) and TQA (58.89 vs. 56.92), attributed to the preservation of texture and exact shape details often abstracted away by semantic-only encoders.

4.3. Image Generation

On ImageNet@256 and @512:

Without Guidance: The baseline Twins model with MSE loss shows poor FID (degraded compared to Flux.2 VAE baseline). Applying Focal Loss significantly improves the prediction of VAE features, reducing the FID from 14.41 (MSE) to 3.84 (Focal Loss) at 80 epochs.
With Guidance: Twins achieves a strong gFID of 1.59 (ImageNet@256) and 1.79 (ImageNet@512) with classifier-free guidance.
Trade-off Resolution: While RAE achieves a slightly lower FID, Twins significantly outperforms it in reconstruction quality (PSNR), successfully narrowing the gap between understanding-oriented and generation-oriented representations.

5. Significance and Claims

The paper claims that Twins successfully breaks the "impossible triangle" of visual tokenization by constructing a unified representation that explicitly fuses semantic-rich ViT features with detail-preserving VAE features.

The significance lies in:

Simplicity: It avoids complex architectural redesigns or additional residual encoders, relying instead on the direct concatenation of off-the-shelf encoders.
Balanced Optimization: It identifies and solves the specific failure mode where DiTs ignore high-frequency details in unified spaces, providing a generalizable solution (Focal Loss) for joint modeling of heterogeneous features.
Unified Interface: It enables a model to both comprehend and generate in the same token space, eliminating the need for extra decode–encode steps and ensuring representational consistency across tasks.

The authors conclude that while jointly modeling heterogeneous features is non-trivial, the proposed strategy establishes a new baseline for unified representation generation, enabling high-quality prediction of both semantic and fine-grained latents in a single framework.