Rethinking and Improving Relative Position Encoding for Vision Transformer

In plain words

A transformer that looks at an image cuts it into small squares and treats them as an unordered bag of tokens unless it is told where each one sits. The usual fix gives every square a fixed "address" from the top-left corner. This paper asks a different question: does it help more to tell the model how far apart two squares are, and in which direction? That is a relative position.

The proposed method, image RPE (iRPE), adds a small learned number to the attention score for every pair of squares, based on their offset in rows and columns. Offsets are grouped into a limited set of buckets. Nearby offsets get their own bucket, and far offsets share buckets more coarsely. Because the same tables can be read in many ways, the authors test several variants: whether the bias depends on the content of the tokens, whether the direction of an offset matters, and which of the query, key or value vectors receive the extra term.

On image classification the method improves a small vision transformer by about one point and needs very little extra compute. On object detection it helps the DETR detector too. The authors also find that relative encodings can replace the absolute address for classification, but not for detection, where the absolute address still matters.

TL;DR

iRPE adds a learned relative position term to every self-attention layer of a vision transformer, indexed by a bucket for each pair of patch offsets in rows and columns. The main variant, contextual Product, has an index function that is exact for small offsets and log-spaced beyond a threshold, with one bucket per (row, column) bucket pair, and it is shared across heads. Author-reported gains: DeiT-S top-1 79.9 to 80.9 (keys only) or 81.4 (keys, queries and values), and DETR AP 39.5 to 40.8 at 150 epochs. Read this if you add position encodings to a ViT or DETR; the bucket function and the weight-initialisation details matter for reproducing the numbers.

Problem

Self-attention is invariant to the order of its inputs, so position has to be injected somehow §1. Absolute encodings add one vector per position to the input once §2.2, eq. 4. Relative encodings add a term to the keys and values of every attention layer, indexed by the offset between two positions §2.2, eqs. 5–6. In language models relative encodings are established. In vision the paper reports that the evidence was contradictory: one prior work found no gain from relative over absolute encodings, another found a gain, and a third reported that relative encodings do not match absolute ones §1.

The existing relative methods were designed for 1D token sequences. Their clip functions give every offset beyond a threshold the same encoding, which loses long-range information, and they do not separate a horizontal offset from a vertical one §3.1, §3.2. Computing one term per pair of positions costs O(n² d) time for n patches with d channels, which becomes large at high resolution §3.2.

The experiments use DeiT-S with a 14×14 token grid for classification §4.1 and DETR with a ResNet-50 backbone for COCO detection [supp. Training and Test Settings of DETR].

Method

Background. Standard attention computes each output as a softmax-weighted sum of value vectors, with logits from scaled query–key dot products §2.1, eqs. 1–3. The relative formulation adds learned vectors to the values and to the queries and keys §2.2, eqs. 5–6.

Prior 1D methods, as analysed by the authors §3.1:

• Shaw et al. add learned vectors indexed by the clipped distance to keys and values; distances beyond a threshold k share one vector [eqs. 7–9].
• Transformer-XL uses a sinusoidal encoding of the distance projected by a learned matrix, plus two learned query-side vectors eq. 10.
• Huang et al. use one shared relative vector interacting with queries and keys, with a correction term eq. 11.
• SASA applies a 2D encoding that splits into horizontal and vertical 1D encodings, concatenated and added to keys only eq. 12.

Unified form of the proposed method §3.2. The attention logit becomes the scaled dot product plus a learned term b_{ij} eq. 13. Two modes are defined:

• Bias mode: b_{ij} is one learned scalar per bucket, independent of the input eq. 14.
• Contextual mode: b_{ij} is the query dotted with a learned vector per bucket eq. 15. Keys and queries can both carry such a term, each with its own table; the query-side table is dotted with the key eq. 16.
• Value term: the output adds a learned vector per bucket, weighted by the attention eq. 17. Only contextual mode is defined for values; bias mode is not.

The paper uses one index function for all mappings. It is exact up to α, then log-spaced up to γ, capped at β, and rounded eq. 18. The paper motivates it by keeping long-range offsets distinguishable while the clip function merges them §3.2. The index is an integer in [−β, β].

Four mappings from a 2D offset to a bucket §3.2:

• Euclidean (undirected): bucket from the Euclidean distance of the offset [eqs. 19–20]. The number of buckets is 2β+1 §3.2.
• Quantization (undirected): like Euclidean, but first maps the set of distances {0, 1, 1.41, 2, 2.24, …} to consecutive integers, so adjacent offsets such as (1, 0) and (1, 1) get different buckets eq. 21. The paper gives no formula for its bucket count; the supplement says it has more buckets than Euclidean [supp. Visualization of 2D Relative Position].
• Cross (directed): one bucket for the row offset and one for the column offset, with their encodings added [eqs. 22–24]. The number of buckets is 2×(2β+1) §3.2.
• Product (directed): one bucket per (row bucket, column bucket) pair, so the encoding is a single table entry eq. 25. The number of buckets is (2β+1)² §3.2.

Efficient implementation for contextual mode §3.2. The term query·bucket vector is computed once per bucket, which costs O(nkd) for k buckets, and then gathered into all n² pairs [eqs. 26–28]. The paper gives the total cost as O(nkd) instead of O(n²d).

Other design choices stated in the paper:

• The table can be shared across heads or separate per head. In bias mode sharing costs accuracy; in contextual mode the gap is negligible §4.2, tab. 2.
• The class token gets one extra bucket for all its pairs §4.1.
• The relative encoding is added to every self-attention layer; by default only to keys §4.1.
• The tables are initialised to zero [supp. Weight Initialization].

# Piecewise index g(x) for one axis offset  [eq. 18]
def piecewise(offset, alpha, beta, gamma):
    # Offsets up to alpha keep their own integer index  [eq. 18]
    if abs(offset) <= alpha:
        return round(offset)
    # Beyond alpha, log-spaced up to gamma, capped at beta  [eq. 18]
    log_share = log(abs(offset) / alpha) / log(gamma / alpha)
    spread = alpha + log_share * (beta - alpha)
    # Whole index; the code truncates toward zero
    # [code: iRPE/DeiT-with-iRPE/irpe.py#L40-L43]
    return int(sign(offset) * min(beta, round(spread)))

# Bucket id for every pair of positions, Product mode  [eq. 25]
# Input:  height, width of the token grid; alpha, beta, gamma
# Output: bucket ids (N, N), N = height * width
def product_bucket_ids(height, width, alpha, beta, gamma):
    # Row and column offset of every ordered pair of positions  [§3.2]
    offsets = pairwise_offsets(height, width)       # (N, N, 2)
    # Whole-number half-width of the table; beta is truncated
    # [code: iRPE/DeiT-with-iRPE/irpe.py#L194-L195]
    half = int(beta)
    span = 2 * half + 1
    # Map each axis offset to an index in [-half, half]  [eq. 18] [eq. 23]
    row_index = piecewise(offsets[..., 0], alpha, beta, gamma)
    col_index = piecewise(offsets[..., 1], alpha, beta, gamma)
    # Combine both axes into one id in [0, span * span)  [eq. 25]
    return (row_index + half) * span + (col_index + half)

# Key-side bias and logits, contextual mode  [eq. 13] [eq. 15]
# Input:  queries, keys (B, H, N, d); table (H or 1, d, K); bucket_ids (N, N)
# Output: logits (B, H, N, N)
def attention_logits(queries, keys, table, bucket_ids):
    # Dot every query with every bucket vector once  [eq. 26] [eq. 27]
    scores = queries @ table                        # (B, H, N, K)
    # Read off, for query i and key j, the score of bucket(i, j)  [eq. 28]
    bias = gather_buckets(scores, bucket_ids)       # (B, H, N, N)
    # Scaled content term plus the position term  [eq. 13]
    content = queries @ keys.transpose(-2, -1)      # (B, H, N, N)
    return (content + bias) / sqrt(queries.shape[-1])

Applying iRPE to a DeiT-style model, as in the code §4.1; code: iRPE/DeiT-with-iRPE/rpe_vision_transformer.py#L79:

1. For each attention layer, build one zero-initialised table per enabled position (queries, keys, values), with K buckets: (2β_int+1)² plus one extra for the class token.
2. For the token grid, compute the bucket id of every pair of patches once, and reuse it for all images with the same grid.
3. Compute the content logits as usual, and add the key-side bias from attention_logits, before dividing by √d eq. 13. A query-side term uses the same lookup with keys in place of queries and its own table eq. 16; code: iRPE/DeiT-with-iRPE/rpe_vision_transformer.py#L83.
4. If enabled, add the value term eq. 17 to the output, weighted by the attention.
5. Train with the DeiT recipe in the Hyperparameters section.

Notation. \tilde{x}_i, \tilde{y}_i are the x and y coordinates of patch i on the image plane §3.2. \beta_{int} is the integer part of β irpe.py.

Hyperparameters

Classification (DeiT), ImageNet §4.1. Values not stated in the paper are marked.

Object detection (DETR), COCO 2017 [supp. Training and Test Settings of DETR; code: iRPE/DETR-with-iRPE/README.md#L130]. The DETR settings are the original DETR settings, except RPE in the encoder.

Results

Author-reported. ImageNet top-1 for DeiT; COCO AP for DETR. The paper does not state the ImageNet split used. Setup: DeiT-S trained 300 epochs on 8 V100 GPUs §4.1.

Table 1 — mapping and mode, DeiT-S, ImageNet tab. 1. Δ is relative to the absolute-position DeiT-S.

Table 2 — shared versus unshared heads, DeiT-S, 50 buckets, three runs tab. 2.

Table 3 — clip versus piecewise index, DeiT-S, product shared-head, 50 buckets tab. 3.

Table 4 — component-wise, DeiT-S, contextual Product shared-head, 50 buckets tab. 4. Δ is relative to row 1.

Table 5 — ImageNet comparison tab. 5. Top-1 accuracy, parameters, input size and MACs as reported.

⁺ The paper applies its Product method with the clip function, shared across heads, to adapt the 1D encoding tab. 5 footnote; supp. Injecting Previous RPE Methods into DeiT. * DeiT-S with SASA's encoding tab. 5 footnote.

Table 6 — DETR on COCO 2017, component-wise tab. 6. Δ is AP relative to the row's epoch baseline (rows 1 and 8).

Rows 1 and 8 are the DETR baselines; row 1 is the official result, row 8 the official ablation tab. 6 comments.

Headline numbers (author-reported):

• ImageNet, keys only: DeiT-Ti 72.2 → 73.7; DeiT-S 79.9 → 80.9; DeiT-B 81.8 → 82.4 tab. 5; §4.3. Keys, queries and values on DeiT-S: 81.4 tab. 5.
• COCO: DETR 39.5 → 40.8 AP at 150 epochs; 40.6 → 42.3 AP at 300 epochs tab. 6.
• Swin-T (supplement): 80.5 → 81.9 top-1 with contextual QKV; bias mode alone gives +0.8 [supp. The Effectiveness on Other Vision Transformers].
• Fine-grained transfer, DeiT-B with keys: Stanford Cars 92.1 → 93.4; CUB-200 83.4 → 84.9 [supp. Transfer Learning on Fine-grained Datasets].
• Cost: at most 1% extra MACs with the efficient implementation §4.2 "Complexity Analysis"; fig. 4. The supplement's reference for runtime is DeiT-S without RPE at 1,096 images/s and 8,930 Mb fig. cost caption, supp. Inference Performance.

Findings the paper draws §4.2, §6: directed mappings beat undirected ones; contextual beats bias; head sharing costs little in contextual mode; the clip and piecewise functions differ little for classification but matter for detection tab. 3, tab. 6 rows 5–6; relative encodings replace absolute ones for classification but not for DETR tab. 4, tab. 6 rows 2–4; relative encodings push shallow layers toward local attention fig. 5.

Paper vs code

Code paths are relative to the repository root at code_commit. The repository is microsoft/Cream, and the iRPE code lives in iRPE/. The root LICENSE is MIT. Each iRPE sub-project (DeiT-with-iRPE, DETR-with-iRPE) carries its own Apache-2.0 LICENSE, which is the license recorded on this card.

Discrepancies:

• Quantization mapping. The paper maps distances to consecutive integers after taking square roots eq. 21. The code applies the piecewise index to the squared distance irpe.py. A helper that does rank quantisation exists irpe.py, but nothing calls it in the DeiT or DETR code; I checked by grep.
• Euclidean rounding order. The paper applies the index to the real-valued distance [eqs. 18–20]. The code rounds the distance to an integer first irpe.py. They can differ only for distances beyond α. I did not measure the effect.
• Bucket counts in the Euclidean, Quantization and Cross configs. Table 1 reports 20, 51 and 56 buckets tab. 1 caption. The configs with those names pass ratios 20, 33 and 20 rpe_models.py rpe_models.py rpe_models.py. By the code's formula, with β = 2r and bucket count 2·int(β)+1 irpe.py irpe.py, that gives 81, 133 and 81 buckets, plus one for the class token. I worked this out by hand and did not run it. The names do not match the bucket counts, so do not reuse the config names as bucket counts.
• Product and the ratio. The paper gives α : β : γ = 1 : 2 : 8 and says the bucket count is set by β §4.1, §3.2. The code scales all three by a ratio r, and the DeiT Product config uses r = 1.9 rpe_models.py. The DETR README uses r = 2.0 for 9×9 buckets README.md. The paper does not state r.
• Non-integer β. At r = 1.9, β = 3.8. The bucket count uses int(β) = 3 irpe.py irpe.py, so the Product table has 7 × 7 + 1 = 50 entries. The piecewise formula uses 3.8 as a float irpe.py.
• Weight decay on RPE tables. The supplement says no weight decay is applied to RPE weights [supp. Weight Initialization]. The repo's no_weight_decay() returns only pos_embed and cls_token rpe_vision_transformer.py. The optimizer is built by timm's create_optimizer main.py. I did not verify timm's grouping: timm is unpinned (timm>=0.3.2) requirements.txt, and timm is not installed here. The lookup tables are 2-D or 3-D, so timm's usual rule would apply weight decay to them. Treat this as unverified.
• Learning rate. The paper gives 1×10⁻³ §4.1. The code base LR is 5×10⁻⁴ and is scaled by batch × GPUs / 512 main.py main.py. This gives 1×10⁻³ at a total batch of 1024, the README setting of 8 GPUs × 128 README.md.
• Model zoo versus Table 5. The DeiT README's DeiT-S QK, QKV and DeiT-B K rows report 81.0, 81.2 and 82.3 README.md. Table 5 reports 81.1, 81.4 and 82.4 with the same MAC counts tab. 5. The Tiny-K and DeiT-S-K rows match.
• DETR model zoo versus Table 6. The README's 9×9 rows give 41.0 AP at 150 epochs and 42.2 at 300 README.md. Table 6 gives 40.8 and 42.3 tab. 6 rows 6, 9.
• Baselines and other experiments are not in the repo. Shaw's, Transformer-XL, Huang's and SASA's encodings are not implemented: rpe_models.py only defines iRPE configs. The Swin, fine-grained transfer and inference timing experiments have no code in the repo; it contains only DeiT and DETR.
• Settings not stated in the paper. Mixup and CutMix strengths, warmup start LR, and DETR gradient clipping come from code defaults, cited in the Hyperparameters section.
• Attention structure. The code's key term uses rpe_k(q) and the query term uses rpe_q(k * scale) rpe_vision_transformer.py rpe_vision_transformer.py. The value term is rpe_v(attn) rpe_vision_transformer.py. These match eqs. 15–17.
• Bias mode on values is rejected by the code irpe.py. This matches the paper, which defines values only in contextual form §3.2, eq. 17.

Pitfalls

• Bucket edges are sharp. I worked through eq. 18 with the code's constants for r = 1.9 (α = 1.9, β = 3.8, γ = 15.2). Offsets 2 and 3 both map to index 2, and every offset from 4 up to the grid's largest offset maps to index 3. Per axis, the 7 indices therefore cover offsets {0}, {1}, {2, 3} and {4, …}. This is by hand, not by running the code. With r = 2.0 (DETR 9×9), offsets 2 and 3 also share a bucket §3.2, eq. 18.
• Use β as the code does. Use int(β) for the table size, not β itself. A table sized from β = 3.8 would not match the 50 entries of the trained model.
• Rounding before the index. The Euclidean path rounds the distance before applying the index irpe.py. Matching the code is simpler than matching the paper's real-valued index if you want the reported numbers.
• Unverified weight decay. If you follow the paper and exclude the RPE tables from weight decay, you must add that yourself; the code does not.
• Bucket ids are cached per grid. The code keys its cache on (height, width, device) irpe.py, so a new resolution rebuilds them. Contextual mode costs O(B·H·N·d·K) in the matmul over all buckets irpe.py.
• Speed without the CUDA op. If the rpe_ops extension is not built, the code uses a PyTorch index path and prints a warning irpe.py. The README says its backward pass is very slow README.md. Build rpe_ops for training README.md.
• Run counts. Only Table 2 states a run count (three runs, spread ±0.04 to ±0.16). The paper does not state run counts for the other tables. Differences of 0.1 to 0.3 points in Tables 1 and 4 are the size of the spread Table 2 reports.
• Comparisons with prior 1D methods are adapted, not original. Table 5's Shaw, Transformer-XL and Huang rows use the authors' Product method with the clip function, so they are not the original implementations tab. 5 footnote.
• Head sharing has a cost in bias mode. Shared heads lose about half a point in bias mode tab. 2. In contextual mode, sharing saves about 0.2M parameters at equal MACs tab. 2; supp. Computation Complexity.
• Cost claims. The "at most 1%" claim is about extra MACs for DeiT-S across the resolutions in fig. 4 §4.2. The supplement's inference figure gives only the reference numbers [supp. Inference Performance].
• DETR batch size. The repo default is two images per GPU main.py. Only the 8-GPU setting gives the mini-batch of 16 in the paper.

Lineage

• Builds on: Shaw et al. (relative encodings with clipping), Transformer-XL, Huang et al., Axial-DeepLab (named in the paper, no arXiv id in the bibliography) §3.1; the piecewise index is inspired by T5 §3.2; DeiT as the baseline, arXiv:2012.12877 §4.1.
• Compared against: SASA (arXiv:1906.05909) §3.1, tab. 5 and CPVT (arXiv:2102.10882) tab. 5.
• Applied to: DETR (/papers/2005.12872) in §4.4 and Swin Transformer (/papers/2103.14030) in the supplement.