EfficientViT: Cascaded Group Attention

Chapter 1: Memory-Bound Attention

Every operation on a GPU is either compute-bound (limited by how fast the arithmetic units can multiply-add) or memory-bound (limited by how fast data moves between high-bandwidth memory and the compute units). The deciding quantity is arithmetic intensity — the ratio of arithmetic operations to bytes of memory traffic.

If a kernel does a lot of math per byte it touches (high intensity), it is compute-bound and FLOPs predict its time. If it does little math per byte (low intensity), it is memory-bound and FLOPs are almost irrelevant — the time is set by memory bandwidth.

A big matrix multiply like Q·K⊤ is the high-intensity case: each output element reuses a whole row and column, so the bytes loaded are amortized over many multiply-adds. Matmuls are compute-bound and FLOPs do predict them. But the other operations in attention are the low-intensity case.

The hidden cost: everything that is not a matmul

Walk through one self-attention layer on a feature map of T tokens and C channels. Beyond the two matmuls (scores and the value mix), the layer also performs:

The softmax does only O(T²) arithmetic but reads and writes the entire O(T²) attention matrix — its arithmetic intensity is roughly 1, the textbook definition of memory-bound. Reshapes do essentially zero arithmetic but copy the whole tensor through memory. LayerNorm and the residual add are element-wise: one pass each way over the feature map for a handful of FLOPs per element.

Worked example: the intensity gap

Take a modest stage: T = 196 tokens (a 14×14 grid), C = 192 channels, fp16 (2 bytes/element). Compare the score matmul against the softmax that follows it.

Score matmul Q·K⊤: output is T×T = 196×196. Each entry is a dot product over C = 192 channels, so 2×192 FLOPs (one multiply, one add per channel). Total ≈ 196×196×384 ≈ 14.7M FLOPs. Memory: read Q and K (2×196×192×2 bytes ≈ 0.15 MB), write the [196,196] matrix in fp16 (≈ 0.077 MB) ≈ 0.23 MB. Intensity ≈ 14.7M / 0.23M ≈ 64 FLOPs/byte — comfortably compute-bound.

Softmax over the same matrix: roughly 5 FLOPs per entry (max, subtract, exp, sum, divide) × 196×196 ≈ 0.19M FLOPs. Memory: read the [196,196] matrix and write it back ≈ 2×0.077 MB ≈ 0.15 MB. Intensity ≈ 0.19M / 0.15M ≈ 1.3 FLOPs/byte — squarely memory-bound.

The matmul does ~50× more arithmetic per byte than the softmax. On hardware that can do ~100 FLOPs per byte of bandwidth, the matmul saturates the compute units while the softmax stalls waiting for memory — and that stall is pure latency the FLOP counter never sees.

Chapter 2: The Sandwich Layout

If attention is the expensive, memory-bound part of a block, the natural question is: how many attention layers do you actually need? A standard ViT block alternates one attention with one FFN, every layer, throughout the network. EfficientViT breaks that one-to-one rhythm.

The sandwich layout replaces the usual (attention, FFN) pair with a block that has a single self-attention layer sandwiched between N feed-forward layers on each side. Schematically, a sandwich block computes:

More precisely, each sub-layer is a residual update. Writing Φ^F for an FFN sub-layer and Φ^A for the self-attention sub-layer, a sandwich block applies:

The single attention layer does the expensive global token mixing once. The many FFN layers — which are pure pointwise matmuls, fully compute-bound and cheap per the roofline — do the channel mixing and the bulk of the representational work. Each FFN sub-layer is also wrapped with a depthwise convolution token interaction, so even between attention layers the model still gathers some local spatial context without paying for global attention.

Why this is a strict win

It looks like you are removing attention layers — won't accuracy collapse? The paper's ablation says no, and here is the intuition. Attention's job is global, long-range token mixing. But once a layer has mixed tokens globally, the next layer's attention sees a feature map that already encodes those long-range relationships; a second consecutive attention is largely redundant. The FFNs, by contrast, are doing complementary work — nonlinear channel mixing — and they are cheap. So replacing redundant attention with extra FFN capacity keeps the accuracy while removing the memory-bound tax.

Worked example: counting memory-bound passes

Suppose a stage needs 6 layers of "work." A standard ViT spends them as 3 (attn, FFN) pairs — 3 attention layers, each paying the reshape+softmax+norm overhead. An EfficientViT sandwich block with N = 2 FFNs per side spends the same 6 layers as 2 FFN + 1 attn + 2 FFN + (1 more FFN to reach 6) — 1 attention layer. Same depth, same FFN count roughly, but one third of the memory-bound attention passes. If each attention contributes, say, 0.6 ms of memory-bound overhead, that is 1.2 ms saved per stage with no accuracy cost.

python
import torch.nn as nn

class SandwichBlock(nn.Module):
    """One memory-bound attention wrapped in 2N cheap FFN layers."""
    def __init__(self, dim, n_ffn=2):
        super().__init__()
        # N FFN layers before the attention ...
        self.pre  = nn.ModuleList([FFNwithDWConv(dim) for _ in range(n_ffn)])
        # ... exactly ONE self-attention (the only memory-bound part) ...
        self.attn = CascadedGroupAttention(dim)
        # ... and N FFN layers after.
        self.post = nn.ModuleList([FFNwithDWConv(dim) for _ in range(n_ffn)])

    def forward(self, x):
        for f in self.pre:  x = f(x)   # cheap, compute-bound
        x = self.attn(x)                  # expensive, memory-bound — ONCE
        for f in self.post: x = f(x)   # cheap, compute-bound
        return x

Chapter 3: Head Redundancy

The sandwich layout reduces how often we run attention. The next two chapters reduce how much each attention wastes. The waste has a name: head redundancy.

Recall standard multi-head attention. The C-channel input X is projected into h heads. Every head receives the same full input X and computes its own Q, K, V from it. The hope is that the heads specialize — one tracks objects, one tracks texture, one tracks position. The reality, Liu et al. measured, is that many heads in the same layer learn to compute nearly the same attention map. They are redundant: you pay for h heads of compute and memory, but the effective diversity is much lower.

Measuring redundancy

How do you quantify "two heads are similar"? Compare their attention matrices. Head a produces A^(a) of shape [T,T]; head b produces A^(b). Flatten each into a vector and take the cosine similarity:

If sim(a,b) is close to 1, the two heads attend almost identically — the second head adds little. The paper visualizes the average pairwise head similarity across layers and finds it is high in vanilla MHSA: the heads cluster. The model is computing h attention maps but only a few of them are distinct.

Why redundancy happens

Two heads start from the same input X. Their only difference is their projection matrices W_Q, W_K. During training, nothing in the loss explicitly pushes the heads apart — gradient descent is perfectly happy to let several heads converge to a useful-but-similar pattern, because each one independently lowers the loss. There is no diversity pressure. So redundancy is the default, not the exception.

The redundancy ↔ diversity tradeoff

The cure cannot simply be "use fewer heads" — that throws away the multi-head benefit of attending to several relationship types. Nor can it be a diversity loss term — that adds training complexity and hyperparameters. EfficientViT's insight is architectural: give each head a different input. If head a and head b never see the same features, they cannot collapse to the same attention map. Diversity becomes structural, guaranteed by construction rather than coaxed by a loss.

That single idea — feed each head a different slice of the feature channels — is group attention, and it is the subject of Chapter 4. Chaining the groups so each head also refines the next is the cascade of Chapter 5.

Chapter 4: Group Attention

Standard multi-head attention feeds all C channels of the input to every head. Group attention does the opposite: it splits the C input channels into h disjoint groups, one per head, and feeds head j only its own group.

Let the input be X of shape [T, C]. Split the channel axis into h slices:

Head j computes its query, key, and value only from its own slice:

where each projection W_·^j now has shape [C/h, d] instead of [C, d]. Then each head runs ordinary scaled dot-product attention on its own slice:

and the heads are concatenated and projected back to C channels by W_O, exactly as in standard MHA.

Two wins at once

Win 1 — guaranteed diversity. Head 1 sees channels 1..C/h; head 2 sees channels C/h+1..2C/h. They process disjoint information, so they cannot collapse to the same attention map. The redundancy of Chapter 3 is eliminated by construction — no diversity loss required.

Win 2 — cheaper projections. In standard MHA, each head's W_Q is [C, d] — it reads all C channels. In group attention it is [C/h, d] — it reads only C/h. With h heads, the total Q-projection cost drops from h·C·d to h·(C/h)·d = C·d, an h× reduction in the QKV projection compute and parameters.

Worked example: the projection saving, by the numbers

Let C = 192, h = 4 heads, head dimension d = 16. Standard MHA: each head's W_Q is [192, 16] = 3,072 params; with 4 heads and 3 matrices (Q,K,V) that is 4×3×3,072 = 36,864 params for the QKV projections. Group attention: each head's W_Q is [48, 16] = 768 params (because C/h = 192/4 = 48); 4×3×768 = 9,216 params. That is a 4× reduction — exactly h× — with more head diversity, not less. The freed budget is reinvested elsewhere (Chapter 6).

python
import torch

# Group attention: split C channels into h disjoint groups
X = torch.randn(T, C)          # [tokens, channels]
groups = X.chunk(h, dim=-1)   # h slices, each [T, C/h]

outs = []
for j, Xj in enumerate(groups):
    Qj = Xj @ WQ[j]            # [T, C/h] @ [C/h, d] -> [T, d]
    Kj = Xj @ WK[j]
    Vj = Xj @ WV[j]
    A  = (Qj @ Kj.T / d**0.5).softmax(-1)
    outs.append(A @ Vj)          # head j sees ONLY its slice
out = torch.cat(outs, -1) @ WO  # concat heads, project back to C

Chapter 5: The Cascade

Group attention guarantees diversity, but it has a downside: by giving each head only C/h channels, each head's query is computed from a thinner, lower-capacity slice. A head that only ever sees 48 of 192 channels has a narrower view. EfficientViT recovers the lost capacity — and squeezes out even more diversity — with the cascade.

The idea, in one sentence: add the output of head j into the input of head j+1. Instead of all heads running independently on their own static slice, the heads run in sequence, and each later head sees its own raw slice plus the refined output of the head before it.

The cascade equation

Let X_j be head j's input slice (its group of C/h channels), as in Chapter 4. In group attention head j computes its query from X_j directly. In cascaded group attention, head j's input is augmented by adding the previous head's output:

Read it left to right. Head 1 runs on its own slice. Its output is added to head 2's slice before head 2 projects Q/K/V. Head 2's output is added to head 3's slice, and so on down the chain. The Concat-and-W_O step at the end is unchanged — all the head outputs are still concatenated and projected back to C.

Why the cascade is clever — three effects at once

1. It recovers capacity without recovering cost. Each head still projects from C/h channels (so the projection stays h× cheaper), but because head j's input now contains head j−1's output, the effective receptive field over channels grows as you go down the chain. Later heads implicitly draw on the work of all earlier heads. The paper notes this increases the network's depth (the cascade is a chain of dependent computations) without increasing its width or its projection cost.

2. It increases attention diversity even further. Group attention made the heads' inputs disjoint. The cascade makes them progressively different: head 3 sees something head 1 and head 2 never saw (their combined refinement), so its attention map is even less likely to echo theirs. Diversity compounds along the chain.

3. It is a parameter-free interaction. The cascade adds no new parameters — it is just an addition of an already-computed tensor. You get richer, deeper, more diverse heads for the cost of a few element-wise adds, which are negligible even on the memory-bound side.

python
import torch

def cascaded_group_attention(X, WQ, WK, WV, WO, h, d):
    """X: [T, C].  WQ/WK/WV: per-head [C/h, d].  WO: [h*d, C]."""
    groups = X.chunk(h, dim=-1)  # h slices of [T, C/h]
    prev = None
    outs = []
    for j in range(h):
        Xj = groups[j]
        if prev is not None:
            Xj = Xj + prev          # THE CASCADE: add previous head's output
        Qj = Xj @ WQ[j]; Kj = Xj @ WK[j]; Vj = Xj @ WV[j]
        A  = (Qj @ Kj.T / d**0.5).softmax(-1)
        oj = A @ Vj                  # [T, d]
        outs.append(oj)
        prev = oj                    # feed forward to the next head
    return torch.cat(outs, -1) @ WO   # [T, C]

Chapter 6: Parameter Reallocation

Group attention freed up an h× chunk of the QKV-projection budget. The cascade spent zero parameters. So EfficientViT has spare capacity — where should it go? The third design principle, parameter reallocation, answers this by moving parameters from where they are wasted to where they matter.

The authors used a structured pruning / sensitivity analysis (Taylor importance) to ask, layer by layer: which channels actually contribute to the output, and which are dead weight? The findings drive three reallocations:

The asymmetry between Q/K and V

This is the subtle, important insight. Q and K exist only to compute the attention scores Q·K⊤ — a [T,T] matrix of similarities. The score is a scalar per token-pair; you do not need a wide Q or K to produce a good similarity. So Q and K can live in a small dimension d with little accuracy loss. The value V, by contrast, is the information that actually gets aggregated and passed forward — its width directly bounds how much content each token can carry. So V should be wide.

Standard MHA ties them together: d_q = d_k = d_v = C/h. EfficientViT unties them — narrow Q/K, wide V — putting parameters where they buy accuracy and starving the parts that don't.

Worked example: where the parameters go

Take C = 192, h = 3 heads. Standard: d_q = d_k = d_v = 64 per head; the Q, K, V projections each cost C×(h·d) = 192×192 = 36,864 params, and a 4× FFN costs 2×C×4C = 2×192×768 ≈ 295K params. Reallocated: set d_q = d_k = 16 (narrow), d_v = 64 (wide); now Q and K each cost 192×48 = 9,216 params (4× less), V stays at 36,864, and a 2× FFN costs 2×192×384 ≈ 147K params (half). The total attention+FFN budget drops sharply, and the params that remain are concentrated in V — the high-impact component. Accuracy holds; latency falls.

python
# EfficientViT unties Q/K width from V width
# Standard MHA: d_q = d_k = d_v = C // h   (tied, uniform)
# EfficientViT: narrow Q/K, wide V         (reallocated)
C, h = 192, 3
d_qk = 16          # small: Q,K only need to produce a similarity score
d_v  = 64          # large: V carries the content that gets aggregated

WQ = nn.Linear(C // h, d_qk, bias=False)   # [64, 16] per head
WK = nn.Linear(C // h, d_qk, bias=False)
WV = nn.Linear(C // h, d_v,  bias=False)   # [64, 64] per head — wider
# FFN expansion ratio 2 instead of the usual 4
ffn = nn.Sequential(nn.Linear(C, 2 * C), nn.GELU(), nn.Linear(2 * C, C))

Chapter 7: Results & Ablations

EfficientViT assembles the three ideas — sandwich layout, cascaded group attention, parameter reallocation — into a family of models, EfficientViT-M0 through M5, built on a hierarchical (three-stage) backbone with overlapping-patch downsampling. The headline claim is a better speed–accuracy tradeoff than prior efficient models on both GPU and CPU.

The accuracy–throughput frontier

On ImageNet-1K, the EfficientViT-M family pushes the Pareto frontier of top-1 accuracy versus measured throughput. The largest model, EfficientViT-M5, reaches roughly 77–78% top-1 accuracy. Against strong efficient baselines — MobileNetV3, MobileViT, EfficientNet, and prior efficient ViTs — EfficientViT reports comparable or higher accuracy at substantially higher measured throughput on a V100 GPU and on an Intel CPU. Crucially, the comparison is on throughput, not FLOPs — the whole point of the paper is that throughput is what users feel.

The ablation: which idea earns its keep

The paper ablates the three components by starting from a vanilla efficient-ViT baseline and adding each idea in turn. The qualitative picture (each component contributes, none is redundant):

The cascade is the part that matters most for accuracy: turning redundant parallel heads into a diverse, dependent chain measurably raises top-1, because the model now extracts more distinct relationships from the same attention budget. The sandwich layout is the part that matters most for speed: it directly removes the memory-bound passes the profiling in Chapter 1 identified as the bottleneck. Parameter reallocation tightens both.

Beyond classification

The paper also evaluates the EfficientViT backbone on downstream tasks — object detection and instance segmentation — where it transfers favorably, indicating the efficiency gains are not specific to classification but properties of the backbone. The same memory-aware design that speeds up ImageNet inference speeds up detection too.

Chapter 8: CGA Explorer

This is the payoff chapter. The interactive below assembles the whole mechanism — channel groups, per-head attention, and the cascade — into one controllable system. You drive every dial the paper exposes and watch how redundancy, diversity, and the cascade chain respond in real time.

Six tokens flow through cascaded group attention. Choose the number of heads h (how finely the C channels are split), toggle the cascade on or off, and pick which head to inspect. The top panel shows each head's channel slice and the cascade arrows feeding outputs forward; the middle panel shows the selected head's attention map; the bottom panel shows the pairwise head-similarity bar — your live readout of whether the heads are diverse (low bars) or redundant (high bars).

Things to try, and what they reveal:

• Toggle the cascade off. The similarity bars climb — without the chain, the heads (which already have disjoint slices) drift toward similar maps. This is the diversity the cascade buys, made visible.
• Increase h. Each head's channel slice gets thinner. With the cascade on, the chain compensates by feeding richer inputs down the line; with it off, thin slices mean weaker, noisier heads.
• Inspect later heads. With the cascade on, head 3's and head 4's maps look least like the others — they sit at the end of the chain and have integrated the most prior refinement, so they are the most distinct.

Chapter 9: Limitations & Connections

EfficientViT is a precise, hardware-grounded contribution — and precisely because it is grounded in 2023-era hardware and a specific size regime, it has clear limits worth naming.

Limitations

• The cascade is sequential. The whole point of the cascade is a chain of dependent heads. That chain is, by construction, not parallel across heads within an attention layer — head j must wait for head j−1. For small h (the paper uses few heads per layer) this serialization is cheap, but it is a real dependency that pure parallel MHA avoids. The net win holds because each head is tiny and there are few of them, but the tradeoff is there.
• Hardware-specific tuning. "Memory-bound" is a statement about a particular memory-bandwidth-to-compute ratio. The sandwich ratios and reallocation widths were tuned against specific GPUs/CPUs of the era; a very different accelerator (huge bandwidth, or specialized attention kernels like FlashAttention that fuse the memory-bound passes) shifts the optimum.
• Small-model regime. EfficientViT targets the efficient/mobile regime (a few hundred million FLOPs, ~77–78% top-1). The redundancy and memory-bound arguments are sharpest there; at very large scale the budget tradeoffs look different, and the head-redundancy story interacts with phenomena that only appear in big models.
• Fixed group split. Channels are split into contiguous, equal groups. A learned or content-adaptive grouping might extract more diversity, but the paper keeps the split static for simplicity and speed.

How it connects to the rest of the field

EfficientViT sits at the intersection of three lineages — the Transformer itself, efficient/mobile vision, and the hardware-aware view of attention. Follow these to place it:

Two more threads worth pulling: the memory-bound attention problem is attacked from the kernel side by FlashAttention (fuse the softmax/reshape passes so they never hit slow memory) — a complementary cure to EfficientViT's architectural one. And the head-redundancy observation echoes the head-pruning literature (Michel et al., "Are Sixteen Heads Really Better Than One?"), which removed redundant heads after training; EfficientViT instead prevents redundancy by construction.

Cheat sheet

You now understand EfficientViT end to end: why FLOPs lie, why attention is memory-bound, how the sandwich layout removes the bottleneck, why heads are redundant, and how cascaded group attention turns that waste into diverse, deep, parameter-free representation. The same questions — "where does the time actually go, and what is the budget really buying?" — are the right questions to ask of any architecture you design.