Attention Via Convolutional Nearest Neighbors

TL;DR

ConvNN unifies convolution and self-attention under a single $k$-NN neighborhood aggregation framework, recovering both as special cases and enabling interpolation along a spectrum between local spatial and global feature-based neighbor selection. The method introduces similarity-based neighbor selection, a weighted aggregation via $\rho$, and a Conv1D-based fusion, with extensions for positional encoding, sparse search, and a hybrid branching layer that blends ConvNN with standard convolution. Empirical results on CIFAR-10/100 show consistent gains from hybrid branching in VGG-11 and competitive, often superior, performance to attention baselines in ViT-Tiny, with notable regularization benefits. The work offers a principled perspective on CNNs and Transformers, suggesting new avenues for interpretable and efficient vision architectures that balance locality and global context.

Abstract

The shift from Convolutional Neural Networks to Transformers has reshaped computer vision, yet these two architectural families are typically viewed as fundamentally distinct. We argue that convolution and self-attention, despite their apparent differences, can be unified within a single k-nearest neighbor aggregation framework. The critical insight is that both operations are special cases of neighbor selection and aggregation; convolution selects neighbors by spatial proximity, while attention selects by feature similarity, revealing they exist on a continuous spectrum. We introduce Convolutional Nearest Neighbors (ConvNN), a unified framework that formalizes this connection. Crucially, ConvNN serves as a drop-in replacement for convolutional and attention layers, enabling systematic exploration of the intermediate spectrum between these two extremes. We validate the framework's coherence on CIFAR-10 and CIFAR-100 classification tasks across two complementary architectures: (1) Hybrid branching in VGG improves accuracy on both CIFAR datasets by combining spatial-proximity and feature-similarity selection; and (2) ConvNN in ViT outperforms standard attention and other attention variants on both datasets. Extensive ablations on $k$ values and architectural variants reveal that interpolating along this spectrum provides regularization benefits by balancing local and global receptive fields. Our work provides a unifying framework that dissolves the apparent distinction between convolution and attention, with implications for designing more principled and interpretable vision architectures.

Figures (10)

• Figure 1: The ConvNN operation
• Figure 2: Neighbor candidate selection strategies for $r=16$. Candidates features are highlighted in gray.
• Figure 3: Loss comparison on VGG-11 training. In branching, we set $\lambda = 0.5$. For visualization purposes, all loss curves are smoothed using a moving average with a window size of 10.
• Figure 4: Impact of $k$ to ConvNN's performance on VGG-11 training against the Conv2D baseline.
• Figure 5: Impact of $r$ to ConvNN's performance on VGG-11 training. The Random and Spatial candidate selection curves are slightly offset vertically for clarity, had the same runtime throughout the experiments.