Twins: Unified Visual Tokens with Focal Loss for Balanced Multimodal Learning

Paper: 2607.22531 Authors: Kaixiong Gong, Xin Cai, Bin Lin, Hao Wang, Yunlong Lin, Mingzhe Zheng, Bohao Li, Jian-Wei Zhang, Miles Yang, Zhao Zhong Categories: cs.CV

The Gap

Multimodal vision models have split into two camps: understanding models (like CLIP/ViT) that learn high-level semantic features, and generation models (like Stable Diffusion) that work with low-level pixel latents from a VAE. The dream is a single model that does both — one shared token space that serves understanding and generation simultaneously.

Prior approaches took two paths. Discrete methods (like VQVAE-based models) share a codebook, but quantization bottlenecks limit expressiveness. Continuous methods avoid quantization but end up maintaining two separate feature streams — ViT embeddings for understanding, VAE latents for generation — which never really talk to each other. You’re essentially running two models in parallel, paying the compute cost twice.

This paper identifies a specific, previously under-discussed problem: when you naively concatenate ViT and VAE features and train them jointly with standard MSE loss in a diffusion framework, the optimization goes haywire. The model learns ViT features quickly but barely makes progress on VAE features. The authors trace this to three structural mismatches between the two feature types and propose a surgical fix.

Problem: Two separate feature spaces (ViT vs VAE)
         for understanding vs generation
                            |
                            v
Assumption: Concatenating them into one token
            space should "just work"
                            |
                            v
Reality:    Naive MSE training causes severe
            optimization imbalance
                            |
                            v
Root Cause: Frequency bias + dimensionality
            mismatch + uncertainty differences
                            |
                            v
Method:     Focal regression loss that upweights
            hard (VAE) dimensions
                            |
                            v
Evidence:   10.57 gFID improvement on ImageNet
            + competitive understanding benchmarks
                            |
                            v
Conclusion: Unified continuous tokens work IF you
            handle the heterogeneity properly

The Increment

One sentence: Before this paper, unifying understanding and generation features in a single continuous token space led to optimization collapse on the generation side; after, a focal loss that dynamically reweights hard dimensions makes joint training viable.

Core Mechanism

The architecture has two main stages. First, an input image is processed by two encoders in parallel: a ViT encoder produces semantic features capturing high-level structure and meaning, while a VAE encoder produces low-level latent features capturing texture and pixel details. These two feature maps are aligned on the same spatial grid and concatenated channel-wise, producing a single sequence of “Twins” tokens — each position carries both semantic and pixel information without increasing sequence length.

This unified token sequence then feeds into a Diffusion Transformer (DiT). During training, a flow matching objective teaches the model to denoise from a noisy version back to the clean Twins representation. The key innovation is in the loss: instead of MSE which treats all channels equally, they use focal regression. Channels where the model’s prediction error is large (typically the VAE dimensions) get upweighted, while channels the model already predicts well (typically ViT dimensions) get downweighted. This keeps the optimization gradient flowing to both components throughout training.

The authors identify three specific reasons why VAE dimensions are harder to learn:

1. Frequency bias: ViT features are smoothed by attention operations; VAE features contain high-frequency detail that’s harder to regress.
2. Intrinsic dimensionality: VAE latents have more degrees of freedom to capture; ViT features are more compressed and structured.
3. Uncertainty type: ViT features are condition-aligned (tightly coupled to the semantic content), while VAE features include condition-independent variability (random texture details not determined by the input).

Image
  |
  +--[ViT Encoder]--> semantic features (smooth, structured)
  |                          |
  |                    [Channel Concat]
  |                          |
  +--[VAE Encoder]--> pixel features (noisy, high-freq)
                               |
                               v
                    Twins Token Sequence
                    (same spatial grid,
                     doubled channels)
                               |
                               v
                    [Diffusion Transformer]
                               |
                               v
                    Flow Matching + Focal Loss
                    (upweight hard VAE dims)
                               |
                               v
                    Unified Representation
                    (understanding + generation)

Here’s a structural metaphor to make this concrete. Imagine you’re teaching two twins to play piano together as a duo. Twin A (ViT) has natural rhythm and picks up melodies quickly. Twin B (VAE) struggles with fine motor coordination and keeps hitting wrong notes.

If you give them equal practice time, Twin A gets bored and Twin B falls further behind. The piece sounds unbalanced — one side polished, the other sloppy.

The focal loss is like a smart practice coach who watches each measure and dynamically shifts attention. Measures where Twin B is struggling get repeated more. Measures Twin A already nails get less rehearsal time. The result: both twins improve at the rate needed for a cohesive performance.

Without this coaching strategy, you’d either: (a) abandon the duo idea entirely (the discrete approach), or (b) have them practice separately and hope they sync up on stage (the two-stream approach). Neither works well. The focal loss makes joint practice viable.

Key Concepts

• Flow Matching: Think of it like teaching someone to navigate from a random point in a field to a specific destination. Instead of showing them the full path at once (which is hard), you show them “at any location, here’s the direction to move.” Flow matching trains a neural network to predict this direction field. During inference, you start at a random point and follow the learned directions step by step until you arrive. It’s a cleaner alternative to diffusion’s noise schedule — continuous, mathematically elegant, and often faster.

• Focal Loss: Originally from object detection (Lin et al., 2017), the idea is simple: focus training effort on what the model gets wrong. Standard loss weights all examples equally. Focal loss multiplies each loss term by a factor that’s large when the model is wrong and small when it’s already right. It’s like a teacher who spends more time on problems students struggle with, not the ones they’ve mastered. Here, it’s adapted from classification to regression — the “hard” examples are specific feature dimensions with high prediction error.

• Channel-wise Concatenation vs Sequence Extension: This is a subtle but important design choice. You could unify ViT and VAE by doubling the sequence length (more tokens). But attention cost scales quadratically with sequence length, so this doubles your compute in a painful way. Concatenating along the channel dimension keeps the sequence length identical — each token is just wider. The Transformer still sees the same number of positions, and attention cost stays the same. It’s like packing two items into one box instead of using two boxes.

Framework Shift

Before (mainstream):                    After (this paper):

[Image]                                 [Image]
   |                                       |
   +->[ViT]->[Semantic Space]-->           +--[ViT Encoder]--+
   |     (understanding)                    |                  |
   |                                       |   [Concat channels]
   +->[VAE]->[Pixel Space]---->            |                  |
        (generation)                        +--[VAE Encoder]--+
                                                |
Separate representations.                   [Single Unified Space]
Cross-task transfer limited.                    |
                                        [Diffusion Transformer]
                                        [Focal Loss: balance dims]
                                                |
                                        Single model, both tasks

From two separate token spaces to one unified space, the core shift is recognizing that concatenation alone isn’t enough — you need a loss function that respects the statistical differences between the two feature types.

Expert Assessment

Problem choice: This is a real gap. The tension between understanding-oriented and generation-oriented representations has been a persistent friction point in multimodal modeling. Most work either picks one side or bolt two models together. The specific observation that joint optimization collapses on the VAE side is genuine and practically important — it’s the kind of thing that would make someone abandon the unified approach and conclude “it doesn’t work.” Identifying and fixing this is valuable.

Method maturity: The focal loss adaptation is clever but not revolutionary — it’s a well-known idea applied to a well-understood problem. The three-source decomposition of the imbalance (frequency, dimensionality, uncertainty) is the more interesting intellectual contribution, as it gives practitioners a framework for thinking about *why heterogeneous features resist joint training. The channel concatenation design is simple and effective. I’d call this “smart engineering guided by careful diagnosis” rather than a paradigm shift.

Experimental integrity: The 10.57 gFID improvement on ImageNet without classifier-free guidance is substantial and meaningful. The ablations likely show the contribution of each component. One concern: the paper benchmarks understanding and generation separately rather than showing emergent capabilities from unification (e.g., text-to-image conditioned on understanding features). This makes the “unified” claim somewhat incremental — it’s two tasks in one model, not something qualitatively new emerging from the union. The comparison baselines should be checked carefully: are they comparing against models of similar parameter count and compute?

Writing quality: The paper reads clearly and the problem diagnosis is well-structured. The weakness is in the evaluation narrative — the understanding benchmarks feel like an afterthought rather than a first-class contribution. A rewrite of the experiments section that explicitly tests “does unification enable new capabilities neither model could do alone?” would significantly strengthen the contribution claim.

Verdict: weak accept — The diagnosis of the optimization imbalance and the focal loss fix are genuinely useful contributions that will help others attempting unified representations. The work is solid but not transformative.

Takeaways

1. When fusing heterogeneous feature streams, watch for optimization imbalance. If one stream dominates gradients, your loss function is the first place to look. Focal-style reweighting is a lightweight fix worth trying before redesigning the architecture.

2. Channel concatenation as a compute-preserving fusion strategy. When you want to combine two feature types without inflating sequence length (and attention cost), concatenating along channels is a clean pattern. This transfers to any Transformer-based multimodal fusion.

3. Diagnose before you fix. The paper’s three-factor decomposition (frequency bias, intrinsic dimensionality, uncertainty type) is a useful mental model for anyone debugging why multi-task or multi-representation training fails. Not every problem needs a new architecture — sometimes the loss function is the bottleneck.