Twins: Learn to Predict Unified Representations with Focal Loss

AI summary

Overview

Research area: Computer vision, specifically unified multimodal models (UMMs) — the problem of finding a single visual tokenization and representation space that serves both image understanding and image generation.

Technical level: Intermediate. The paper assumes familiarity with Vision Transformers (ViT), Variational Autoencoders (VAE), Diffusion Transformers (DiT), and flow matching, but its central ideas — concatenating two feature streams and reweighting a regression loss — are conceptually simple.

Scope: The paper proposes the "Twins" representation (channel-wise concatenation of SigLIP2 ViT features and FLUX.2 VAE latents on the same token grid) plus a focal-regression loss for flow matching, and evaluates it on ImageNet reconstruction, ImageNet class-conditional generation, and multimodal understanding benchmarks.

What This Paper Is About

Unified multimodal models need one visual representation that both supports semantic understanding and enables high-fidelity image generation. Existing continuous approaches instead use two mismatched spaces — a ViT feature space for understanding and a VAE latent space for generation — forcing extra decode–encode round trips and breaking representational consistency. The paper's goal is to fuse both into a single continuous token space that a Diffusion Transformer can learn to predict, and to fix the training imbalance that arises when it tries.

Key Contributions

1. A simple unified representation. The authors form a shared continuous token space, Twins, by channel-wise concatenating SigLIP2 ViT features with FLUX.2 VAE latents on the same token grid. Because fusion is in the channel dimension rather than the sequence dimension, the token sequence length is unchanged, so the quadratic attention cost with respect to token count does not increase.

2. Identification and systematic analysis of an optimization imbalance. When training a DiT to predict Twins jointly, the model fits the SigLIP2 component well but underfits the VAE component. The paper traces this to three sources of heterogeneity: spectral characteristics, intrinsic dimensionality, and condition-aligned versus condition-independent uncertainty.

3. A focal reweighting objective for flow matching. Inspired by Focal Loss, the authors adapt a feature-level focal regression loss that upweights large-error VAE dimensions, balancing optimization across the ViT and VAE components. The exponent γ is set to 0.5 in their experiments.

4. Empirical validation across three fronts. The paper reports generation gains over an MSE baseline, understanding performance comparable to or better than a strong SigLIP2 single-encoder baseline, and state-of-the-art reconstruction fidelity.

Main Findings

• Large generation gain from focal reweighting. On ImageNet, the focal regression objective yields up to 10.57 gFID gain over naive MSE loss without classifier-free guidance. In the ImageNet@256 setting, the Twins MSE baseline reaches gFID 14.41 at 80 epochs without guidance, while Twins with Focal Loss reaches gFID 3.84.

• The component-wise failure mode is real and quantified. In Figure 2, focal loss reduces the VAE FID by Δ = 5.36 at 80 epochs while maintaining performance on the SigLIP component.

• Spectral bias. FFT radial power spectra show SigLIP features are dominated by low-frequency content with rapid spectral decay, whereas FLUX VAE features retain significantly more energy in high-frequency bands, behaving as broadband signals. Under spectral bias theory, networks learn low-frequency functions first, so DiT fits SigLIP early and treats VAE high frequencies as difficult noise.

• A dimensionality paradox. Using Two-Nearest Neighbors estimation, SigLIP has a much higher physical dimension (D = 768) but a significantly lower class intrinsic dimension (ID_SigLIP ≈ 15) than the VAE (D = 128, ID_VAE ≈ 35). SigLIP's global ID (≈ 8.5) is lower than its single-class ID (≈ 14.7); the authors argue the single-class ID is a more faithful proxy for conditional DiT optimization difficulty. Features with higher ID (e.g., Flux.2) show significantly higher FID and slower convergence than lower-ID features (e.g., DINOv2-B).

• Conditional alignment versus condition-independent uncertainty. PCA cumulative explained variance collapses drastically for SigLIP under single-class conditions, indicating a highly structured, condition-aligned space. VAE dimensionality stays high under conditioning, indicating uncertainty statistically independent of the class condition.

• Focal loss also slightly improves SigLIP prediction. The authors attribute this to gradient balancing: down-weighting well-converged SigLIP gradients suppresses dominance of "easy" gradients, and modeling fine-grained VAE details forces the shared backbone to learn more robust multi-scale representations.

• Reconstruction results. On the 256×256 ImageNet-1K validation set, Twins achieves PSNR 31.46, SSIM 0.90, and rFID 0.11. For comparison, RAE reaches PSNR 18.83, SSIM 0.50, and rFID 0.57; Wan2.2 reaches PSNR 31.25; SD-VAE 3 reaches 31.29. UniFlow (InternViT) reports a higher PSNR of 33.23, SSIM 0.96, and rFID 0.26.

• Understanding results. With Qwen2.5-7B as the language backbone at resolution 384, Twins (SigLIP2-SO400M + Flux.2 VAE) reports POPE 87.82, GQA 64.93, TQA 58.89, MMB 77.00, MME-S 1971.0, and MME-P 1588.8. The SigLIP2-only baseline row reports POPE 87.15, GQA 64.54, TQA 56.92, MMB 77.00, MME-S 1826.8, and MME-P 1512.1 — the paper highlights the fine-grained gains on GQA (64.93 vs. 64.54) and TQA (58.89 vs. 56.92).

• Guidance comparison. With classifier-free guidance, Twins achieves gFID 1.59 on ImageNet@256. RAE reaches a lower gFID (1.13 at 800 epochs) but at a much poorer reconstruction PSNR of 18.83.

• Scaling trends hold at 512×512. On ImageNet@512, Twins with Focal Loss reaches gFID 3.78 without guidance and 1.79 with guidance, versus 6.80 for the MSE baseline without guidance; RAE reports 1.13 at 400 epochs. The authors state the trends are consistent with the 256×256 results.

Methodology in Plain English

The team starts from two off-the-shelf encoders that already work well: a SigLIP2 Vision Transformer, which produces semantically rich features aligned with language, and a FLUX.2 VAE, which preserves fine pixel-level detail. Instead of training a new tokenizer or aligning the two spaces, they simply glue the two feature vectors together along the channel dimension at every spatial token. Because both encoders use the same patch size and token grid, the number of tokens stays the same, so attention costs do not grow.

They then train a Diffusion Transformer (using flow matching with linear interpolation between data and noise) to predict these glued features as if they were a single data distribution. Training reveals a problem: the model learns the ViT half quickly but neglects the VAE half, producing blurred images.

Rather than redesigning the architecture, the authors diagnose why. They use FFT to show the ViT features are mostly low-frequency and the VAE features are broadband. They use a Two-Nearest-Neighbors estimator to show the ViT features actually live on a much lower-dimensional manifold than their nominal 768 dimensions suggest, while VAE features are intrinsically higher-dimensional and therefore harder to fit. They use PCA under class conditioning to show ViT features collapse onto a low-dimensional subspace given a label, while VAE features retain high uncertainty unrelated to the label.

Their fix is a focal-style reweighting of the flow-matching regression: each VAE channel's squared error is multiplied by the error raised to a power 2γ (with γ = 0.5), so channels the model is currently getting wrong receive larger gradients. This is applied to the VAE channel indices only. The result is that the model no longer takes the easy path of fitting only the semantic half.

For evaluation, they plug the Twins tokens into two pipelines: a DiT for class-conditional ImageNet generation (using the DDT head design from DDT and autoguidance as the guidance method), and a LLaVA-style VLM with Qwen2.5-7B as the language backbone, pretrained on LLaVA-558k and instruction-tuned on Cambrian-737k with LLaVA-1.5 settings.

Why This Matters

Impact on research. The paper reframes a practical training failure — a DiT ignoring part of its target representation — as a measurable property of the feature spaces themselves (spectral content, intrinsic dimension, conditional dependency). It offers a lightweight, training-free way to build a shared continuous space from existing encoders, and shows that a loss-level fix, rather than an architectural redesign, can resolve the imbalance. This is a different route than discrete codebook unification or joint-objective alignment used by prior unified tokenizers.

Real-world applications:

• Unified multimodal assistants that both describe and generate images without a separate decode–encode round trip between the understanding and generation stacks.
• Image editing and restoration pipelines that need pixel-accurate reconstruction (rFID 0.11, PSNR 31.46) while retaining semantic grounding.
• Class-conditional and text-conditional image synthesis systems that want a single latent space shared with a language model.
• Fine-grained visual question answering or detail-sensitive recognition tasks, where the paper reports gains on GQA and TQA from retaining low-level detail.

Industry relevance. The work comes from Tencent Hunyuan with academic collaborators, and code is released at github.com/Tencent-Hunyuan/Twins. Because it reuses widely adopted components (SigLIP2, the FLUX.2 VAE) and adds only a loss change, it is a low-friction upgrade path for teams already running ViT-plus-VAE dual-encoder unified models.

Future Directions

• Extending beyond class conditioning. The analysis is built on ImageNet class labels as the generation condition; whether the same spectral, intrinsic-dimension, and conditional-dependency story holds for long, free-form text prompts is not established in this paper.
• Closing the remaining gFID gap. RAE reports a lower FID (1.13 on ImageNet@256 with guidance; 1.13 on ImageNet@512 at 400 epochs) than Twins, so the trade-off between reconstruction fidelity and sample quality is not fully resolved.
• Beyond two encoders. The method concatenates exactly two feature streams; whether the focal reweighting scheme scales to three or more heterogeneous representations with different intrinsic dimensions is an open question.
• Task breadth. Understanding is evaluated only through a LLaVA-style pipeline on standard benchmarks; higher-resolution editing, video, and interleaved generation-understanding tasks are not tested.

Target Audience

Researchers and engineers working on unified multimodal models, visual tokenizers, and latent diffusion who want a compact diagnosis of why heterogeneous feature spaces are hard to co-model, plus a minimal loss-level intervention. It is also useful for practitioners building production systems that need one representation shared between understanding and generation, and for readers interested in applying intrinsic-dimension and spectral analysis as diagnostic tools for training dynamics. Readers without background in diffusion models or ViT/VAE tokenization will need to consult the cited prior work (RAE, UniFlow, FLUX, DiT) first.

Authors’ abstract

Unified multimodal models seek a shared visual token space that supports both multimodal understanding and image generation. Discrete methods unify the interface via a shared codebook, whereas continuous pipelines often rely on two disparate representations -- semantic features (e.g., ViT) for understanding and low-level latents (e.g., VAE) for synthesis -- resulting in mismatched latent spaces. We propose Twins, a unified continuous token space formed by channel-wise concatenating ViT and VAE features on the same token grid, so the sequence length is unchanged and attention cost does not increase. However, jointly modeling Twins in a Diffusion Transformer exposes a severe optimization imbalance: the model fits the ViT component well but struggles to match the VAE latent distribution. We trace this imbalance to three sources of heterogeneity: frequency bias, intrinsic dimensionality, and condition-aligned vs condition-independent uncertainty. To address it, we adapt a focal regression objective for flow matching that upweights large-error VAE dimensions, better balancing optimization across the ViT and VAE components. On ImageNet, this yields up to 10.57 gFID gain over naive MSE loss without classifier-free guidance. Twins also performs competitively on multimodal understanding benchmarks and improves reconstruction fidelity, narrowing the gap between understanding- and generation-oriented representations.