vision-transformers-registers

name	vision-transformers-registers
title	Vision Transformers Don't Need Trained Registers
version	0.0.2
engine	skillxiv-v0.0.2-claude-opus-4.6
license	MIT
url	https://arxiv.org/abs/2506.08010
keywords	["vision-transformers","attention-mechanisms","register-neurons","outlier-tokens","inference-optimization"]
description	Apply test-time register token injection to pre-trained Vision Transformers without retraining, eliminating high-norm outlier artifacts and improving attention map quality.

Vision Transformers Don't Need Trained Registers

Core Concept

Vision Transformers suffer from high-norm outlier tokens that create noisy attention maps and degrade downstream tasks. Previous solutions required retraining models with dedicated register tokens. This paper identifies register neurons—sparse sets of neurons that generate these outliers—and shows that by shifting activations to untrained tokens at test time, pre-trained ViTs achieve comparable or better performance without any retraining. This approach is training-free, applicable to any pre-trained ViT, and provides a simple but elegant solution to a fundamental architectural problem.

Architecture Overview

• Register Neuron Detection: Identifies MLP neurons with consistently high activations at outlier patch positions
• Attention Sinks: High-norm tokens that attract excessive softmax probability, creating artifacts
• Test-Time Token Injection: Appends untrained dummy tokens that capture outlier activations
• Neuron Activation Shifting: Copies maximum register neuron activation to test tokens, zeros elsewhere
• Zero Retraining: Fully compatible with existing pre-trained checkpoints and inference pipelines

Implementation

Step 1: Identify Register Neurons

Analyze ViT activations to find neurons responsible for outliers:

import torch
import torch.nn as nn
import numpy as np
from typing import Tuple, List

class RegisterNeuronFinder:
    """Find neurons that generate high-norm outlier tokens in ViTs"""

    def __init__(self, model, device='cuda'):
        self.model = model
        self.device = device
        self.layer_info = {}

    def extract_mlp_activations(self, images_dataloader, target_layer_idx=6):
        """
        Collect MLP activations across images to identify outlier generators.

        Typically, outliers emerge after layer 6 in OpenCLIP ViTs.
        """
        all_activations = {}
        outlier_positions = []

        for batch_idx, images in enumerate(images_dataloader):
            images = images.to(self.device)

            # Forward pass with activation capture
            activations = self.capture_layer_activations(images, target_layer_idx)
            # Shape: (batch_size, num_patches, hidden_dim)

            # Identify outlier patches (high norm tokens)
            norms = torch.norm(activations, dim=-1)  # (batch_size, num_patches)
            threshold = norms.mean() + 3 * norms.std()

            # Find positions where norms exceed threshold
            batch_outliers = torch.where(norms > threshold)
            outlier_positions.extend(batch_outliers[1].cpu().tolist())  # patch indices

            all_activations[batch_idx] = activations.detach().cpu()

         all_activations, outlier_positions

     ():
        
        activations = 

         ():
             activations
            activations = output[]  

        
        target_layer = .model.transformer.resblocks[layer_idx]
        hook = target_layer.register_forward_hook(hook_fn)

        
         torch.no_grad():
            _ = .model(images)

        hook.remove()

         activations

     ():
        
        register_neuron_indices = []

         layer_idx, activations  all_activations.items():
            
            outlier_activations = activations[:, outlier_positions, :]
            

            
            neuron_strength = outlier_activations.().mean(dim=(, ))
            

            
            threshold = np.percentile(
                neuron_strength.numpy(), threshold_percentile
            )
            strong_neurons = torch.where(neuron_strength > threshold)[]

            register_neuron_indices.append(strong_neurons.tolist())

         register_neuron_indices

Step 2: Implement Test-Time Register Token Injection

Create the core algorithm that uses dummy tokens to capture outliers:

class TestTimeRegisterInjection(nn.Module):
    """Apply register token injection at test time (no training required)"""

    def __init__(self, model, register_neurons_per_layer, num_registers=1):
        super().__init__()
        self.model = model
        self.register_neurons_per_layer = register_neurons_per_layer
        self.num_registers = num_registers

        # Create test-time register tokens (trainable but only used at test time)
        self.register_tokens = nn.Parameter(
            torch.randn(1, num_registers, model.config.hidden_size) * 0.02
        )

    def forward(self, images):
        """
        Forward pass with test-time register injection.

        At inference, we append dummy tokens to capture outlier activations
        instead of letting them pollute image patch embeddings.
        """
        # Standard forward pass through ViT
        with torch.no_grad():
            embeddings = self.model.embeddings(images)
            # (batch_size, 1 + num_patches, hidden_dim)

        # Inject register tokens
        batch_size = embeddings.shape[0]
        register_tokens = self.register_tokens.expand(batch_size, -1, -1)

        embeddings_with_registers = torch.cat(
            [embeddings, register_tokens], dim=1
        )

        # Forward through blocks with register neuron shifting
        x = embeddings_with_registers

        for layer_idx, block in enumerate(self.model.transformer.resblocks):
            x = block(x)

            
             layer_idx  .register_neurons_per_layer:
                x = .shift_register_neurons(
                    x, layer_idx, register_token_start_idx= + embeddings.shape[]
                )

         x[:, , :]  

     ():
        
        num_registers = x.shape[] - register_token_start_idx

        
        register_neurons = torch.tensor(
            .register_neurons_per_layer[layer_idx]
        ).to(x.device)

         (register_neurons) == :
             x

        
         reg_idx  (num_registers):
            token_pos = register_token_start_idx + reg_idx

            
            image_activations = x[:, :register_token_start_idx,
                                 register_neurons]
            

            
            max_activation = image_activations.(dim=-)[]  
            max_activation = max_activation.(dim=-)[]  

            
            x[:, token_pos, register_neurons] = max_activation.unsqueeze(-)

            
            x[:, :register_token_start_idx, register_neurons] = 

         x

Step 3: Analyze Attention Maps Before and After

Visualize the improvement from register token injection:

class AttentionAnalyzer:
    """Analyze attention map quality improvements"""

    def __init__(self, model):
        self.model = model

    def extract_attention_maps(self, images, layer_idx=11, head_idx=0):
        """Extract attention maps from specific layer and head"""
        attention_maps = []

        def hook_fn(module, input, output):
            # output is (attention_output, attention_weights)
            attention_maps.append(output[1])

        # Register hook on attention layer
        target_layer = self.model.transformer.resblocks[layer_idx].attn
        hook = target_layer.register_forward_hook(hook_fn)

        with torch.no_grad():
            _ = self.model(images)

        hook.remove()

        return attention_maps[0]  # (batch, num_heads, seq_len, seq_len)

    def compute_attention_quality_metrics(self, attention_maps):
        """
        Compute metrics for attention map quality.

        Good attention maps have:
        - Dispersed patterns (not concentrated)
        - Clean spatial structure
        """
        batch_size, num_heads, seq_len, _ = attention_maps.shape

        metrics = {}

        # 1. Entropy (higher = more dispersed = better)
        entropy = -torch.sum(attention_maps * torch.log(attention_maps + 1e-10),
                            dim=-1).mean()
        metrics['entropy'] = entropy.item()

        # 2. Concentration (lower = better, less "attention sinks")
        max_attention = attention_maps.(dim=-)[].mean()
        metrics[] = max_attention.item()

        
        spatial_coherence = .compute_spatial_coherence(attention_maps)
        metrics[] = spatial_coherence.item()

         metrics

     ():
        
        batch_size, num_heads, seq_len, _ = attention_maps.shape

        
        distances = torch.zeros(seq_len, seq_len)
        grid_size = (np.sqrt(seq_len - ))  

         i  (, seq_len):
             j  (, seq_len):
                i_pos = i - 
                j_pos = j - 

                i_x, i_y = i_pos // grid_size, i_pos % grid_size
                j_x, j_y = j_pos // grid_size, j_pos % grid_size

                dist = np.sqrt((i_x - j_x)** + (i_y - j_y)**)
                distances[i, j] = dist

        
        coherence = torch.corrcoef(
            torch.stack([distances.flatten(), attention_maps.mean(dim=(, )).flatten()])
        )[, ]

         coherence

Step 4: Evaluation on Downstream Tasks

Compare performance with and without register injection:

def evaluate_with_registers(model_original, model_with_registers, eval_dataloader):
    """Compare ViT performance with and without test-time registers"""

    results = {'original': {}, 'with_registers': {}}

    # Evaluate original model
    print("Evaluating original model...")
    original_acc = evaluate_model(model_original, eval_dataloader)
    results['original']['accuracy'] = original_acc

    # Evaluate with registers
    print("Evaluating with test-time registers...")
    registers_acc = evaluate_model(model_with_registers, eval_dataloader)
    results['with_registers']['accuracy'] = registers_acc

    # Compute improvement
    improvement = (registers_acc - original_acc) / original_acc * 100

    print(f"\nResults:")
    print(f"Original: {original_acc:.4f}")
    print(f"With Registers: {registers_acc:.4f}")
    print(f"Improvement: {improvement:+.2f}%")

    return results

def evaluate_model(model, dataloader):
    """Standard accuracy evaluation"""
    model.eval()
    correct, total = 0, 0

    with torch.no_grad():
        for images, labels in dataloader:
            images = images.to(model.device)
            labels = labels.to(model.device)

            outputs = model(images)
            _, predicted = torch.(outputs, )

            correct += (predicted == labels).().item()
            total += labels.size()

     correct / total

Practical Guidance

• Training-Free: No retraining required; works with any pre-trained ViT checkpoint
• Register Count: Typically 1-4 registers per image; 1 often sufficient
• Layer Selection: Register neurons typically emerge after layer 6 in vision models
• Threshold Tuning: 95th percentile works well; adjust based on model/task
• Computation Cost: Negligible; only adds a few hundred tokens to sequence
• Performance Boost: Typically 1-3% improvement on vision tasks
• Attention Quality: Significantly cleaner attention maps (higher entropy, lower concentration)
• Generalization: Works across different ViT architectures (ViT-B, ViT-L, OpenCLIP variants)

Reference

• Register neurons are sparse sets of neurons with specific geometric structure
• Test-time injection exploits the observation that outliers are not task-critical information
• Unlike trained registers, untrained dummy tokens work just as well, suggesting outlier problem is structural
• Shifting activations preserves information while improving attention map interpretability