XCiT: Cross-Covariance Image Transformers

In plain words

A transformer for images looks at every small patch of the picture against every other patch. That works very well, but the work grows very fast as the picture gets bigger, so high-resolution images become slow and memory-hungry.

XCiT changes what the attention compares. Instead of asking which patches relate to which, it asks which feature channels relate to which, looking across all the patches at once. The table it builds has a fixed size, so it does not get bigger when the image has more patches. The network adds two small steps: a light convolution so that neighbouring patches can exchange information directly, and the usual per-patch feed-forward step. The authors' point is that this keeps the accuracy of a transformer, scales to large images more like a convolutional network, and copes better when the test image size differs from the training size.

TL;DR

XCiT is a vision transformer whose core operation, cross-covariance attention, computes attention between feature channels instead of between image patches. Its cost is linear in the number of patches, so high-resolution inputs stay affordable §1, §3.2. The authors describe the aim as combining transformer accuracy with the scalability of convolutional networks §1. They report accuracy on par with DeiT and CaiT in comparable settings, and better dense-prediction results than ResNet and several transformer backbones §1. Before implementing, read Paper vs code: the equation's softmax axis and temperature form differ from the released code, and the FPN adapter and several training details are only in the code.

Problem

Token self-attention compares all N tokens with each other. Its time cost is O(N²d) and its memory cost is O(hN² + Nd) for h heads §3.2, Complexity analysis. For a w×h image the paper writes this as O(w²h²) §1. This makes the method expensive for high-resolution images and dense tasks such as detection and segmentation §1.

Existing remedies are approximate attention and pyramidal designs. The authors argue that they either lose accuracy or stay too costly for very large images §1. Learned position embeddings must also be interpolated when the input resolution changes §3.3, Handling images of varying resolution.

Method

XCiT is a columnar network: it keeps the same spatial resolution (token count) across all layers §3.3. Each XCiT layer has three residual sub-blocks in this order: cross-covariance attention (XCA), local patch interaction (LPI) and a feed-forward network (FFN). Each sub-block is preceded by LayerNorm fig. 1 §3.3. Classification adds class-attention layers on top §3.3.

1. Token self-attention, the baseline. Tokens X (N × d) are projected to queries, keys and values. Attention is a softmax over Q Kᵀ / √d_k, then a weighted sum of V §3.1.
2. Cross-covariance instead of Gram matrix. The token-side matrix Q Kᵀ is N × N. The paper uses the d × d matrix Kᵀ Q instead, which can be computed in time linear in N §3.1. The motivation is the known relation between the Gram and covariance matrices of X §3.1.
3. Unit-normalise channels. Each channel (column of length N) of Q and K is scaled to unit ℓ2 norm over the tokens. Every entry of the d × d matrix then lies in [−1, 1] §3.2. The authors report that this norm control strongly stabilises training, especially with variable token counts §3.2.
4. Temperature. A learnable temperature τ scales the products before the softmax, so the distribution can be sharper or flatter. The paper's equation divides by τ §3.2 eq. 1. The code multiplies by a per-head parameter initialised to 1 (see Paper vs code).
5. Softmax and output. Attention weights come from the softmax of the cross-covariance. Each output feature is a convex combination of the value features in its group §3.2 eq. 1.
6. Heads. Channels are split into h groups, each with its own Q, K and V projections. Attention is block-diagonal across groups. This reduces the aggregation cost by a factor h and the authors found it easier to optimise and better in accuracy §3.2. Total cost is O(N d² / h) time and O(d² / h + N d) memory §3.2.
7. Local patch interaction (LPI). Tokens are reshaped onto their 2-D grid. Two depth-wise 3 × 3 convolutions with Batch Normalization and GELU in between are applied, then the tokens are reshaped back. This adds explicit communication between neighbouring patches, which XCA alone does not provide §3.3.
8. Feed-forward network. A point-wise MLP with one hidden layer of 4d units §3.3.
9. Patch embedding. The default patch size is 16 × 16. The paper also uses 8 × 8 patches, which it notes can improve performance §3.3. The embedding is a convolutional projection, as in LeViT, rather than a linear projection §3.3 tab. 4.
10. Positional encoding. A 2-D sinusoidal encoding is generated in a 64-d intermediate space and then linearly projected to d. It is added to the tokens and needs no interpolation when the image size changes §3.3 app. B.1.
11. Class attention for classification. Two layers (default in the code) let a CLS token read the patch tokens, one head at a time, in the CaiT style §3.3 xcit.py. The count of these layers is not stated in the paper.
12. LayerScale. Each residual branch is scaled by a learnable per-channel value, initialised to a model-specific ε §4.1 app. B.3 tab. B.1.
13. Head. The CLS token after the final LayerNorm feeds a linear classifier xcit.py.

The core operation, in the paper's notation eq. 1:

\text{XC-Attention}(Q, K, V) = V\,\mathcal{A}_{\text{XC}}(K, Q), \qquad \mathcal{A}_{\text{XC}}(K, Q) = \text{Softmax}\left(\hat{K}^\top \hat{Q} / \tau\right) eq. 1

Here Q̂ and K̂ are the channel-normalised queries and keys from step 3. The orientation of the softmax is ambiguous in the notation. The pseudocode below follows the released code, which is the reading consistent with the text's "convex combination" wording.

# Cross-covariance attention (XCA), one block  [eq. 1] [alg. 1]
# Input:  tokens (B, N, d)
# Output: tokens (B, N, d)

def xca(tokens, num_heads, temperature):
    # Project every token to queries, keys and values  [§3.2]
    queries = linear_q(tokens)                  # (B, N, d)
    keys = linear_k(tokens)                     # (B, N, d)
    values = linear_v(tokens)                   # (B, N, d)
    # Split the channels into heads, d / h channels each  [§3.2]
    queries = split_heads(queries, num_heads)   # (B, h, N, d/h)
    keys = split_heads(keys, num_heads)         # (B, h, N, d/h)
    values = split_heads(values, num_heads)     # (B, h, N, d/h)
    # Scale each channel to unit length over the N tokens  [§3.2]
    # [code: xcit.py#L247-L248]
    queries = l2_normalize(queries, axis=tokens)
    keys = l2_normalize(keys, axis=tokens)
    # Per-head d/h x d/h cross-covariance, times a learned temperature
    # [§3.2] [code: xcit.py#L250]
    scores = transpose(queries) @ keys * temperature   # (B, h, d/h, d/h)
    # Softmax over key channels, separately for each query channel
    # [§3.2] [code: xcit.py#L251]
    weights = softmax(scores, axis=key_channels)
    # Each output channel is a convex mix of its head's value channels
    # [§3.2]
    mixed = values @ transpose(weights)         # (B, h, N, d/h)
    return merge_heads(mixed)                   # (B, N, d)

# One XCiT block: XCA, then LPI, then FFN  [fig. 1] [§3.3]
# Input:  tokens (B, N, d); patch grid height and width
# Output: tokens (B, N, d)

def xcit_block(tokens, grid_height, grid_width):
    # Cross-covariance attention on the normalised tokens, with
    # LayerScale  [§3.3] [§4.1] [code: xcit.py#L288-L291]
    tokens = tokens + layerscale_attn * xca(layer_norm(tokens))
    # Depth-wise local patch interaction on the 2-D grid  [§3.3]
    tokens = tokens + layerscale_lpi * lpi(
        layer_norm(tokens), grid_height, grid_width
    )
    # Point-wise feed-forward network, 4d hidden units  [§3.3]
    tokens = tokens + layerscale_ffn * ffn(layer_norm(tokens))
    return tokens

Hyperparameters

Model configurations tab. 1. Depth, width d, heads h and parameter counts for the seven models in the paper. Top-1 is ImageNet-1k validation accuracy. "dist." means trained with hard distillation from a convolutional teacher §4.1. "↑" means fine-tuned and tested at 384 × 384 tab. 1 caption.

Stochastic depth and LayerScale ε tab. B.1 app. B.3. Per model, for patch sizes 8 and 16 unless stated. The paper calls the rate "drop rate" and ε the LayerScale initial value.

ImageNet-1k classification training §4.1.

• Stated: 400 epochs; AdamW optimiser; cosine learning-rate decay; the DeiT training recipe; LayerScale; stochastic depth adjusted per model (values above); images cropped with ratio 1.0 for evaluation, following CaiT §4.1.
• Not stated in the paper: batch size, base learning rate, weight decay, warm-up, label smoothing, mixup/cutmix and augmentation strengths (they are said to follow the DeiT recipe, which this paper does not restate).
• Defaults in the released training script are listed in Paper vs code; they are not the paper's stated values.

Object detection and instance segmentation §4.2 app. B.2.

• Mask R-CNN with FPN on COCO, 36 epochs (3× schedule) after ImageNet-1k pre-training.
• AdamW, learning rate 1e-4, weight decay 0.05, batch size 16.
• Multi-scale training and augmentation following DETR §4.2.
• Features are taken from layers [4, 6, 8, 12] for 12-layer models and [8, 12, 16, 24] for 24-layer models, and rescaled to strides 4, 8, 16 and 32 app. B.2.

Semantic segmentation (ADE20k) §4.3.

• Semantic FPN for 80k iterations and UperNet for 160k iterations.
• AdamW, learning rate 6e-5, weight decay 0.01, batch size 16, single-scale testing.
• Same FPN feature extraction as detection §4.3.

Self-supervised learning (DINO) tab. 3.

• 300 epochs, ImageNet-1k. Crop ratio 0.875 for consistency with DINO, except the last row, which uses 1.0 (80.7% → 80.9%) tab. 3 caption.

Results

Author-reported. Setup: ImageNet-1k (1.28M training, 50k validation, 1000 classes) §4.1; COCO (118k training, 5k validation, 80 classes) §4.2; ADE20k (20k training, 5k validation, 150 classes) §4.3. Memory and throughput were measured at batch size 64 on one V100-32GB fig. 2 tab. D.5. Models are trained on ImageNet-1k only unless stated.

ImageNet-1k and ImageNet-V2 top-1, supervised tab. 2. Columns: parameters, GFLOPs, test resolution, ImageNet-1k validation and ImageNet-V2. "dist." = distilled. "↑" = fine-tuned at 384 and tested at 384.

Self-supervised ImageNet-1k with DINO tab. 3. Linear and k-NN accuracy.

Ablations on XCiT-S12 with ImageNet-1k top-1 tab. 4.

The "Without XCA layer" row removes the XCA layer and reports a 6.1-point drop §4.1 analysis. The authors also state that the convolutional patch projection and LPI speed up training §4.1 analysis.

COCO object detection and instance segmentation, Mask R-CNN with FPN, 3× schedule tab. 5. Box AP (AP^b) and mask AP (AP^m) on mini-val. Backbones are pre-trained on ImageNet-1k. Numbers for Swin are taken from the Swin authors' code §4.2.

ADE20k semantic segmentation, single-scale mIoU tab. 6. Left block: Semantic FPN. Right block: UperNet. Params are the segmentation model's parameter count as listed in the paper. The paper notes that it excludes comparisons with models pre-trained on larger datasets tab. 6 caption.

The PVT-Tiny FPN parameter count is printed as "35.7M" for mIoU in the source; the value above is the number as written there.

Transfer learning after ImageNet-1k pre-training, top-1 tab. D.2. Models fine-tuned on each dataset.

More model sizes, with 16 and 8 patches at several test resolutions tab. D.1. The source labels the head-count column "#Blocks". For 16 × 16 patches the columns are GFLOPs at 224, top-1 at 224, top-1 at 224 with distillation, and top-1 at 384 (fine-tuned, distilled). The same four values for 8 × 8 patches follow the GFLOPs column.

Inference memory and throughput at batch size 64 on V100-32GB tab. D.5. Each cell shows images per second, then peak memory in MB fig. 2. OOM means out of memory.

Image retrieval (appendix) tab. D.4 app. D.3. Models pre-trained on ImageNet-1k only, using the CLS token, VLAD or ASMK aggregation over patch features. The best Holidays mAP in the table is 92.3 (XCiT-S12/8, ASMK, k=65536). On ROxford5k, the best medium score is 42.0 (XCiT-S12/8, ASMK, k=65536) and the best hard score is 15.0 (XCiT-S12/16, ASMK, k=65536). These are author-reported, and the table's default resolution is 768 tab. D.4 caption. The paper says the method is slightly below a ResNet-50 baseline on ROxford and better on Holidays app. D.3.

Robustness to test resolution fig. 3. XCiT-S12 keeps more accuracy than DeiT-S when the test resolution differs from the 224 training resolution, and benefits more from the test-time resolution increase effect (FixRes). The paper describes these curves in words and gives no table of values.

Paper vs code

Compared against the official repository facebookresearch/xcit at code_commit 82f5291f412604970c39a912586e008ec009cdca xcit.py. The paper's sections and tables are cited as in the arXiv version.

• Temperature form. Paper: the cross-covariance is divided by τ §3.2 eq. 1. Code: a per-head parameter initialised to 1 multiplies the scores [code: xcit.py#L230] xcit.py. Appendix Algorithm 1 also multiplies app. C alg. 1. The two forms are reparametrisations of each other; use the code's form.
• Softmax orientation. Paper: Softmax(K̂ᵀQ̂/τ) with output V·A eq. 1. Read literally with row-wise softmax, the weights for each output channel do not sum to one, which contradicts the paper's own "convex combination" wording §3.2. Code: attention is Q̂ᵀK̂ per head with softmax over key channels for each query channel, and output channel i mixes value channels j xcit.py. Appendix Algorithm 1: attention is K̂ᵀQ̂ with softmax over the query channel index app. C alg. 1, so the output channel index is the key channel. The two orientations are different functions. The released code is what produced the published checkpoints, so follow it.
• Positional encoding constants. Paper: sinusoidal, 2-D, 64-d intermediate, then projected app. B.1. Code matches this (32 dims per coordinate, 2π scale, temperature 10000, eps 1e-6) xcit.py. The base temperature 10000, scale and eps are not stated in the paper.
• FPN adapter for detection and segmentation. Paper: upsampling is a single transposed convolution, kernel 2, stride 2 app. B.2. Code, for 16 × 16 patch models, uses two stride-2 transposed convolutions with SyncBatchNorm and GELU in between to reach stride 4 xcit.py. The stride-8 level uses one transposed convolution xcit.py. For 8 × 8 patch models the code uses one transposed convolution for stride 4, an identity for stride 8, and max pooling for strides 16 and 32 [code: detection/backbone/xcit.py#L377-L390]. The segmentation backbone has a similar structure xcit.py.
• Local patch interaction order. Paper: "Batch Normalization and GELU in between", with no order given §3.3. Code: conv → GELU → BatchNorm → conv xcit.py.
• Patch embedding. Paper: convolutional projection "as in LeViT" §3.3. Code: four 3 × 3 stride-2 conv + BatchNorm blocks with GELU between for 16 × 16 patches, three for 8 × 8, and no activation after the last block xcit.py.
• Stochastic depth. Paper: per-model values tab. B.1. Code: one rate for every XCA block, the same in each layer xcit.py. Class-attention blocks are built without a drop-path rate, so they get none xcit.py.
• Token normalisation flag. The paper does not describe it. The code's tokens_norm is False for the nano models and True for the others [code: xcit.py#L415, L424, L433] xcit.py.
• Training recipe. Paper: DeiT recipe, AdamW, cosine decay, 400 epochs §4.1. Code defaults: batch size 64 per process, learning rate 5e-4 scaled linearly by batch × processes / 512 main.py, weight decay 0.05 main.py, warm-up 5 epochs main.py, label smoothing 0.1 main.py, mixup 0.8 and cutmix 1.0 main.py, repeated augmentation on by default main.py, and a teacher model regnety_160 for distillation, with distillation off by default [code: main.py#L135] main.py. The paper names no teacher model in the text. The README example trains XCiT-S12/16 with batch 128 on 8 GPUs and drop path 0.05 README.md, which the linear rule above turns into a learning rate of 1e-3. This is a derived value, not a paper value.
• Evaluation crop. Paper: crop ratio 1.0 §4.1. README: "By default we use the --full_crop flag" README.md. In the script, --full_crop is an opt-in flag with no default main.py. Pass it explicitly.
• Detection schedule. Paper: AdamW, lr 1e-4, weight decay 0.05, 36 epochs, batch 16 §4.2. Code matches, and adds step decay at epochs 27 and 33 [code: detection/configs/xcit/mask_rcnn_xcit_small_12_p16_3x_coco.py#L87-L92], no weight decay on normalisation layers same file, L87-L90, and a per-model drop path (0.05 for S12) same file, L24. The paper does not state these.
• Segmentation schedule. Paper: AdamW, lr 6e-5, weight decay 0.01, batch 16, 80k or 160k iterations §4.3. Code matches, and adds a linear warm-up of 1500 iterations, warm-up ratio 1e-6, and polynomial decay with power 1.0 [code: semantic_segmentation/configs/xcit/upernet/upernet_xcit_small_12_p16_160k_ade20k.py#L27-L36], plus drop path 0.05 for the small model same file, L17. The paper does not state these.
• Object detection comparisons. Swin numbers come from the Swin authors' code §4.2 footnote, not from a re-run by the XCiT authors.
• Class-attention count. The code default is 2 xcit.py. The paper does not give the number.

No discrepancies found for the XCA steps that the released code and the paper both state: ℓ2 normalisation over tokens, block-diagonal heads, linear complexity in N, and the block order XCA → LPI → FFN xcit.py xcit.py.

Pitfalls

• Softmax orientation is the main risk. The paper's equation is ambiguous and the appendix pseudocode differs from the released code (see Paper vs code). An implementation that follows the equation literally, with row-wise softmax, will not match the released checkpoints. Use the code's orientation and check one layer's attention output against it.
• Temperature. The paper gives no initial value for τ. The code initialises a multiplicative per-head parameter to 1 xcit.py. Weight decay on this parameter is not specified; the XCA module's no_weight_decay lists temperature xcit.py.
• ℓ2 normalisation epsilon. Not stated in the paper. The code calls PyTorch's normalize with default arguments xcit.py. Without normalisation the authors report failure for S12/16 tab. 4. They also say convergence without it was possible with hyperparameter tweaks §4.1, so the normalisation is a stability choice, not a hard requirement.
• Cost is linear in tokens, not free. Doubling the token count doubles the attention cost, but the FFN and LPI are also linear in tokens. 8 × 8 patches quadruple the token count of 16 × 16 at the same resolution, so 384-pixel 8 × 8 models cost many more GFLOPs. Compare S12/8 at 384 (55.6 GFLOPs) with S12/16 at 224 (4.8 GFLOPs) tab. 1.
• Block-diagonal heads. Each head's cross-covariance is d/h × d/h, so the number of heads sets the size of the attention map. The paper's Table 1 gives the head counts per model tab. 1.
• Resolution changes. The XCA softmax always runs over the same number of channels, so the attention map keeps its shape at any resolution §3.3. The positional encoding is computed on the fly for the current grid [code: xcit.py#L35-L55]. Neither needs interpolation.
• SyncBatchNorm. The code uses nn.SyncBatchNorm in the patch embedding, LPI and FPN layers xcit.py xcit.py. Single-process training needs a regular BatchNorm instead. The paper does not discuss this.
• Class-attention blocks. Their normalisation behaviour depends on tokens_norm: when it is False only the CLS token is normalised after the attention [code: xcit.py#L208-L211]. The paper does not mention this.
• Stochastic depth schedule. The paper gives per-model rates, not a per-layer schedule. The code uses one flat rate across XCA blocks and none in class-attention blocks xcit.py xcit.py.
• ImageNet hyperparameters. The paper names only the recipe. Batch size, base learning rate, weight decay, warm-up and augmentation strengths are not stated. The code defaults are the values to check before reproducing any number.
• Distillation. The teacher is not named in the paper. Distilled numbers in tab. 1 and tab. 2 depend on a convolutional teacher that the paper does not identify.
• Comparisons across papers. The COCO results are at 3× schedule after ImageNet-1k pre-training. The paper notes that larger results come from larger pre-training datasets or longer schedules and are not directly comparable §4.2.
• Figure and table numbers. These were derived by counting captions in the arXiv v2 LaTeX source, not checked against the compiled PDF. The appendix uses section letters (for example, Table B.1). The README cites Table B.1 and Figure 3, which agree with this count.

Lineage

• Builds on: token self-attention [Vaswani et al., 2017, cited in §3.1]; ViT (Dosovitskiy et al., 2020) for the columnar patch transformer §3.3; DeiT (Touvron et al., 2020) for the training recipe and hard distillation §4.1; CaiT (Touvron et al., arXiv:2103.17239) for class-attention layers and LayerScale §3.3 §4.1; LeViT (Graham et al., 2021) for the convolutional patch projection §3.3; Group Normalization for the idea of grouped channels §3.2; DINO (/papers/2104.14294) for the self-supervised setup §4.1 tab. 3.
• Compared against: ResNet, ResNeXt, EfficientNet, RegNet and NFNet for classification tab. 2; DeiT, CaiT and Swin (/papers/2103.14030) for classification and detection tab. 2 tab. 5; PVT, ViL and Swin for detection and segmentation tab. 5 tab. 6; Linformer, efficient attention and axial attention for approximate self-attention app. A.2 tab. A.1.
• Superseded by: not stated in this wiki.