FastVAR in Infinity: What the Code Actually ChangesFastVAR 在 Infinity 代码里到底改了什么

This note looks at the FastVAR implementation on top of Infinity rather than only the paper algorithm. The short version is: FastVAR does not train a new Infinity model. It swaps the text-to-image Transformer block, wraps self-attention, cross-attention, and FFN with token pruning and cached restoration, and skips the two largest scales in the Infinity inference loop.

The relevant code sits in FastVAR/Infinity/infinity/models. Compared with the original Infinity/infinity/models, two new files matter most: fastvar_basic.py and fastvar_utils.py. The first replaces the block implementation; the second implements pivotal token selection and cache-based restoration.

Where the Hook Is

• Original Infinity uses CrossAttnBlock. In the baseline, each text-to-image block runs full self-attention, full cross-attention, and full FFN on the complete token map.
• FastVAR imports FastVARCrossAttnBlock. The FastVAR version changes the model constructor so text-to-image blocks come from fastvar_basic.py instead of the original basic.py.
• The weights remain compatible. The module names still contain the same attention, cross-attention, FFN, normalization, and adaptive modulation pieces, so this is a post-training inference-time change rather than a retraining recipe.
• The inference loop also changes. In autoregressive_infer_cfg, the FastVAR version directly skips scales whose width is 48 or 64, implementing the paper's aggressive 100% pruning for the last two Infinity scales.

Token Selection

The selection rule is implemented in masked_previous_scale_cache. Given the current token map cur_x, the code reshapes it back to a spatial grid, applies global average pooling, and treats this average token as a low-frequency proxy.

Then it computes sum((cur_x - mean_x) ** 2) for every token and sorts scores descending. Tokens far from the global average are treated as high-frequency or detail-heavy tokens, which corresponds to edges, texture, object boundaries, and local changes.

The merge closure uses torch.gather to keep only selected token indices. That reduced sequence is what enters the expensive Transformer operations.

Where Pruning Is Enabled

• Only selected resolutions are pruned. compute_merge currently checks whether the current width is in [32, 40].
• The ratios are hard-coded. Width 32 uses pruning ratio 0.4; width 40 uses 0.5. In other words, 60% or 50% of tokens are kept for those scales.
• Only later layers are pruned. FastVARCrossAttnBlock enables pruning when layer_idx is in range(3, 28). Early layers and very late layers stay intact.
• Other scales use identity functions. When pruning is not enabled, compute_merge returns do_nothing for merge, unmerge, and index lookup, so the block behaves like the original path.

Cache Restoration

After a module runs on the reduced token set, FastVAR must restore a full 2D token grid. Otherwise the next layer, the VAE-side scale transition, and the image lattice would no longer match the expected shape.

The unmerge closure takes a cached token map from an earlier scale, reshapes it to B x C x H x W, interpolates it to the current resolution with area mode, flattens it back, and scatters the newly computed selected tokens into their original positions.

In this implementation, each block keeps three caches: previous_scale_cache_self_attn, previous_scale_cache_cross_attn, and previous_scale_cache_ffn. The cache source is fixed to 24 x 24 via cached_size = [24, 24]. When the current sequence length equals 24 * 24, the block saves the outputs for later restoration.

Applied to All Three Submodules

The important engineering detail is that FastVAR does not only prune self-attention. Inside FastVARCrossAttnBlock.forward, the same pattern is applied three times:

• Self-attention. Normalize the gathered tokens, run self-attention, then unmerge with the self-attention cache.
• Cross-attention. Gather the current visual tokens before text cross-attention, run cross-attention on fewer visual queries, then restore with the cross-attention cache.
• FFN. Gather again before the feed-forward network, run the MLP on fewer tokens, then restore with the FFN cache.

This is why the speedup is not limited to attention. FFN cost also drops because the dense MLP is evaluated on a shorter sequence.

RoPE Has to Follow the Pruned Indices

One easy-to-miss detail is positional encoding. If the visual sequence is gathered, the rotary position cache must be gathered with the same indices. Otherwise the selected tokens would carry the wrong 2D positions.

FastVAR introduces apply_rotary_emb_fastvar. It slices the original RoPE cache for the current scale and, when a pruning index function is active, gathers RoPE entries with the selected token indices before applying rotation to Q and K.

Last Two Scales Are Skipped

The most aggressive part is in the autoregressive loop. For Infinity, when the current scale width is 48 or 64, the FastVAR code executes continue. That means the Transformer does not run for those scales at all.

This matches the paper's Infinity schedule {40%, 50%, 100%, 100%}. The 32 and 40 scales are partially pruned inside the block. The 48 and 64 scales are treated as fully pruned, relying on accumulated VAE latent interpolation rather than another Transformer refinement step.

Why It Speeds Up

• Sequence length becomes smaller at costly scales. Attention and FFN no longer process all tokens at width 32 and 40.
• KV cache grows more slowly. Self-attention caches fewer current-scale keys and values when pruning is active.
• The largest scales are removed from Transformer compute. Skipping width 48 and 64 avoids exactly the stages that dominate runtime and memory.
• It remains compatible with FlashAttention. Token importance is feature-based rather than attention-map-based, so FlashAttention can stay enabled.

My Take

The code is more direct than the paper description suggests. It is not an adaptive system that learns a schedule online. It is a small set of concrete inference-time decisions: prune widths 32 and 40, cache 24 x 24 layer outputs, prune only middle layers, and skip widths 48 and 64.

That directness is also the strength. FastVAR uses the coarse-to-fine structure of Infinity as an engineering prior: preserve early structure, spend compute only on high-frequency late tokens, and let cache interpolation fill stable low-frequency regions. The method is therefore easy to graft onto an existing Infinity checkpoint, but the exact scale list, cache scale, and ratios are backbone-specific knobs rather than universal constants.