MaxViT

MaxViT combines a convolutional MBConv stage with two complementary sparse self-attention mechanisms — block attention and grid attention — to obtain global receptive fields at linear complexity in the number of tokens. The backbone is a four-stage pyramid, and each stage stacks repeats of the MaxViT block consisting of a MobileNetV3-style MBConv followed by block and then grid attention.

Given a feature map of shape $(H, W, C)$, block attention partitions it into non-overlapping $P \times P$ windows (default $P = 7$) and runs full self-attention within each window — identical to Swin's W-MSA. Grid attention takes the same feature map and reshapes it into a dilated $P \times P$ grid: tokens that share the same intra-window position across all windows are gathered into a single attention group. Formally, if we tile the spatial dimensions as $H = h \cdot P$, $W = w \cdot P$, block attention groups along the $(h, w)$ axis and grid attention groups along the $(P, P)$ axis:

Together the two operations realise dense local plus dense global self-attention, but each individual attention call has $\mathcal{O}(P^2)$ cost per token, giving overall $\mathcal{O}(HW)$ complexity. Variants Tiny / Small / Base / Large simply scale stage depths and channel widths along this shared multi-axis attention design.