LookHere: Vision Transformers with Directed Attention Generalize and Extrapolate

Authors: anthony  fuller
daniel g kyrollos
yousef  yassin
james r green
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 02:30:26

Abstract: High-resolution images offer more information about scenes that can improve model accuracy. However, the dominant model architecture in computer vision, the vision transformer (ViT), cannot effectively leverage larger images without finetuning — ViTs poorly extrapolate to more patches at test time, although transformers offer sequence length flexibility. We attribute this shortcoming to the current patch position encoding methods, which create a distribution shift when extrapolating.

We propose a drop-in replacement for the position encoding of plain ViTs that restricts attention heads to fixed fields of view, pointed in different directions, using 2D attention masks. Our novel method, called LookHere, provides translation-equivariance, ensures attention head diversity, and limits the distribution shift that attention heads face when extrapolating. We demonstrate that LookHere improves performance on classification (avg. 1.6%), against adversarial attack (avg. 5.4%), and decreases calibration error (avg. 1.5%) — on ImageNet without extrapolation. With extrapolation, LookHere outperforms the current SoTA position encoding method, 2D-RoPE, by 21.7% on ImageNet when trained at $224^2$ px and tested at $1024^2$ px. Additionally, we release a high-resolution test set to improve the evaluation of high-resolution image classifiers, called ImageNet-HR.

Main paper: 
 introduction: 
 There is a decades-long trend in computer vision towards higher-resolution imagery, which contains more detailed scene information. Increasing resolution is a reliable way to improve model accuracy [13, 14, 15, 16, 17, 18, 19, 20] , but this comes at a cost; training models for hundreds of epochs on large-scale datasets is expensive, especially at high-resolutions. There are two ways to reduce this cost and still see accuracy benefits from high-resolutions: 1 high-resolution finetuning, which pretrains models at a lower resolution, like 224 2 px, then finetunes them at a higher resolution, like 384 2 px; and 2 extrapolating, which deploys models at a higher resolution, without further training. Of these two options, we should aim for models that can effectively extrapolate, as it presents a zero-cost solution that does not require finetuning at every target resolution. Finetuning costs aside, improvements to extrapolation should benefit high-resolution finetuning since models that are better at extrapolating can adapt to higher resolutions more easily. Although extrapolation is a significant and exciting challenge, state-of-the-art (SoTA) model architectures extrapolate poorly.Vision transformers (ViTs [9] ) offer SoTA performance on many computer vision tasks. ViTs are simple; they split images into non-overlapping patches, linearly project pixels to form patch embeddings, and process these "tokens" with a stack of architecturally identical transformer layersmaintaining a constant feature map size throughout. This non-hierarchical design enables learning patch representations, which are useful for dense prediction tasks [21, 22, 23] and are fundamental Figure 1 : ViT-B/16 models trained for 150 epochs on ImageNet at 224 2 px and tested up to 1024 2 px. Model architectures are consistent between runs other than position encoding methods. We perform an 8-run hyperparameter sweep, per method, to ensure fair comparisons. Our three LookHere variants improve extrapolation ability, with more narrow fields of view performing best at 1024 2 . for vision-language models [24, 25, 26] . The design enables efficient processing of only a subset of patches, known as token dropping [27, 28] . Lastly, it enables model scaling by increasing the embedding size and the layer count [29, 30] .Image-size extrapolation with ViTs can be achieved in three ways:1 increasing the patch size, which packs more pixels into each patch embedding; 2 increasing the "patchification" stride, which skips-over pixels; and 3 increasing the number of patches. Of these three options, we should aim for models that can effectively ingest more patches -called "sequence length extrapolation" in the natural language processing (NLP) community [31] -as a greater number of patches presents models with more (uncompressed) information that we hope to leverage into higher accuracy. Furthermore, methods that improve sequence length extrapolation, like our proposed method, can be fused with methods that adjust patch sizes, like FlexiViT [32] . We strongly believe that patch position encoding is a primary cause of the poor sequence length extrapolation ability of ViTs -like it is in NLP, where significant advancements have been made by improving position encoding [31, 33, 34, 35] .Adding learnable or fixed sinusoidal position embeddings to patch embeddings before the first layer is the most common way ViTs encode positions. Recently, the rotary position embeddings (RoPE [36] ) used in SoTA language models [37, 38] were extended to ViTs, as 2D-RoPE [7] , showing exciting results. RoPE is a different approach to position encoding that injects positional information in each self-attention layer by rotating queries and keys with fixed sinusoidal embeddings. But for these methods to ingest more patches at test time, they must either introduce new position embeddings or modify existing embeddings -both options create a significant distribution shift. Motivated by these observations and more, we make the following contributions: when tested at the training resolution; b significantly improves performance when tested beyond the training resolution; and c increases its performance advantage after high-resolution finetuning. 3 Extrapolation Insights -We show that extrapolation: a benefits images with small objects the most, as they occupy more patches at test time; b produces class-level and dataset-level effects; and c creates distribution shifts that can be visualized via attention maps. 4 ImageNet-HR -We introduce the first natively high-resolution ImageNet test set (1024 2 px) aimed to benchmark classifiers on images that were not upsampled to achieve the target image size. 
background and related work: 
 A ViT splits an image into a grid of non-overlapping patches, flattens the grid into a sequence, and flattens the patches into vectors; i.e., R Y ×X×C → R Ny×Nx×P 2 ×C → R (Ny•Nx)×(P 2 •C), where Y is the image-height, X is the image-width, C is the number of channels, N y is the grid-height, N x is the grid-width, P is the patch height and width. A linear layer maps each vector of pixels to a patch embedding; i.e., RP 2 •C → E patch i ∈ R D ,where D is the embedding dimension also known as the transformer width. We define i and (i y , i x ) as the sequence position and the 2D position of the i th patch, respectively, where N is the total number of patches, equal to N y • N x , i ∈ {1, 2, . . . , N }, i y ∈ {1, 2, . . . , N y }, and i x ∈ {1, 2, . . . , N x }. Finally, sequence length extrapolation occurs when N test > N train .A patch embedding represents the content of a patch, and contains no information representing its original location within the image. Thus, we must encode patch positions to enable spatial reasoning; otherwise, a ViT will operate on a bag of patches.We define a "plain ViT" as attention-only and non-hierarchical. Our primary goal is to improve the extrapolation ability -i.e., generalize to more patches at test time -of plain ViTs. Our work is motivationally aligned with FlexiViT [32] and NaViT [6] , improving the flexibility of plain ViTs. Next, we briefly describe seven position encoding methods and refer the reader to the cited studies for further details; we include them all in our controlled experiments. Another method, iRPE [39] , is also compatible with plain ViTs. However, we exclude it because it is more than twice as slow as other methods; nonetheless, we benchmark iRPE with our best training recipe in Appendix A.2.1.Input Embeddings. This group leverages learned or fixed position embeddings, E pos i ∈ R D , that are added to patch embeddings at the transformer input; i.e., z i = E patch i + E pos i , where z is the input to the first transformer layer. Position embeddings represent the absolute positions of patches in an image.1 1D position embeddings [9] (1D-learn for short) map i to learnable embeddings. 2 2D sinusoidal embeddings [8] (2D-sincos for short) individually map i y and i x to fixed 1D-sinusoidal embeddings(E y i , E x i ∈ R D 2), then concatenate them along the embedding dimension. 3 Factorized position embeddings [6] (Factorized for short) individually map i y and i x to learnable embeddings (E y i , E x i ∈ R D ), then add them. 4 Learnable Fourier features [11] (Fourier for short) map (i y , i x ) to Fourier features [40, 41] , then to embeddings with a multi-layer perceptron (MLP).Attention Biases. This group leverages learned or fixed operations that encode positions by modifying the pairwise interactions between patches in self-attention without adding position embeddings to patch embeddings. Recall that self-attention first applies three separate linear transformations to project internal patch representations and splits the resultant vectors into H smaller vectors of length D H ; i.e., R N ×D → R 3×N ×H×D H -creating queries, keys, and values for each attention head. We denote a specific head by h. Next, attention scores (A ∈ R H×N ×N ) are calculated by measuring the similarity between all pairs of queries (q hi ∈ R D H ) and keys (k hj ∈ R D H ), separately, for each head; i.e., a hij = q hi • k hj / √ D H , where i and j are query and key sequence positions, and we define (i y , i x ) and (j y , j x ) as their 2D positions. Attention scores (a hij ) represent the amount of information moving from patch position j to i -whereas values (v hj ∈ R D H ) represent the content of the moving information.(ALiBi [31] ), 2D-ALiBi [12] penalizes attention scores as a function of the Euclidean distance between (i y , i x ) and (j y , j x ), and a head-specific scalar, called a slope. Slopes bias attention heads at different rates. 7 A 2D extension of rotary position embeddings (RoPE [36] ), 2D-RoPE [7] rotates queries and keys as a function of their positions. Each query is rotated by the sinusoidal embedding of i y for half its dimensions and the sinusoidal embedding of i x for the other half of its dimensions; likewise, keys are rotated as a function of j y and j x .Non-plain ViTs. Many hybrid or hierarchical architectures have been invented that often encode positions differently [42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53] . Although these architectures may be favored in some circumstances, the plain ViT is the most common single architecture due to its simplicity, flexibility, and scalability. We benchmark many non-plain ViTs and large SoTA ViTs on extrapolation in Appendix A.2.1.ViT Extrapolation. Some ViTs have been tested at higher resolutions than they were trained [54, 43, 12, 55] . NaViT [6] benchmarked input embedding methods on extrapolation, none see the gains at higher resolutions that we observe.3 LookHere Design Motivation. We introduce 2D attention masks that assign each attention head a direction and a FOV, preventing attention outside the head's FOV. Within a head's FOV, attention scores are penalized based on relative patch distances. Three ideas motivate this design. 1 Attention head diversity: heads often learn redundant algorithms that can be pruned with little accuracy penalty [56, 57, 58] . Head redundancy has also been observed in NLP [59, 60, 61] , where diversity-encouraging loss functions have been leveraged to improve generalization [62, 63, 64, 65] . From a mechanistic point of view, we can think of attention heads as an ensemble of sub-networks that "operate completely in parallel, and each add their output back into the residual stream," [66] and the residual stream is mapped to logits. Diversity has long been a desirable property of ensembles [67, 68] , and constraining attention heads to focus in different directions ensures it. 2 Attention head consistency: heads often learn interpretable spatial algorithms, like "attend to the area above the query," which reliably retrieves information from the internal representations above the query; however, we believe these types of spatial algorithms might fail when new or modified position embeddings are introduced to encode new patch positions during extrapolation -misleading the model about the information above the query, for example. We believe hard-coding both directions and distances (via attention masks and biases) will reduce the need for models to learn their own spatial algorithms. 3 Translation-equivariance has long been a desirable property of vision models, contributing to the success of convolutional networks [69, 70, 71] . ViTs are critiqued for weak inductive biases, leading to poor sample efficiency when trained from scratch [72, 73, 74] . We believe that LookHere's stronger inductive biases, achieved via directional masking and distance penalties, can improve ViT sample efficiency.Design Specifics. Let H be the number of heads, L be the number of layers, and N be the number of patches (plus one for the CLS token). We denote the LookHere matrices by A FIX ∈ R L×H×(N +1)×(N +1) . We encode positions by subtracting the LookHere matrix for a layer l, A l FIX , from the learned attention matrix, A l LRN = QK T / √ D H , before the softmax that normalizes the attention matrix prior to multiplying it by values [75] , i.e., A l = softmax(A l LRN -A l FIX ). We do not add position embeddings to patch embeddings.Let i and j be query and key sequence positions, respectively, with 2D-coordinates (i y , i x ) and (j y , j x ). Crucially, j is visible to i if j lies within i's FOV. This attention masking technique is inspired by the 1D causal masks used in autoregressive transformer decoders used in NLP [75] . When j is visible, we bias the attention score based on the Euclidean distance between i and j to encode the relative distance between patches. We scale distances via a slope function m :N L × N H → R, m(l, h) = s l (l) • s h (h)• s g that strengthens or weakens the distance penalty as a function of the head (s h : N H → R) and layer (s l : N L → R), scaled by a global slope s g ∈ R. Finally, the CLS token is visible to all positions.EQUATIONDistance(i, j) = (i y -j y ) 2 + (i x -j x ) 2(2) Design Ablations. We offer four takeaways through extensive ablations (Appendix A.6): 1 LookHere is robust to the choice of slope function. We set our default s l to linearly decrease from 1.5 to 0.5 with increasing depth (inspired by depth-wise attention distance findings [76] ). This helped in preliminary experiments, but the benefits disappear in our ablations. We arbitrarily set our default s h to ( 1 2 , 1 8 , 1 32 , 1 128 ) for the four undirected heads, but distance penalties on undirected heads can be removed entirely. We set s g = 1; LookHere is also robust to the choice of the global slope. We believe precisely tuning slopes is unnecessary because models can learn to scale attention logit magnitudes. 2 Increasing penalties with the square or square root of the distance harms extrapolation. 3 Removing all distance penalties harms extrapolation. 4 Our main contribution, 2D directional masks, are crucial to retain performance, but our method is robust to many directional configurations.(0, 0) (0, 1) (0, 2) (1, 0) (1, 1) (1, 2) (2, 0) (2, 1) (2, 2) CLS CLS (0, 0) (0, 0) (0, 1) (0, 1) (0, 2) (0, 2) (1, 0) (1, 0) (1, 1) (1, 1) (1, 2) (1, 2) (2, 0) (2, 0) (2, 1) (2, 1) (2, 2) (2, 2) K Q CLS CLS (0, 0) (0, 0) (0, 1) (0, 1) (0, 2) (0, 2) (1, 0) (1, 0) (1, 1) (1, 1) (1, 2) (1, 2) (2, 0) (2, 0) (2, 1) (2, 1) (2, 2) (2, 2) K Q ∞ ∞ A LRN ← QK T √ d k A FIX ← LookHereCompute. A FIX is precomputed and fixed, subtracting it element-wise from the learned attention matrices A LRN only costs H • (N + 1) • (N + 1) floating point operations (FLOPs) per layer. For a ViT-B/16 model, these subtractions account for 0.016% of the total FLOPs. LookHere reduces FLOPs by not adding position embeddings to patch embeddings, but this amount is also negligible. Additionally, LookHere matrices offer structured sparsity (up to 7/8 for a 45°FOV) that can speedup attention -although exciting, this speedup requires custom kernels that we leave for future work. 
experiments: 
 Deep neural networks -including ViTs -can be sensitive to seemingly minor hyperparameter changes when trained from scratch. Dosovitskiy et al. [9] finetuned the original ViT at a higher resolution, reaching 77.9% top-1 accuracy on ImageNet (we refer to ILSVRC2012 or ImageNet-1k as ImageNet). Steiner et al. [77] searched 28 hyperparameter configurations, achieving best and average runs of 80.0% and 76.9%, respectively (average calculation omits runs without data augmentation, as they were poor). Touvron et al. [78] ablated repeat augmentation [79] , dropping accuracy by 4.8%. Touvron et al. [17] replaced cross-entropy loss with binary cross-entropy loss, raising accuracy by 1.3%. Importantly, these are all ViT-B/16 models trained from scratch for 300 epochs on ImageNet. Informed by these observations and more, we design a controlled experiment: We search 8 hyperparameter configurations for each position encoding method using a single codebase; this offers an apples-to-apples comparison between our three LookHere variants and seven baselines. For each configuration, we train a ViT-B/16 on 99% of the ImageNet training set, holding the last 1% as a validation set called "minival", following [77, 81] (see Appendix A.4.1 for other hyperparameters). We train all models from scratch for 150 epochs on 224 2 px images. Our results are competitive and sometimes surpass ViTs trained for much longer, which validates our setup. The best models (according to minival accuracy), among our 8-run hyperparameter sweep per method, are always trained using 3-Augment [17] , a 3.0 • 10 -3 learning rate, and a 0.05 weight decay.Test sets. We test all 80 models on six ImageNet test sets. This includes 1 the original "validation" set used as a test set (Val for short [1] ), 2 the reassessed labels of the original validation set (ReaL for short [4] ), 3 the independently collected and in-distribution test set (v2 for short [2] ), 4 the natural adversarial test set (-A for short [3] ), 5 the ImageNet rendition test set (-R for short [5] ), and 6 the high-resolution test set that we introduce (-HR for short). ImageNet-HR. Since there are no natively high-resolution ImageNet test sets, there are two options to test the extrapolation ability of models trained on ImageNet: 1 upsample existing test sets to higher resolutions, and 2 collect a high-resolution test set ourselves. However, upsampling lowresolution images introduces another distribution shift (i.e., interpolated pixels) that we may not want to test. Thus, we collect a high-resolution test set to remove this confounding variable from our analysis. We manually collect 5 images for each ImageNet class, resulting in 5k total images, and manually crop them to 1024 2 px. This is smaller than other test sets (v2 is 30k images, -A is 7.5k images). However, we invest considerable resources to ensure its quality with two priorities: annotation accuracy and image diversity. See Appendix A.1 for details. ImageNet-HR can be accessed: https://huggingface.co/datasets/antofuller/ImageNet-HR Adversarial Attacks. We perform Fast Gradient Sign Method (FGSM [82] ) adversarial attacks with two strengths ( 1 255 , 3 255 ) on all models using Val images. Calibration Estimates. We calculate the Expected Calibration Error (ECE [83] ) with 15 bins of all models using Val images.Higher-Resolution Finetuning. With the best model per method, we continue training on ImageNet for 5 epochs at 384 2 px. We test at 384 2 px without extrapolating.Segmentation. With the best model per method, we finetune following the Segmenter protocol with a linear decoder [84] . Additionally, we probe the patches by only training a linear layer to produce a low-resolution logit map which is upsampled to obtain a full resolution segmentation map, following [85] . We run these experiments on ADE20k [86] at 512 2 px and Cityscapes [87] at 768 2 px.Patch Logit-lens. Inspired by interpretability research [88] , we evaluate the quality of the learned patch representations for models leveraging LookHere compared with other methods. Following prior work [89, 90] , we project frozen patch representations onto the learned class embedding space using the MLP classifier head that was learned for the CLS token. We leverage the ImageNet-S dataset [91] , which contains partial segmentation maps for 12k images from Val, covering 919 ImageNet classes.Extrapolating. With the best model per method, we test on images larger than 224 2 px, increasing the number of patches and we test on images smaller than 224 2 px, decreasing the number of patches; for both experiments, no further training is performed -the models are tested on their resolution generalization ability. For 1D-learn and 2D-sincos, we bilinearly interpolate the position embeddings used during training. For Factorized, we linearly interpolate the position embeddings for each axis. Fourier does not require adjustment since fractional positions along each axis are used as input. For RPE-learn, we interpolate the learned relative biases using the official BEiT implementation [10] . 2D-ALiBi does not require adjustment either. However, we tune a parameter on minival that scales the distance penalty at each test resolution. For 2D-RoPE, we tune its base frequency on minival -this is a SoTA method to extrapolate RoPE used in NLP [33] . Lastly for LookHere, we tune the global slope on minival. The benefits of tuning slopes are minimal, see Appendix A.4.4.Table 1 : Top-1 acc. (%) for ViT-B models trained on ImageNet for 150 epochs; trained and tested at 224 2 . We report the best and average results across our 8-run hyper-parameter sweep.Val [1] ReaL LookHere improves ViT adversarial robustness and model calibration (Tables 2 3) ; both have been linked to ensemble diversity [92, 93, 94] , which we offer as a potential explanation. This is an interesting finding because adversarial robustness and calibration can be at odds with accuracy [95, 96] . We show that LookHere learns more diverse attention heads by measuring the generalized Jensen-Shannon divergence [97] between heads (Figure 4 ). In the Appendix A.8, we measure more properties of models leveraging different position encoding methods. LookHere significantly outperforms other methods on segmentation linear probing, demonstrating its ability to learn spatially- High-resolution finetuning increases the performance advantage of all three LookHere variants over 2D-RoPE (Table 5 ). This aligns with our intuition that improving extrapolation methods can improve high-resolution finetuning. Lower initial finetuning loss has been linked to better retaining the general representations learned during pretraining [98] , and better extrapolating models have lower initial loss at a higherresolution, by definition.[4] v2 [2] -A [3] -R [Using a "logit lens" [88] approach, we project patch representations onto the class embedding space [89] . We observe that LookHere encodes semantic information in its patches faithful to the original patch location; these patch-level predictions act as a segmentation map that can be generated without additional training. The officer in Figure 5 is not a one-off example; using ImageNet-S [91] , we see that LookHere outperforms 2D-RoPE by at least 22% mIoU using this patch-projection method (Figure 5 ). Our best explanation is that, by restricting attention, LookHere 1D-learn (4.1%) 2D-sincos (2.1%) Factorized (8.7%) Fourier (3.3%) RPE-learn (6.1%)2D-ALiBi (6.5%) 2D-RoPE (8.7%) LH-180 (31.2%) LH-90 (31.1%) LH-45 (30.7%)Figure 5 : We apply frozen MLP classifying heads (learned on the CLS token) on frozen patch representations. We visualize ImageNet class predictions: assault rifle (red), bulletproof vest (green), crash helmet (blue), and holster ( white ). In parentheses, we show mIoU results (@224px) on ImageNet-S [91] , where we apply this technique to segment images without training. prevents the attention collapse at deeper layers observed in Figure 4 that divorces patch representations from their original patch locations; this collapse has been observed in other ViTs [99, 100] . We also expect that preventing attention collapse will benefit vision-language models, where frozen patch representations are used as "image tokens" that should represent their original patch locations [24, 25, 26] . More examples and detailed analysis are in Appendix A.7LookHere significantly improves extrapolation ability (Figure 1 ). Our smallest FOV variant (LH-45) sees improving relative performance as resolution increases. LH-45 outperforms 2D-ALiBi, which is equivalent to LookHere without our 2D directional masks, by 9.5% on Val at 1024 2 px. These two results demonstrate the extrapolation benefits of restricting attention to fixed FOVs. LH-45 gains 1.3% on Val when extrapolating from 224 2 to 384 2 px; this is the largest gain we find in the literature, including our extensive benchmarking of SoTA models in Appendix A.2. LookHere also outperforms other methods when tested on smaller images, but the advantage narrows (Figure 6 ).Interestingly, smaller objects benefit most from extrapolation (Figure 7 ), which are distributed over more patches at test time. We believe this effect also explains the 6 -8% that LookHere models gain when extrapolating on ImageNet-A from 224 2 to 448 2 px; by inspection, ImageNet-A seems to have small objects, and other work found zooming-in on center-cropped ImageNet-A images improves [8] 1D-learn [9] RPE-learn [10] Fourier [11] 2D-ALiBi [12] Figure 7 : The effect of object size on accuracy gains or losses due to extrapolation. Object size is measured using annotations from Kaggle's ImageNet Object Localization Challenge [101] .1 st 2 nd 3 rd 4 th 5 th -2 -1 0 1 2 3 Top-1 Accuracy [%] ∆ Acc. 224 2 →performance [102] . Finally, all LookHere variants outperform other methods on ImageNet-HR, indicating better handling of interpolated pixels generated when upsampling lower-resolution imagery is not the reason why LookHere extrapolates better. Reducing the distribution shift faced by attention heads during extrapolation is our best explanation for LookHere's large relative improvement. Figure 8 shows attention maps that are "unflattened" to visualize the image regions to which heads attend, averaged over the same 5k images. We show one head per model that exhibits similar behavior at a 224 2 resolution. Models leveraging RPE-learn and 2D-ALiBi learn variants of an algorithm that retrieve information from above the query; however, both models retrieve information elsewhere in the image when extrapolating. LookHere hard-codes this type of algorithm, which it continues to execute when extrapolating. In Appendix A.8 we find more examples of interesting attention head behaviour.Extrapolation affects different datasets differently; it also affects different classes differently. For example, when extrapolating, all models underpredict certain classes (bakery, church, and tights) and overpredict other classes (mobile home, threshing machine, and sports car). This investigation is inspired by the class-level effects of data augmentation [103] . In Appendix A.9 we find more class-level effects of extrapolation. 
closing: 
 Limitations. The primary limitation of LookHere is it requires hand-designed directional masks and distance penalties. However, our extensive ablations demonstrate that LookHere is robust to the choice of directional masks and distance penalties. The primary limitation of our experiments is we do not scale ViTs to giant sizes. Instead, we select the most common size, the ViT-B/16, and focus our computational resources on a controlled experiment -that extensively and fairly tunes the appropriate baselines for plain ViTs; this allows us to make confident conclusions based on our thorough experiments.LookHere position encoding significantly improves the ability of plain ViTs to make inferences when provided a greater number of patches than seen during training. We thoroughly demonstrate that LookHere outperforms other methods with and without extrapolation on standard image benchmarks and our high-resolution ImageNet test set called ImageNet-HR. We provide new insights into ViT extrapolation by showing object-size, class-level, and dataset-level effects. We believe LookHere will help the vision community transform higher-resolution into higher accuracy.Future Work. We are excited to realize the computational gains that LookHere makes available via sparse attention kernels, as well as bring LookHere to video and 3D point-cloud applications.Figure 9 : LH-180 bias matrices for query patch (11, 8) , grid size of 14x14.Figure 10 : LH-90 bias matrices for query patch (11, 8) , grid size of 14x14.Figure 11 : LH-45 bias matrices for query patch (11, 8) , grid size of 14x14.Recipe. Our training recipe that is consistent across configurations:• AdamW [109] -using the default PyTorch implementation that does not fully decouple learning rate and weight decay • Binary cross-entropy loss -summing along the class dimension, averaging along the batch dimension• Linear warm-up for 10% of steps and cool-down using a cosine decay schedule to a zero learning rate• Batch size of 2048• Mixup [110] α = 0.8, cutmix [111] α = 1• CLS token with an MLP classifying head -final linear layer weights are initialized to 0 and biases to -6.9 (so all class probabilities start at 1 1000 ) • layer drop rate of 0.1 and MLP dropout of 0• Train for 150 epochs on the first 99% of ImageNet-1k -using Huggingface's datasets library, i.e., load_dataset("imagenet-1k", split="train[:99%]")• Choose checkpoint according to the best minival top-1 accuracy (run after each epoch), where minival is the last 1% of the ImageNet-1k training set, i.e., load_dataset("imagenet-1k", split="train[99%:]")Training takes around 3 days on an RTX 4090 GPU. Thus, all 80 training runs take around 240 GPU-days. We spend another 54 GPU-days on 18 ablations. Ablations and our iRPE run always use our best training recipe, which is 3-Augment [17] data augmentation, 3 • 10 -3 learning rate, and 0.05 weight decay. iRPE [39] takes around 7 days on an RTX 4090 GPU, even with the official custom CUDA kernel. As a result, we exclude it from our apples-to-apples comparisons.Following DEiT III's finetuning recipe [17] , we increase the drop rate to 0.2 and the weight decay to 0.1, and fix the learning rate to 10 -5 with a 512 batch size.For 2D-ALiBi, 2D-RoPE, and LookHere models, we tune a single parameter at the target resolution on minival (Table 9 ). LookHere models benefit less from tuning than 2D-ALiBi and 2D-RoPE models. For example at a 512 2 resolution, the difference in top-1 accuracy on minival when using the tuned parameter versus the default value is 2.1% for 2D-ALiBi, 1.3% for 2D-RoPE, and 0.15% for LH-45. Thus, LookHere does not require tuning its global slope value to effectively extrapolate. For both linear probing and full finetuning we use a linear decoder. The linear decoder consists of a linear layer applied to the frozen patch representations which is then upsampled to the original image size. Similar to [85] we add a BatchNorm layer before the linear layer.For full finetuning, we followed the Segmenter training recipe [84] exactly. For ADE20k, the base learning rate is 10 -3 for 160k iterations with a batch size of 8, at 512 2 px. For Cityscapes, the base learning rate is 10 -2 for 80k iterations with a batch size of 8, at 384 2 px. We train with SGD.For linear probing, we freeze the backbone and pre-compute the patch representations.We use the AdamW optimizer [109] and sweep the following learning rates: {0.0001, 0.0002, 0.0005, 0.001, 0.002, 0.005, 0.01, 0.02, 0.05, 0.1, 0.2, 0.3, 0.5}. For both ADE20k and Cityscapes we set the batch size to 16 and train the linear decoder for 50 epochs.Table 10 : First half of our hyper-parameter sweep. ViT-B models trained on ImageNet for 150 epochs; trained and tested at 224 2 . RA is for RandAugment and 3A for 3-Augment.Val [1] ReaL [4] v2 In our paper, we show that LookHere prevents attention collapse measured by JSD (Figure 4 ). Here, we measure attention diversity using L 1 and L 2 distance. We also measure attention distances and patch-wise representational similarity -both at 224 2 px (Figure 15 ) and at all resolutions tested (Figure 16 & 17). • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [Yes] Justification: Our dataset consists of 5k images, we manually downloaded them from the internet and reviewed them several times. We do not believe any images pose a safety risk. Our models do not pose any more of a safety risk than other ViT-B models.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: We properly credit creators of datasets and models. We share the websites used to collect images and their licenses.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.LookHere -We introduce a novel position encoding method for plain ViTs that restricts attention heads to fixed fields of view (FOV) and points them in different directions via2D masks. This design provides: a translation-equivariance, b attention head diversity, c improved interpretability, and d limits the distribution shift that attention heads face when extrapolating.2 Controlled Experiments -We perform an apples-to-apples comparison between seven position encoding methods for plain ViTs alongside our three LookHere variants. We demonstrate that LookHere: a improves classification, segmentation, adversarial robustness, and model calibrationLearnable relative position encoding[10] (RPE-learn for short) biases attention scores by mapping all possible relative positions between queries and keys to learnable embeddings (B ij ∈ R H ); i.e., biases are a function of i y -j y , i x -j x , and h.A 2D extension of Attention with Linear BiasesFigure 33: Cross-plots continued.Figure 34: Cross-plots continued.Figure 35: Cross-plots continued. 
 
 Appendix:  


Decision: In the discussion period, the authors addressed the concerns of the reviewers and the reviewers are happy to accept this paper. The AC also enjoyed reading this paper and the proposed idea of directionally limiting the fields of view to better process high-resolution images.

Reviews: review: paper_summary: paper_summary: Vision transformer is known for its constrained scalability across various image resolutions and this work is designed to address the generalization ability of ViT at high-resolution images. Specifically, the authors propose LookHere, a drop-in replacement for the positional encoding of standard ViTs, which restricts attention heads to fixed fields of view oriented in different directions using 2D attention masks. It offers translation-equivariance, guarantees attention head diversity, and minimizes the distribution shift encountered by attention heads during extrapolation. Extensive experiments validate the effectiveness of the proposed method.
strength_weakness: strength_weakness: 1. The generalization ability of LookHere to smaller-scale images remains unverified. Although the authors particularly focus on testing at high-resolution images, it would be better if LookHere could be generalized to smaller images as well.

2. The ablation study on the design choices seems missing. Although the authors provide abundant experiments and analysis, the ablation study on the specific designs is missing and their efficacy is unverified.

3. The size of the proposed ImageNet-HR is too small (5 images per class). The limited number of images in each class may not be representative enough and could have a strong bias implicitly.
questions: questions: Apart from the questions in weakness, the reviewer has one additional question:

Is there a particular reason that the authors did not compare LookHere with NaViT, which can also be tested at higher resolutions?
limitations: limitations: N/A
paper_summary: paper_summary: The authors explore an alternative to existing positional encoding methods for vision transformers. Within attention operations, 2D attention masks (subtracted from attention matrix i.e. key-query inner product, prior to softmax) limits feature mixing to fixed fields of view pointed in different directions. ViTs resulting from proposed method exhibit improved performance across diverse tasks, especially when extrapolating to input image dimensions larger than the train data. The authors also contribute a new high-resolution ImageNet dataset variant for evaluation.
strength_weakness: strength_weakness: 1. **Translation Equivariance:** Authors claim this to be a result of proposed method. However, this is not proved theoretically or demonstrated through experiments. Please explain this in method section, compare to CNNs, and provide experiments to quantify this is possible. In case I missed something related from appendix, please highlight it better in main paper. 

2. **[Minor] Adversarial Attacks:** Only FGSM - maybe try newer methods like AutoAttack?
questions: questions: 1. **Learning Parameters**: Have you tried learning the slope values? And maybe even the entire attention masks? These could be learned using some SSL setup like MAE.
limitations: limitations: Discussed.
paper_summary: paper_summary: This paper introduces Lookhere, a novel mask-based positional encoding designed to address the performance degradation of ViT when resolution changes. Lookhere restricts the receptive field of each head in the attention mechanism and enhances the diversity of information perceived by each head. Extensive experiments have demonstrated the effectiveness of Lookhere. Compared to other positional encodings, Lookhere has significantly better resolution scalability.
strength_weakness: strength_weakness: In summary, I believe this is a very solid work with no obvious issues. I have only one personal concern. From Figure 6, it can be observed that the essence of the proposed positional encoding seems to be limiting the receptive field of each token, so that when the resolution changes, the number of tokens each token attends to does not change too drastically. However, this approach seems to reduce the model's overall receptive field. As a result, the fine-tuning results at a resolution of 384 do not show a particularly significant advantage compared to other positional encodings. I suspect that as the fine-tuning time increases, this advantage may disappear or even be worse than RoPE. Could the authors fine-tune for a longer period (for example, 30 epochs) to verify that this approach does not negatively impact the model's accuracy?
questions: questions: Please refer to the weakness
limitations: limitations: N/A
paper_summary: paper_summary: This paper proproses LookHere, a novel positional encoding method for Vision Transformers for dealing with high-resolution images. Specifically, LookHere explicitly constratin attention heads to attend certain directions via 2D masks. With comprehensive experiments. LookHere demonstrates strong performance across different benchmarks. It also achieves better performance than existing positional encoding methods on high-res images.
strength_weakness: strength_weakness: 1. Actually the authors have clearly dicussed the weakness of their method, e.g. hand-designed masks, single scale of ViT. Beside that, have the authors considered a discussion with ConViT [50]? ConViT introduces convolutional bias to attention head, which also restricts the head to only attend to certain areas.
2. (Minor) It would be more promising to see if this technique can be applied into multi-modal LLMs.
questions: questions: See the weakness.
limitations: limitations: Yes

Comments: title: Post-rebuttal comment
Official Comment from Reviewer
final review
Final Review
comment: The reviewer appreciates the detailed response provided by the authors. The concerns are addressed and the rating is upgraded to weak accept.
Thanks for the authors' rebuttal. They have addressed my concerns.
Thanks the author for the rebuttal. My concerns are well addressed. I have raised my score to 7.
The authors address all concerns; no changes to my original rating of 8 (strong accept). 

Highly encourage authors to release results on stronger adversarial attacks at a later point in time as feasible.