MViTv2: Improved Multiscale Vision Transformers

Chapter 1: Pooling Attention (MHPA)

Standard self-attention takes an input X of shape [T, D] (T tokens, D channels), projects it into queries, keys, and values, and computes:

The cost is dominated by the score matrix QK^T, which is [T, T]. If T is large, this is the bottleneck. MViT's core idea is simple to state: before computing attention, pool each of Q, K, and V along the spatial dimensions, shrinking their token count. This is Multi-Head Pooling Attention (MHPA).

Formally, MViT introduces a pooling operator P( · ; Θ) parameterized by a kernel size k, a stride s, and padding p — exactly like a conv/pool layer, but applied to a sequence reshaped back to a 2D (or 3D) grid. Each of Q, K, V gets its own pooling:

Here is the crucial asymmetry that drives the whole architecture. The query pooling stride s_Q sets the output resolution of the block: if you pool Q with stride 2, you get half as many output tokens. The key/value pooling stride s_KV sets the cost of attention without changing the output count: pooling K and V with stride 4 means the score matrix is [T_Q, T_KV] with a far smaller T_KV.

Worked example: where does each shape go?

Take a tiny grid: T = 16 tokens arranged 4×4, with d = 8 channels per head. Project to Q, K, V, each [16, 8]. Now pool:

Pool Q with kernel 1, stride 1 — no change, so Q̂ is [16, 8] (16 output tokens, the full 4×4 grid). Pool K and V with kernel 2, stride 2 — the 4×4 grid becomes 2×2, so K̂ and V̂ are [4, 8]. The score matrix Q̂K̂^T is therefore [16, 4] — 64 entries, not the 256 of full [16,16] attention. Softmax over the 4 pooled keys (each row sums to 1), then multiply by V̂ [4, 8] to get the output [16, 8].

So a query at every one of the 16 positions still produces an output, but it "sees" the image through only 4 pooled key/value summaries. We cut attention cost 4× while keeping full output resolution. That is the entire trick — and notice it cannot be done by pooling the block output, because then you would have computed full attention first and gained nothing.

What kind of pooling? MViT uses a learnable, channel-wise (depthwise) convolution as the pooling operator, followed by layer normalization, rather than plain max- or average-pooling. A depthwise conv of kernel k and stride s acts independently on each channel — it both downsamples and learns a small spatial filter, all with very few parameters (k×k per channel). This is strictly more expressive than max-pool while costing almost nothing.

python
import torch
import torch.nn as nn
import torch.nn.functional as F

def attention_pool(x, pool, hw):
    """x: [B, heads, T, d]  reshape to grid, depthwise-pool, flatten back.
       pool: a depthwise Conv2d(d, d, k, stride=s, groups=d).
       hw:   (H, W) of the current token grid (T == H*W)."""
    B, N, T, d = x.shape
    H, W = hw
    # [B, N, T, d] -> [B*N, d, H, W]  (channels-first 2D grid)
    g = x.reshape(B * N, H, W, d).permute(0, 3, 1, 2)
    g = pool(g)                          # depthwise conv, stride s
    Hp, Wp = g.shape[-2], g.shape[-1]    # pooled grid
    # back to [B, N, Hp*Wp, d]
    g = g.permute(0, 2, 3, 1).reshape(B, N, Hp * Wp, d)
    return g, (Hp, Wp)

def pooling_attention(Q, K, V, pQ, pK, pV, hw):
    Q, hwQ = attention_pool(Q, pQ, hw)   # s_Q sets OUTPUT resolution
    K, _   = attention_pool(K, pK, hw)   # s_KV sets attention COST
    V, _   = attention_pool(V, pV, hw)
    d = Q.size(-1)
    scores = (Q @ K.transpose(-2, -1)) / d ** 0.5   # [B,N,Tq,Tkv]
    A = scores.softmax(dim=-1)
    return A @ V, hwQ                  # output grid = hwQ

Chapter 2: The Stage Hierarchy

One pooling-attention block coarsens the keys but (with s_Q = 1) keeps output resolution fixed. To build a real pyramid we need the resolution to actually drop and the channels to actually widen as we go deeper. MViT organizes its blocks into stages — groups of blocks that all share the same resolution and channel width. At each stage transition, exactly one block uses query pooling (s_Q = 2) to halve the spatial side, and the channel dimension is doubled.

This is the ConvNet recipe, expressed in attention. ResNet has stages res2–res5; MViT has scale stages. Concretely, the MViTv2-Base (24-block) configuration looks like this:

Trace the spatial side: 56 → 28 → 14 → 7, halving three times. Trace the channels: 96 → 192 → 384 → 768, doubling three times. Trace the head count: 1 → 2 → 4 → 8 — MViT keeps the per-head channel width roughly constant (here 96) by doubling heads whenever it doubles channels. This is deliberate: a fixed per-head dimension keeps the √d scaling and attention statistics stable across stages.

Worked example: token budget across the pyramid

Start at scale 2 with 56×56 = 3136 tokens at 96 channels. The "size" of the activation is tokens × channels = 3136 × 96 ≈ 301k. After scale 3's transition: 784 tokens × 192 channels ≈ 150k. After scale 4: 196 × 384 ≈ 75k. After scale 5: 49 × 768 ≈ 38k. Each stage roughly halves the activation size — the pyramid is monotonically cheaper as it deepens, even as it gets semantically richer. Compare this to a flat ViT, where 196 × 768 ≈ 150k stays constant through all 12 blocks.

How does the input become 56×56 tokens? A patchify stem — a single overlapping conv with kernel 7, stride 4, padding 3 — turns the 224×224×3 image into a 56×56×96 feature map in one shot. (For video, a 3D conv also strides over time.) Overlapping patches, unlike ViT's non-overlapping 16×16 chop, give a smoother, convolution-like inductive bias at the very first layer, which MViT found helps optimization.

One more design choice worth naming: the class token. MViT keeps an optional learnable class token prepended to the sequence (as in ViT) for the classification head, but for detection it discards the class token entirely and feeds the four stages' feature maps directly into a Feature Pyramid Network. Because the architecture is natively multi-scale, those four stage outputs are already an FPN's P2–P5 — no extra resampling needed. We will return to that in Chapter 6.

Chapter 3: Q-Pooling, Stride, and the FLOP Budget

We have two pooling strides, s_Q and s_KV. This chapter makes their cost concrete so you can predict, and tune, the FLOP budget of any MViT configuration. This is where "pooling attention" stops being a slogan and becomes an engineering dial.

The dominant cost of attention is the two matmuls around the score matrix: Q̂K̂^T and A·V̂. With T_Q queries, T_KV keys/values, and per-head dimension d, each matmul costs roughly T_Q · T_KV · d. In a flat ViT, T_Q = T_KV = T, so the cost is T²d. In MViT:

So pooling Q by s_Q and K/V by s_KV divides the attention cost by s_Q² · s_KV². The squares matter: stride 2 on a 2D grid removes 3 of every 4 tokens.

Worked example: the first stage is the expensive one

Scale 2 runs at 56×56 = 3136 tokens. Full attention there would be 3136² ≈ 9.8 million score entries per head — ruinous. MViTv2 applies s_KV = 8 in scale 2, pooling the 56×56 grid down to 7×7 = 49 keys. Now the score matrix is 3136 × 49 ≈ 154k entries — a 64× reduction (8²). Crucially s_Q = 1, so all 3136 high-resolution queries survive. Every fine-grained query still attends globally, but to a 49-token summary.

MViT uses an adaptive s_KV: large in the high-resolution early stages (where it is most needed and least harmful), shrinking as the grid coarsens. A common schedule is s_KV = 8, 4, 2, 1 across scales 2–5: by scale 5 the grid is already 7×7 = 49, so no key pooling is needed at all and attention is full and global.

The pooled-Q reshape, by hand

When s_Q = 2 at a stage transition, the 28×28 grid pools to 14×14. Concretely the depthwise pool with kernel 3, stride 2, padding 1 maps output position (i, j) to input window rows 2i−1..2i+1 and cols 2j−1..2j+1. Overlapping (kernel > stride) windows mean adjacent output tokens share input context — another small convolution-like smoothing that MViTv2 found stabilizes training versus non-overlapping kernels.

A subtlety with the class token: it is not part of the spatial grid, so it cannot be pooled by a 2D conv. MViT handles it by splitting it off, pooling only the grid tokens, then concatenating the class token back. Forget this and the reshape from [T, d] to [H, W, d] fails because T = H·W + 1 is not a perfect grid. This is a real implementation gotcha.

python
def pool_with_cls(x, pool, hw, has_cls=True):
    # x: [B, heads, T, d]; T = H*W (+1 if class token present)
    if has_cls:
        cls_tok, x = x[:, :, :1], x[:, :, 1:]   # peel off [CLS]
    x, hw_pooled = attention_pool(x, pool, hw)  # 2D pool the grid only
    if has_cls:
        x = torch.cat([cls_tok, x], dim=2)      # stitch [CLS] back on
    return x, hw_pooled

Chapter 4: Decomposed Relative Position

Here is MViTv2's first headline improvement over v1. MViTv1 used absolute positional embeddings — one learned vector per absolute grid location, added to the tokens at the input. Two problems. First, absolute position does not transfer: an embedding trained for a 14×14 grid means nothing on a 28×28 detection grid. Second, vision is largely shift-invariant — a cat is a cat whether at the top-left or bottom-right — so what attention really wants is the relative offset between a query and a key, not their absolute coordinates.

The fix is a relative position embedding added inside the attention logits. For a query at grid position p and a key at position q, we add a learned bias R that depends only on the offset (p − q):

Now translate the whole image and every (p−q) offset is unchanged — the position bias is identical. Shift-invariant by construction, and it generalizes to any resolution because it is indexed by offset, not absolute coordinate.

Why "decomposed"? The naive cost is brutal

A 2D relative embedding has to cover every possible offset. On an H×W grid the offsets span (2H−1) in height and (2W−1) in width, so a full 2D table needs (2H−1)(2W−1) entries per channel. For a 56×56 grid that is 111×111 ≈ 12,321 vectors. Worse, computing the bias term Q_p·R_p−q for every (p, q) pair is itself O(T²d) — it re-introduces the exact quadratic cost pooling was meant to remove.

MViTv2's insight: decompose the 2D offset into its height and width components and learn two separate 1D tables:

R^h has (2H−1) entries, R^w has (2W−1) entries. For the 56×56 grid that is 111 + 111 = 222 vectors instead of 12,321 — a 55× reduction in parameters. And the decomposition lets the bias be computed with two cheap factored matmuls (one along height, one along width) instead of a dense O(T²) gather, bringing the relative-position overhead down to O(T·(H+W)d) — negligible.

Worked example: the offset index

On a 4×4 grid, height offsets Δh = p_h − q_h range over −3..+3 — that is 2×4−1 = 7 values. We store R^h as a length-7 table and index it by Δh + 3 (so −3 maps to row 0, 0 maps to row 3, +3 maps to row 6). Same for width. The bias for query (1,2) attending to key (3,0) uses R^h[1−3+3 = 1] + R^w[2−0+3 = 5]. Two small lookups, added. That is the entire mechanism.

python
def decomposed_rel_pos(attn, q, Rh, Rw, H, W):
    """attn: [B, heads, H*W, H*W] logits (pre-softmax, pre-scale)
       q:    [B, heads, H*W, d] queries
       Rh:   [2H-1, d] height table,  Rw: [2W-1, d] width table"""
    # Build index grids of offsets, shifted to be non-negative
    dh = (torch.arange(H)[:, None] - torch.arange(H)[None, :]) + (H - 1)
    dw = (torch.arange(W)[:, None] - torch.arange(W)[None, :]) + (W - 1)
    rh = Rh[dh]   # [H, H, d]  -- query-h vs key-h bias
    rw = Rw[dw]   # [W, W, d]  -- query-w vs key-w bias
    qg = q.reshape(*q.shape[:2], H, W, -1)  # [B,heads,H,W,d]
    # Factored einsums: height bias along rows, width bias along cols
    bh = torch.einsum('bnhwd,hkd->bnhwk', qg, rh)  # [B,n,H,W,H]
    bw = torch.einsum('bnhwd,wkd->bnhwk', qg, rw)  # [B,n,H,W,W]
    # add as separable bias to the [..,H*W,H*W] attention map
    attn = attn.reshape(*attn.shape[:2], H, W, H, W)
    attn = attn + bh[..., :, None] + bw[..., None, :]
    return attn.reshape(*attn.shape[:2], H * W, H * W)

Chapter 5: The Residual Pooling Connection

MViTv2's second headline fix is even smaller and even more important for training. Recall that the query Q is pooled to Q̂ before attention. After attention, the output has Q̂'s resolution. MViTv2 adds the pooled query Q̂ back into the attention output as a residual connection — inside the attention operator, before the usual block-level residual:

That extra "+ Q̂" is the residual pooling connection. It looks trivial. It is not. Here is why it matters.

The gradient and information argument

Pooling the query is a downsampling — a lossy, low-pass operation. Without the shortcut, the only path from a block's input to its output runs through the softmax-attention bottleneck, which itself only sees pooled keys. Information about the query's own fine structure can get washed out, and the gradient back to Q has to survive the softmax. Adding Q̂ gives a clean identity path: the block can, in the limit, behave like "pooled query, lightly adjusted by attention," which is a much easier function to learn early in training. It is precisely the ResNet identity-shortcut argument, applied to the pooling operator.

There is a second, quieter benefit. The bias R from Chapter 4 is added as Q·R, so it interacts multiplicatively with the query. The residual + Q̂ ensures the query signal stays strong enough through the stack for that relative-position interaction to remain meaningful at every depth. The two MViTv2 fixes are not independent — the residual pooling protects the query magnitude that the decomposed relative position relies on.

The full MViTv2 attention block, assembled

Putting Chapters 1–5 together, one MViTv2 block is: layer-norm, then pooling attention (with decomposed relative position bias and the + Q̂ residual pooling connection), then the block residual, then layer-norm, then an MLP, then another block residual. The two innovations live entirely inside the pooling-attention call.

python
class MViTBlock(nn.Module):
    def forward(self, x, hw):
        # --- pooling attention sub-block ---
        Q, K, V = self.qkv(self.norm1(x))
        Qp, hwQ = pool(Q, self.pQ, hw)   # s_Q sets output res
        Kp, _   = pool(K, self.pK, hw)   # s_KV cuts cost
        Vp, _   = pool(V, self.pV, hw)
        attn = Qp @ Kp.transpose(-2, -1) / self.scale
        attn = add_decomposed_rel_pos(attn, Qp, self.Rh, self.Rw, *hwQ)
        z = attn.softmax(-1) @ Vp
        z = z + Qp                       # RESIDUAL POOLING CONNECTION
        z = self.proj(z)
        # --- block residual (x may need pooling to match hwQ) ---
        x = pool_shortcut(x, hw, hwQ) + z
        # --- MLP sub-block ---
        x = x + self.mlp(self.norm2(x))
        return x, hwQ

Chapter 6: One Backbone for Detection & Video

The payoff of being natively multi-scale: the same MViTv2 backbone plugs straight into detection and video pipelines that a flat ViT cannot serve without surgery.

Detection: the pyramid IS the FPN

An object detector like Mask R-CNN needs a Feature Pyramid Network (FPN) — feature maps at strides 4, 8, 16, 32 (called P2, P3, P4, P5), each feeding a detection head tuned to a different object scale. A flat ViT produces only one stride (16), so adapting it (ViTDet) requires building an artificial pyramid by up/down-sampling that single map. MViTv2's four scale stages already are P2–P5: stride 4 (scale 2), stride 8 (scale 3), stride 16 (scale 4), stride 32 (scale 5). You wire the four stage outputs directly into the FPN. No interpolation, no hacks.

One more detection detail: at 1024×1024 input the early-stage grids are gigantic (256×256 at stride 4). MViTv2 uses Hybrid window attention for detection — most blocks attend within local windows (like Swin) to keep the high-resolution stages affordable, with a few global pooling-attention blocks per stage to maintain a global receptive field. This is a detection-only adaptation; for classification the pooling attention is global everywhere.

Video: pooling attention was born for the T-axis

Video adds a time axis, turning the grid into T×H×W. The number of tokens explodes — a 16-frame 224×224 clip at stride 4 is 8×56×56 ≈ 25,000 tokens. Full attention is hopeless; this is exactly the regime pooling attention was designed for. MViT simply makes its patchify stem and its pooling operator 3D (depthwise conv3d), pooling over time as well as space. The keys/values get pooled in all three dimensions, so even a long clip yields a small pooled key set.

The same architecture, with 3D pooling, set strong results on Kinetics-400/600/700 and Something-Something-v2 video classification — with far fewer FLOPs than the prior video transformers (which used either dense space-time attention or factorized approximations). The multiscale prior also matches video's natural structure: motion and fine texture early (high-res, short temporal range), actions and events late (low-res, long temporal range).

Chapter 7: Results & Ablations

What did the two small fixes actually buy? MViTv2's experiments are organized as a clean A/B: start from MViTv1, add decomposed relative position, add residual pooling, measure the delta on ImageNet, then scale to detection and video.

ImageNet classification

On ImageNet-1K, MViTv2 outperforms its v1 predecessor and matches or beats the strong Swin Transformer family at comparable FLOPs. The qualitative headline: MViTv2 reaches higher top-1 accuracy than Swin while keeping the simple, single-operator pooling design (no shifted windows, no relative-position lookup tables of the joint 2D kind). With ImageNet-21K pretraining, the large MViTv2 variant pushes into the high-80s top-1 — competitive with the best vision transformers of its time.

The ablation that justifies the paper

The central ablation adds the two components one at a time to the MViTv1 baseline:

Both components help independently and the gains stack — the residual pooling connection adds accuracy on top of relative position rather than overlapping with it. The paper also shows the decomposed (separable) relative position essentially matches a full joint 2D table on accuracy while being far cheaper, validating the factorization approximation from Chapter 4.

Detection and video

On COCO object detection and instance segmentation (with Mask R-CNN / Cascade Mask R-CNN), MViTv2 backbones set strong box and mask AP, beating comparably-sized Swin backbones, with the largest gains on small objects — the prediction the multiscale design makes. On Kinetics video classification, MViTv2 set state-of-the-art top-1 at substantially lower FLOPs than prior video transformers, and on Something-Something-v2 (a motion-heavy benchmark) it likewise led, confirming the architecture's temporal modeling.

Chapter 8: The Pyramid Explorer

This is the payoff. Build an MViTv2 pyramid live: choose how many stages, set the per-stage query and key/value strides, and watch the resolution coarsen, the channels widen, the token count fall, and the attention cost track your pooling choices — stage by stage. Read off every tensor shape and the running FLOP budget exactly as you would from a config table.

Things to try. Push the input side to 112 and the stages to 4: the first stage holds 112×112 = 12,544 tokens, and you will see why first-stage s_KV = 8 is essential — without it, the score matrix is over 150 million entries per head. Drop s_KV to 1 and watch the full-attention FLOP bar explode while the output resolution stays identical: that is the exact trade pooling attention exists to make. Reduce stages to 2 and the late, semantic-rich coarse maps disappear — you have a shallower, less hierarchical model.

Notice the invariant the builder enforces: at every stage transition the spatial side halves and the channels double, so the activation volume (tokens × channels) roughly halves per stage. That monotone decrease is the pyramid's efficiency signature — deeper means cheaper and richer, the opposite of a flat ViT where the volume is pinned constant. You are looking at the same energy balance that has made convolutional pyramids efficient for decades, now expressed entirely through pooled attention.

Chapter 9: Limitations & Connections

MViTv2 is a beautiful, practical architecture, but it is not the end of the story. Naming its limits sharpens understanding of what it does and does not solve.

Limitations

Pooling is lossy by design. Coarsening keys/values discards fine spatial frequencies; for the very early high-resolution stages this is the right trade, but it does mean a query can never attend to two distant fine details simultaneously at full resolution. Local-window methods (Swin) keep full resolution within a window but lose global reach; MViT keeps global reach but loses resolution. Neither has both for free — that is the fundamental resolution-vs-receptive-field tension of efficient attention.

The decomposition is an additive approximation. R^h + R^w cannot represent a truly non-separable 2D position bias (e.g. a diagonal-only preference). The paper shows this costs almost nothing in practice, but it is an assumption, not a free lunch.

Detection still needs an adaptation. The hybrid window attention used for high-resolution detection is a task-specific patch, not part of the pure architecture — a reminder that "one backbone, three tasks" still requires per-task attention-mode plumbing.

It is still O(T_Q·T_KV). Pooling reduces the constant dramatically but does not change the asymptotic form; for extreme resolutions you eventually want linear-attention or state-space approaches.

Cheat sheet

Where to go next

To ground the pieces this paper builds on and extends, study these companion lessons:

And these Gleams for the underlying mechanics:

• Vision Transformer (Gleam) — patchify and token attention from zero.
• Attention & Transformers (Gleam) — the Q/K/V machinery MViT pools.
• Convolutional Networks (Gleam) — the feature-pyramid / stride / channel-widening recipe MViT borrows.
• Pooling & Aggregation (Gleam) — what pooling does to a feature map, the operation at MViT's core.
• Detection & Segmentation (Gleam) — FPN heads and the small-vs-large object scale problem.