💻 computer science

Unifying Convolution and Attention via Convolutional Nearest Neighbors

This paper introduces Convolutional Nearest Neighbors (ConvNN), a unified framework that demonstrates convolution and self-attention as special cases of a single $k$ -nearest neighbor aggregation process, thereby bridging these two dominant computer vision architectures and achieving state-of-the-art performance improvements on ImageNet-1K.

Original authors: Mingi Kang, Jeová Farias Sales Rocha Neto

Published 2026-07-02

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Mingi Kang, Jeová Farias Sales Rocha Neto

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to understand a complex painting. To do this, you have two main ways of looking at it:

The "Local Neighbor" Approach (Convolution): You stand very close to a specific spot on the canvas and only look at the tiny square of paint immediately around your finger. You ignore everything else. This is how traditional Convolutional Neural Networks (CNNs) work. They are great at seeing textures and edges because they focus on immediate spatial neighbors.
The "Global Similarity" Approach (Attention): You step back and look at the whole painting. You ask, "Which other parts of this painting look like the spot I'm looking at, even if they are on the opposite side?" You then blend those similar parts together. This is how Vision Transformers (ViT) work. They are great at understanding the big picture because they can connect distant, similar features.

For a long time, computer scientists thought these two methods were completely different tools that couldn't be mixed easily.

The Big Idea: The "ConvNN" Tool

This paper introduces a new, universal tool called Convolutional Nearest Neighbors (ConvNN). The authors argue that the two methods above aren't actually different; they are just two ends of the same spectrum.

Think of ConvNN as a smart search engine for pixels.

How it works: When the computer looks at a specific pixel (the "query"), it searches for its "neighbors" to gather information.
The Twist: The computer can decide how to find these neighbors based on a rule you set:
- Rule A (Spatial Proximity): "Find the neighbors that are physically closest to me." (This turns the tool into a standard Convolution).
- Rule B (Feature Similarity): "Find the neighbors that look most like me, no matter how far away they are." (This turns the tool into Self-Attention).

The paper proves mathematically that if you tweak the settings of this one tool, you can make it act exactly like a standard convolution or a standard attention mechanism. It unifies them into a single framework.

How They Tested It

The researchers didn't just do math; they built working models to see if this idea actually helps computers see better. They tested it on ImageNet, a massive database of 1.28 million images used to train AI.

1. Testing on "Local" Models (ResNet-50):
They took a standard image-recognition model (ResNet-50) and added a "hybrid branch." Imagine a worker who usually only looks at the immediate neighborhood (Convolution) but now has a sidekick who can also look for similar-looking objects anywhere in the city (ConvNN).

Result: This hybrid team performed significantly better than the worker alone. They improved accuracy by 3.0%.
Lesson: You don't have to choose between looking locally or globally; doing both together is the winning strategy.

2. Testing on "Global" Models (Vision Transformer):
They took a Transformer model (ViT-Base) and replaced its standard "global search" with their new ConvNN tool.

Result: The new tool beat the original Transformer by 0.7%.
Key Discovery: The standard Transformer treats all its "top matches" equally. The new ConvNN tool learned to give different weights to different matches. It's like a librarian who doesn't just grab the top 10 similar books, but learns to prioritize the most relevant ones among them. This made the model especially good at handling large groups of similar items.

Why This Matters (In Simple Terms)

It's a Unifying Theory: It shows that Convolution and Attention are just different settings on the same machine.
It's Flexible: You can design a system that sits anywhere in the middle, using a mix of spatial closeness and feature similarity.
It's Efficient: The authors created a special, fast version of this tool (using something called a "Triton kernel") that uses less computer memory and runs faster, making it practical for real-world use.

A Note on Training

The paper also noticed something interesting about how the model learns. The new tool starts off slower than the standard methods during the early stages of training. However, as the training gets closer to the finish line (when the computer is fine-tuning its answers), the new tool catches up and eventually surpasses the others. It seems the new tool needs a "low and slow" learning pace to figure out exactly how to weight its neighbors, whereas standard methods are more rigid and learn quickly but stop improving sooner.

In summary: The paper introduces a single, flexible framework that bridges the gap between local and global image processing. By treating both as variations of "finding nearest neighbors," they created a tool that is mathematically sound, faster to run, and more accurate than the current leaders in both categories.

Technical Summary: Unifying Convolution and Attention via Convolutional Nearest Neighbors

Problem Statement

Convolutional Neural Networks (CNNs) and Vision Transformers (ViTs) represent the two dominant architectural families in computer vision, defined respectively by spatially local convolution and global self-attention. Despite their distinct operational definitions, prior work has hinted at a theoretical connection, suggesting that attention with relative positional encoding can express convolution, and that local attention is structurally equivalent to depthwise convolution. However, existing approaches typically treat these operations as separate entities to be bridged externally (e.g., via mixed architectures) or analyze specific pairwise links. There is a lack of a unified mathematical framework that demonstrates both operations as special cases of a single underlying principle, thereby exposing the continuous design spectrum between them.

Methodology: Convolutional Nearest Neighbors (ConvNN)

The authors introduce Convolutional Nearest Neighbors (ConvNN), a unified framework based on $k$ -nearest neighbor ( $k$ -NN) aggregation. The core insight is that both convolution and attention are instances of neighbor aggregation that differ only in their neighbor selection strategy:

Convolution selects neighbors based on spatial proximity.
Attention selects neighbors based on feature similarity.

Framework Mechanics

The ConvNN operation consists of three sequential steps:

Similarity Computation: Input features $X$ are projected into Query ( $Q$ ) and Key ( $K$ ) spaces. A similarity matrix $S$ is computed (defaulting to dot-product, but allowing for spatial distance metrics).
Neighbor Selection and Modulation: For each query, the top- $K$ most similar keys are selected. A weighting function $\rho$ (e.g., softmax or uniform) modulates the corresponding Value ( $V$ ) vectors.
Weighted Aggregation: The selected and weighted neighbors are aggregated using a 1D convolution (Conv1D) with kernel size $K$ and stride $K$ .

Theoretical Unification

The paper provides formal proofs demonstrating that ConvNN recovers existing operations as exact special cases:

Convolution: By restricting projections to raw spatial coordinates, using Euclidean distance for similarity, and selecting $K = R^d$ neighbors (where $R$ is the kernel size), ConvNN exactly recovers standard and depthwise convolution, including correct boundary behavior under zero-padding.
Attention: By using learned projections and scaled dot-product similarity, ConvNN recovers:
- Standard self-attention (when $K=N$ ).
- Sparse attention variants like KVT (when $K < N$ ).
- Structured sparse attention (e.g., BigBird, causal masking) via similarity masking.
- Multi-head attention topologies via algebraic equivalence.

Practical Implementations

To address computational and memory constraints, the authors propose three practical realizations:

Sparse Candidate Search: Instead of exhaustive $O(N^2)$ search, a random subset of $r \ll N$ candidates is selected. This avoids reintroducing spatial bias inherent in structured subsampling while reducing complexity to $O(Nr)$.
Triton-Accelerated Kernel: A fused GPU kernel merges similarity computation, neighbor gathering, and aggregation into a single launch. This eliminates intermediate tensors, reducing forward peak memory by $2.3\times$ and accelerating both forward and backward passes.
Hybrid Branching Layer: A module that processes inputs through parallel branches—one using standard convolution (spatial proximity) and one using ConvNN (feature similarity)—which are then concatenated and mixed via a pointwise convolution.

Key Results

The framework was validated on ImageNet-1K classification across two complementary architectures:

1. Convolutional Architectures (ResNet-50)

Hybrid Branching: Replacing standard $3\times3$ convolutions with a hybrid branching layer ( $\lambda=0.5$ ) improved top-1 accuracy from 76.7% (ResNet-50 baseline) to 79.76% (at $k=16$ ), a gain of 3.06%.
Pure Replacement: A standalone ConvNN replacement without the branching layer achieved ~76.2%, matching the baseline but not exceeding it, indicating that the combination of local and global selection is the driver of performance.
Candidate Selection: Random sparse candidate selection ( $r=32$ ) was computationally efficient and consistently outperformed the baseline, though it underperformed full-feature selection.

2. Vision Transformers (ViT-Base)

ConvNN-Attention: Replacing multi-head self-attention (MHSA) with ConvNN-attention improved top-1 accuracy from 80.94% (MHSA baseline) to 81.64% (at $k=25$ ), a gain of 0.70%.
Comparison with KVT: ConvNN consistently outperformed KVT (which uses fixed uniform aggregation weights) across all tested neighbor counts. The gap widened at larger $k$ (e.g., 81.64% vs. 81.28% at $k=25$ ), demonstrating the value of learned aggregation kernels over fixed uniform sums.
Training Dynamics: ConvNN exhibited a distinct convergence pattern, lagging behind MHSA in early/middle training stages but surpassing it in the final epochs. The authors attribute this to the learned aggregation kernel requiring low-learning-rate fine-tuning to converge to precise per-channel weightings.

Significance and Claims

The paper claims that ConvNN provides a principled foundation for designing operations that bridge convolutional and attention-based computation. Its significance lies in:

Unification: It formally proves that convolution and attention are not fundamentally distinct but are points on a continuous operational spectrum defined by neighbor selection criteria.
Design Space Exploration: It exposes a new design space where practitioners can systematically tune the balance between spatial proximity and feature similarity, rather than choosing one operation over the other.
Performance Gains: The empirical results demonstrate that hybrid configurations (combining both selection modes) and learned aggregation kernels can yield state-of-the-art performance improvements over matched baselines in both CNN and Transformer architectures.

The authors maintain a modest stance, noting that while ConvNN improves upon matched baselines, it does not necessarily target the absolute state-of-the-art achieved by bespoke macro-architectural designs (e.g., CoAtNet, ConvNeXt). Instead, it isolates and optimizes the underlying aggregation operation itself. Two open questions remain: whether an intermediate selection criterion can outperform both endpoints, and how the schedule-dependence of the learned aggregation kernel should be handled in practice.