VIOLIN: Spatial Priors via Space Filling Curves for ViTs

TL;DR

• A Vision Transformer's attention has no built-in notion of which patches are neighbors; only the positional embeddings carry that. Large models trained on large datasets learn locality anyway, but small models and small datasets suffer.
• VIOLIN borrows the decay mask from linear transformers: each attention score is multiplied, before the softmax, by \gamma^{d}, where d is the distance between the two patches along a scan of the image.
• A single scan order distorts 2D distance, so VIOLIN uses eight space-filling curves (Snake, Zig-zag, Hilbert and Peano, each with a transposed copy) and averages their masks. It adds 1,296 parameters to DeiT-B, 0.0015% of the model.
• Added at fine-tuning time, it raises VTAB-1K accuracy on all nine backbones tested, by up to 8.7 points on the Structured tasks that depend on spatial layout. In pretraining it adds 0.8 and 0.9 points on ImageNet for DeiT-T and DeiT-S, and 7.2 points for a DeiT-T trained on individual pixels of CIFAR-100.

Attention does not know which patches are neighbors

A ViT cuts the image into patches, flattens them into a sequence and runs self-attention over it. Attention is permutation equivariant: shuffle the patches and the outputs shuffle the same way. Nothing in the attention itself says that two patches are adjacent; that information comes only from the positional embeddings, and the model has to learn to use it.

CNNs get locality for free from their small kernels. With enough parameters and data a ViT learns it too, but for small models (the paper uses 30M parameters or fewer as the cut-off) or small datasets it does not learn it well. Fine-tuning a large pretrained ViT on a thousand examples is the second case.

A decay mask from linear transformers

Linear transformers such as RetNet order tokens with a decay factor \gamma instead of position embeddings: token i sees token j with weight \gamma^{i-j}. Written as a full, non-causal matrix this is the Kac-Murdock-Szegő mask, and VIOLIN multiplies it into the attention scores:

Distant pairs have their scores shrunk toward zero, so only nearby pairs can keep large scores. The mask does not block attention; it is a soft prior, much like the distance penalty in ALiBi but multiplied into the scores instead of added to them.

The distance |i - j| is measured in the flattened sequence, so it depends on how the image was flattened. Row-major order, which the paper calls the Z-curve and which most ViTs use, keeps horizontal neighbors one step apart but puts vertical neighbors a full row apart. It also places the last patch of one row right next to the first patch of the next, although they sit on opposite edges of the image.

Eight curves instead of one

A space-filling curve visits every cell of a grid exactly once. The paper uses four, each defined in its Appendix B.3 and code:

• Snake: row by row, reversing direction on every other row, so the end of one row is next to the start of the next.
• Zig-zag: along the anti-diagonals, alternating direction, the order JPEG uses for DCT coefficients.
• Hilbert: the recursive quadrant curve, built with the same generalized Hilbert ("gilbert") algorithm as the Hilbert-Guided Sparse Local Attention paper.
• Peano: the paper's name for the Morton or Z-order curve, which sorts patches by interleaving the bits of their row and column. It is a different curve from the row-major "Z-curve" above.

Each curve also has a transposed copy that swaps rows and columns, F_{c^\top}(i, j) = F_c(j, i), giving eight curves in all.

Running attention once per curve would multiply the cost by eight. Instead, VIOLIN computes Q, K and V once in row-major order and moves each curve's mask into that order with a permutation, \widetilde{M}_c = P_c^\top M_c P_c. One QK^\top then serves every curve:

Each curve's decay is learned per head as \gamma_c = \text{sigmoid}(\beta_c), and each head also learns the scale \alpha. That is L \cdot h \cdot (|\mathcal{C}| + 1) parameters: 12 \cdot 12 \cdot 9 = 1{,}296 for DeiT-B, 0.0015% of its 86M, and 0.64% more FLOPs.

Why average the curves

Every curve fails somewhere. The table below uses the demo's model on a 14×14 patch grid with \gamma = 0.9 and lists, over every reference patch, the weakest weight any curve gives to a patch that shares an edge with the reference, and the strongest weight it gives to a patch at least four steps away. These are this page's own calculations, not numbers from the paper.

Hilbert keeps patches inside a quadrant close, but two patches on either side of a quadrant boundary can end up more than 150 steps apart. Row-major and Peano both place some patches that are far apart in the image right next to each other in the sequence. The average cannot fall that far: every horizontal neighbor is adjacent on the Snake curve and every vertical neighbor on the transposed Snake, so each edge neighbor keeps at least \gamma / 8. The paper's Claim C.4 states the general version: no single curve moves the average by more than 1/m for m curves, and if a fraction p of the curves place two patches within r steps, the averaged mask is at least p\gamma^r.

The pretraining ablation (Table 12, DeiT-S on ImageNet, baseline 79.9%) points the same way. Peano alone gives the largest single-curve gain (+0.5), Hilbert alone gives none, and all four together give +0.8.

The learned decays tend to stay close to 1 (the paper's Figure 5). On 196 patches even \gamma = 0.9 gives a patch at the far end of the sequence, 195 steps away, a weight of 0.9^{195} \approx 10^{-9}, so values near 1 are what keep long-range attention alive. A curve whose \gamma_c drops low contributes little to the average, which lets each head and layer favor particular curves.

Implementation

The average gives each head a spatial prior without choosing one scan for every patch pair. To see exactly how it enters attention, follow a standalone PyTorch core from eight paths to one attention output.

Each row pairs an original-paper reference, our tensor explanation and an executable excerpt from the downloadable source. The four images are details cropped directly from arXiv v1. This example uses seeded synthetic Q, K and V to check the attention mathematics and gradients; it does not load a pretrained DeiT or DINO classifier or reproduce the paper's training results.

1. Choose eight paths through the patch grid

Original paper

Enlarge figure

Math and tensor shapes

N=H W,\quad i=yW+x

Assign each patch a row-major index i on a fixed H × W grid. Each path returns a permutation of those N indices: curve position maps to patch index. Its transpose traverses the swapped grid and maps back to the original coordinates.

The eight paths are Snake, JPEG zig-zag, generalized Gilbert/Hilbert and “Peano,” each in both orientations. “Peano” follows the paper's Appendix G.3 and official implementation: it means Morton bit-interleaving here, rather than the classical Peano curve.

Only geometry is cached. A module's grid and optional class-token configuration stay fixed; learned masks are recomputed on every forward pass.

Python implementation

python

@lru_cache(maxsize=128, typed=True)
def curve_order(curve, height, width, transposed=False):
    """Return curve position -> row-major patch index for any rectangle."""
    _positive_integer(height, "height")
    _positive_integer(width, "width")
    if curve not in ("snake", "zigzag", "hilbert", "peano"):
        raise ValueError("curve must be snake, zigzag, hilbert, or peano")
    if not isinstance(transposed, bool):
        raise TypeError("transposed must be a boolean")
    if transposed:
        swapped = curve_order(curve, width, height)
        order = tuple(
            (index % height) * width + index // height for index in swapped
        )
    elif curve == "snake":
        order = tuple(
            row * width + col
            for row in range(height)
            for col in (
                range(width) if row % 2 == 0 else reversed(range(width))
            )
        )
    elif curve == "zigzag":
        coordinates = [
            (row, col) for row in range(height) for col in range(width)
        ]
        coordinates.sort(
            key=lambda rc: (sum(rc), -rc[0] if sum(rc) % 2 == 0 else rc[0])
        )
        order = tuple(row * width + col for row, col in coordinates)
    elif curve == "peano":  # The authors' name for Morton bit-interleaving.
        order = tuple(
            sorted(range(height * width), key=lambda i: _morton_key(i, width))
        )
    else:  # Generalized Gilbert supports non-power-of-two maps.
        path = []
        if width >= height:
            _generate2d(0, 0, width, 0, 0, height, path)
        else:
            _generate2d(0, 0, 0, height, width, 0, path)
        order = tuple(row * width + col for col, row in path)
    if sorted(order) != list(range(height * width)):
        raise ValueError("Curve traversal must be a permutation of every patch")
    return order

2. Invert each path into row-major ranks

Original paper

Enlarge figure

Math and tensor shapes

r_c(o_c(t))=t

The forward order o_c says which patch a path visits at step t. Its inverse r_c says where a given row-major patch occurs on that path. Stacking all inverses gives ranks[8, N].

For a 2 × 3 JPEG zig-zag grid, the order is [0, 1, 3, 4, 2, 5] and its inverse is [0, 1, 4, 2, 3, 5]. Patches 3 and 2 occur at steps 2 and 4, so their curve distance is 2. This counts path steps, not their Euclidean distance in the image.

Using these inverse ranks indexes both axes of the mask in row-major order, implementing P_c^\top M_c P_c without dense permutation-matrix products. Q, K and V keep their original token order.

Python implementation

python

@lru_cache(maxsize=128, typed=True)
def curve_ranks(curve, height, width, transposed=False):
    """Invert the order: row-major patch index -> its curve position."""
    order = curve_order(curve, height, width, transposed)
    ranks = [0] * len(order)
    for position, patch in enumerate(order):
        ranks[patch] = position
    return tuple(ranks)

3. Learn eight decays and one scale per head

Original paper

Math and tensor shapes

\gamma_{h,c}=\sigma(\beta_{h,c})

beta[heads, 8] is learnable; sigmoid puts each decay between 0 and 1. alpha[heads] is a separate learned multiplier. The integer ranks are a registered buffer, so device moves preserve geometry without treating it as a parameter.

n_{\mathrm{added}}=9h\quad\text{per layer}

For DeiT-B's 12 layers and 12 heads this is 12 × 12 × 9 = 1,296 added parameters, excluding its existing QKV and output projections.

The example seeds beta uniformly in [5, 9], following the pretraining range, and deliberately starts alpha at 1 for teaching. Appendix D.1 · PDF page 24 specifies that beta range and the fine-tuning initialization: beta in [15, 20] and alpha centered at 1. The official pretraining initializer instead draws alpha from Normal(0, 0.5).

Python implementation

python

def __init__(self, height, width, heads, class_token=False, seed=0):
    super().__init__()
    _validate_configuration(height, width, heads)
    if not isinstance(class_token, bool):
        raise TypeError("class_token must be a boolean")
    if isinstance(seed, bool) or not isinstance(seed, int):
        raise TypeError("seed must be an integer")
    self.height, self.width, self.heads = height, width, heads
    self.class_token = class_token
    generator = torch.Generator(device="cpu").manual_seed(seed)
    self.beta = nn.Parameter(
        torch.empty(heads, 8).uniform_(5, 9, generator=generator)
    )
    self.alpha = nn.Parameter(torch.ones(heads))  # Teaching initialization.
    ranks = [
        curve_ranks(curve, height, width, transposed)
        for curve, transposed in CURVES
    ]
    self.register_buffer("ranks", torch.tensor(ranks, dtype=torch.long))

4. Form differentiable decay masks

Original paper

Math and tensor shapes

d_c(i,j)=|r_c(i)-r_c(j)|

K_{h,c}[i,j]=\gamma_{h,c}^{d_c(i,j)}

Subtracting the inverse ranks along two broadcast axes gives distances[8, N, N]. Broadcasting the decays over those axes produces K[heads, 8, N, N], symmetric masks whose diagonals equal 1.

In the preceding small-grid example, distance 2 and an illustrative decay of 0.9 give weight 0.9² = 0.81. The CLI instead uses its seeded learned initial values.

Exponentiation stays in autograd: every curve's beta can receive a gradient. These are dense masks with quadratic storage, not a sparse attention pattern.

Python implementation

python

def curve_masks(self):
    """Return differentiable [heads, eight curves, patches, patches]."""
    distances = (self.ranks[:, :, None] - self.ranks[:, None, :]).abs()
    gamma = torch.sigmoid(self.beta)
    # Ranks put each curve's KMS mask into the common row-major basis.
    return gamma[:, :, None, None] ** distances[None, :, :, :]

5. Average the masks before softmax

Original paper

Enlarge figure

Math and tensor shapes

M_h=\frac{1}{8}\sum_{c=1}^{8} K_{h,c}

mean(dim=1) removes the curve axis, leaving M[heads, N, N]. Every curve now contributes to one spatial prior for its head.

This averages the masks before a single softmax. Running eight separate attention operations and averaging their outputs would give a different result, because softmax is nonlinear. The executable checks compare the mask-first implementation with an independent reference.

Python implementation

python

def patch_mask(self):
    """Average the eight aligned masks, before scores or softmax."""
    return self.curve_masks().mean(dim=1)

6. Leave an optional leading class token unattenuated

Original paper

Math and tensor shapes

\widehat M_h[0,j]=\widehat M_h[i,0]=1

\widehat M_h[1:,1:]=M_h

With class_token=True, prepend a column and a row of ones. The mask becomes [heads, N + 1, N + 1], and the patch-to-patch block remains unchanged.

The caller must put exactly one class token first, followed by the row-major patches. The decay mask leaves its row and column unattenuated; the per-head alpha still scales all logits. This core does not implement a second distillation token.

Python implementation

python

def attention_mask(self):
    """Keep a leading class token's row/column equal to one (D.2)."""
    mask = self.patch_mask()
    if not self.class_token:
        return mask
    heads, patches, _ = mask.shape
    column = mask.new_ones(heads, patches, 1)
    row = mask.new_ones(heads, 1, patches + 1)
    return torch.cat((row, torch.cat((column, mask), dim=-1)), dim=-2)

7. Execute one attention product

Original paper

Enlarge figure

Math and tensor shapes

S_h=Q_hK_h^\top/\sqrt{d}

L_h=\alpha_h S_h\odot\widehat M_h

A_h=\operatorname{softmax}(L_h)

Y_h=A_hV_h

Inputs are matching Q, K, V[batch, heads, T, d], with T = N or N + 1. The mask broadcasts across the batch; softmax normalizes each query over the key axis. The output keeps the same shape.

Multiplication shrinks signed logits toward zero. A negative logit can gain relative probability after this change, so a smaller mask value does not guarantee a smaller attention probability.

Place this core between a ViT's existing QKV projection and output projection, preserving the base block and token order. It computes one QK product and one value product; dense attention remains quadratic. Use the authors' implementation for complete classifier integration and training.

Python implementation

python

def forward(self, q, k, v):
    """One QK product; multiply logits by alpha*M before softmax."""
    self._validate_qkv(q, k, v)
    scores = (q @ k.transpose(-2, -1)) / math.sqrt(q.shape[-1])
    mask = self.attention_mask()
    # Negative scores shrink towards zero and can gain relative probability.
    weights = torch.softmax(
        scores * self.alpha[None, :, None, None] * mask, dim=-1
    )
    return weights @ v

Run the attention core

Download violin_attention.py and requirements-violin.txt into the same directory, then create a Python environment and run the CPU example:

bash

python -m venv .venv
source .venv/bin/activate
python -m pip install -r requirements-violin.txt

python violin_attention.py --device cpu
python violin_attention.py --height 2 --width 3 --heads 2 --head-dim 4 --class-token --seed 7 --device cpu

The default is a 14 × 14 grid, three heads and head dimension 16. The second command exercises the leading class token on a small rectangle. The script needs no model weights or dataset; run python violin_attention.py --help for all options. With a CUDA-capable PyTorch installation and GPU, the same check can run on CUDA:

bash

python violin_attention.py --height 2 --width 3 --heads 2 --head-dim 4 --class-token --seed 7 --device cuda

Stdout reports the device, output shape, added parameter count, independent-reference match and maximum output error, mask symmetry and unit diagonal, then finite-gradient checks from a synthetic backward pass. Nondegenerate runs check that beta and alpha both receive nonzero gradients. With just one patch, beta is mathematically inactive; with one total token, alpha is inactive too, and the script reports finite zero gradients with the reason.

These checks establish a runnable attention core, rather than a trained-model result or a performance measurement. The following results come from the paper's classifier training and evaluation.

Results

Fine-tuning on VTAB-1K (Table 3). VTAB-1K has 19 tasks with 1,000 training examples each, in three groups: Natural, Specialized and Structured. VIOLIN's mask is added to a pretrained model, freshly initialized, and fine-tuned together with it. Accuracy is the average over all three groups.

The Structured group (counting objects, locating them, estimating depth and orientation) gains the most, which is what a spatial prior should help. The mask also combines with parameter-efficient fine-tuning on DeiT-B: LoRA goes from 71.04% to 72.55% and DoRA from 70.75% to 71.90%. Masks learned only during fine-tuning beat masks carried over from VIOLIN pretraining (Appendix F.2).

Pretraining (Tables 5, 6 and 17). On ImageNet-1K with the DeiT recipe, DeiT-T goes from 72.2% to 73.0% and DeiT-S from 79.8% to 80.7%. Self-supervised DINO-S improves by 0.3 to 0.7 points in k-NN and linear evaluation. At base size the gain nearly disappears: DeiT-B goes from 81.8% to 81.9%. Training a DeiT-T on individual pixels of CIFAR-100, where no patching supplies any locality, the mask lifts accuracy from 60.8% to 68.0%.

Dense prediction (Table 7). Semantic segmentation on ADE20K with DeiT-B and UPerNet improves from 45.24 to 45.80 mIoU. Detection on COCO with a Swin-T backbone and Mask R-CNN moves from 42.7 to 42.8 box AP and from 39.3 to 39.7 mask AP.

Other locality priors (Table 8). All methods start from the same pretrained DeiT-B and are fine-tuned on the Structured group, where the baseline scores 57.00%.

VIOLIN scores highest, with about a thousand parameters against 105K for Swin's relative position bias and 6.2M for LocalViT. Multiplying the mask into the scores beats adding it, and a mask built from plain 2D Manhattan distance gains far less than any curve-based mask.

Critical analysis

Strengths:

• Cheap and drop-in. About a thousand parameters and no change to the rest of the network. It can be bolted onto any pretrained global-attention ViT at fine-tuning time, alongside LoRA or DoRA.
• Helps where it should. The largest gains are on the Structured tasks and on pixel-level CIFAR-100, the two settings that depend most on spatial layout.
• A clear reason to average. Claim C.4 and the curve ablation both say the same thing: each curve distorts 2D distance in its own way, and the average is more robust than any one of them.

Limitations:

• Gains shrink with scale. ImageNet pretraining gains fall from +0.9 points on DeiT-S to +0.1 on DeiT-B, and VTAB gains from +4.74 on DeiT-III-S to +1.41 on DeiT-III-H. The paper says as much: the method targets small models and small datasets.
• Much of the gain comes from having any decay mask. In the DeiT-S pretraining ablation (Table 13, baseline 79.8%), a mask from row-major order alone gives +0.7 points against +0.9 for all eight curves. On the Structured group a single Peano curve gets within 0.26 points of the full set.
• Not free at runtime. The mask adds 0.64% FLOPs, but DeiT-S inference with a batch of 256 goes from 206.1 to 233.1 ms per batch at 224×224, about 13% slower, and from 1739.3 to 1789.7 ms (about 3%) at 512×512 (Table 2). Memory is essentially unchanged.
• It keeps global attention's cost. The mask is a dense N \times N multiply, so attention stays quadratic. The Hilbert-Guided Sparse Local Attention paper uses the same curve for the opposite purpose: to make local attention fast rather than global attention more local.
• Small gains on dense prediction. The improvements are +0.56 mIoU on ADE20K and +0.1 box AP on COCO.

Related Reading

• Hilbert-Guided Sparse Local Attention: the same generalized Hilbert curve, used to make local attention block-sparse and fast
• Vision Transformer (ViT): the global-attention ViT whose missing locality VIOLIN adds back
• DINO: the self-supervised ViT that VIOLIN improves in both pretraining and fine-tuning
• LoRA: the parameter-efficient fine-tuning method VIOLIN combines with on VTAB-1K
• Swin Transformer: whose relative position bias is one of the locality priors VIOLIN is compared against