NN-kNN: A Neural Network That Explains Itself with Neighbors

A visual, code-level guide to the classification and regression versions of NN-kNN

Drafted from the uploaded NN-kNN papers and code Heuzi/NN-kNN — March 2026

Most machine learning articles force you to choose between two stories.

One story is about performance: better embeddings, better optimization, better benchmarks. The other is about interpretability: examples, prototypes, nearest neighbors, and human-readable reasons.

NN-kNN is interesting because it tries to keep both.

Instead of training a neural model first and then explaining it afterward, NN-kNN builds prediction through retrieved cases. The model still learns representations, feature importance, and case importance end-to-end, but at inference time it can still say: these were the most influential training examples, and this is how they contributed.

That makes NN-kNN a useful bridge between classical k-nearest neighbor reasoning and modern neural learning.

TL;DR

• NN-kNN keeps the case-based transparency of k-NN while learning feature representations and feature weights like a neural model.
• The classification version adds sampling, feature extraction, glocal feature weighting, and differentiable case activation.
• The regression version replaces voting with attention over retrieved labels, adds locality-aware retrieval, and optionally adapts the retrieved estimate with NN-CDH.
• In the uploaded papers Heuzi/NN-kNN, NN-kNN is competitive with strong baselines across classification and regression benchmarks, while still returning influential cases as explanations.

The core idea in one sentence

NN-kNN is a trainable k-NN-like model where the similarity function, feature weighting, and case contribution are learned by backpropagation, but the final prediction is still grounded in retrieved cases.

That last clause is the important one.

A lot of “interpretable” neural systems explain their answers post hoc. NN-kNN is closer to an inherently interpretable pipeline: retrieval is not a side tool attached after the fact; retrieval is the prediction mechanism.

Why combine neural networks and k-NN at all?

Traditional k-NN has a lot going for it:

it is easy to understand,

it naturally supports example-based explanations,

and it makes local decisions instead of relying on a single global formula.

But plain k-NN also struggles when the data are high-dimensional, when the right similarity metric is unclear, or when all features should not count equally.

Neural networks solve the opposite half of the problem. They can learn useful feature spaces, learn which dimensions matter, and scale much better with modern tooling. But their predictions are usually opaque.

NN-kNN is built around a simple bet:

Let the neural part learn how to compare cases, but let the final prediction still be made through cases.

How NN-kNN works for classification

The classification paper extends the earlier NN-kNN design into a more scalable and flexible version. The high-level pipeline is shown below.

*Figure 1. Classification-side NN-kNN pipeline. The model compares the query to sampled cases, computes feature distances, applies glocal feature weighting, activates cases, and aggregates activated cases into a final prediction.**

Here is the intuition behind each stage.

### 1) Sample a manageable set of cases

Instead of comparing the query against the entire training set every time, the model can work with a sampled subset of cases. This is one of the design changes that makes the newer NN-kNN much more practical on larger datasets.

### 2) Pass the query and cases through a feature extractor

For raw tabular features, this extractor can be the identity. For images or text, it can be something learned, pretrained, or jointly trained with NN-kNN.

This matters because vanilla k-NN often fails not because neighbor methods are bad, but because the feature space is bad.

### 3) Compute per-feature distances

NN-kNN does not immediately collapse everything into one Euclidean distance. It first computes distances feature by feature.

That design leaves room for the next step: learned feature weighting.

### 4) Apply glocal feature weighting

One of the most interesting ideas in the classification paper is the global-local (“glocal”) feature weighting scheme.

Instead of forcing every case to use the same global weight vector, or giving every case a completely independent weight vector, NN-kNN mixes a small set of shared global patterns with case-specific coefficients.

In plain language:

• the model learns a few reusable feature-importance patterns,
• each case learns how much it should rely on each pattern,
• and the resulting weighted distance becomes more adaptive without exploding the parameter count.

This is a very practical compromise. It keeps locality, but avoids the brittleness and cost of a fully per-case weighting system.

### 5) Activate cases with a differentiable retrieval function

Instead of hard “top-k only” retrieval, the classification version uses a learned case activation based on distance and case bias. Cases close enough to the query become active; others contribute little or nothing.

That turns neighbor retrieval into something differentiable and trainable.

### 6) Aggregate case evidence into a class prediction

Activated cases contribute to class evidence. The final class comes from the class with the strongest aggregate support.

So even though NN-kNN is trained like a neural model, the actual reasoning trace still looks like a case-based method:

these cases activated, these cases supported these classes, and the winning class had the strongest support.

Why this is different from post-hoc explanations

This distinction is easy to miss, but it is central.

A post-hoc system says:

“The neural network already decided. Now let me retrieve some similar examples to justify the answer.”

NN-kNN says:

“The retrieved cases are part of how the answer was formed.”

That is a much stronger interpretability claim.

What the classification experiments show

The uploaded classification paper shows three useful patterns.

Small and structured datasets

On several small tabular datasets, the newer NN-kNN generally improves on the earlier NN-kNNO formulation, and it performs especially well on the synthetic Zebra datasets where the relevant feature structure is intentionally tricky.

Table 1. Selected classification results adapted from the paper. Accuracy values are rounded from the reported table

Image classification stays competitive

On image tasks, the paper reports that Conv + NN-kNN (500 cases) reaches 0.688 on CIFAR-10 versus 0.689 for ConvNet, and 0.867 on SVHN versus 0.875 for ConvNet.

Table 2. Selected image-classification results adapted from the classification paper

The retrieved examples are actually useful to look at

The paper also includes example-based explanations for CIFAR-10. Even when the model is wrong, its top activated cases often reveal why it was confused.

Figure 2. Example-based explanations on CIFAR-10.

Why regression is harder than classification

At first glance, extending NN-kNN from classification to regression sounds easy.

Replace majority vote with averaging, right?

Not quite.

The regression paper identifies a subtle but important failure mode:

A retrieval-based regressor can make an accurate prediction by combining neighbors whose labels are far apart. Numerically that can work, but interpretability suffers because the retrieved cases are no longer giving locally consistent support.

In other words, a regressor can be right for the wrong explanation.

That is why the regression extension is more than a small patch. It changes the retrieval and aggregation story in meaningful ways.

How NN-kNN works for regression

The regression architecture keeps the same overall spirit but changes three important pieces.

*Figure 3. Regression-side NN-kNN.**

### 1) Case scores become attention over labels

For regression, the model computes a score for each case and normalizes those scores over cases. The paper experiments with softmax and sparsemax.

That gives a probability-like attention distribution over retrieved cases.

The prediction is then just a weighted average of case labels.

### 2) Retrieval is pushed toward label-similar neighbors

This is the key regression-specific idea.

The model is not only asked to predict well. It is also regularized so that the high-attention neighbors are label-similar to the query target.

That improves explanation quality, not just predictive error.

### 3) An optional adaptation module can correct the retrieved estimate

Even if the neighborhood is good, a weighted average can still be systematically off.

To handle that, the regression version optionally adds NN-CDH adaptation, which adjusts the retrieved estimate using the difference between the query and the retrieved neighborhood.

The paper trains this carefully in stages because adaptation can otherwise “rescue” weak retrieval and reduce neighborhood quality.

The regression results in plain English

The benchmark table in the regression paper is large, but the broad message is simple:

NN-kNN is often competitive with or better than weighted k-NN and MLKR, and in several datasets it gets close to or beats MLPs while retaining case-based explanations.

Average RMSE of five runs (lower is better) using raw-scale targets; values are rounded for readability. Dataset abbreviations: CH=California Housing, Di=Diabetes, Ab=Abalone, BF=Body Fat, BS=Bike Sharing, WQ=Wine Quality, AF=Airfoil, SP=Student Performance, Y=Yacht, EE=Energy Efficiency, UTK=UTKFace (ResNet features). Best value per column is in bold, and second-best is underlined

Final takeaway

NN-kNN is not just a neural network with a nearest-neighbor explanation attached. It is a trainable retrieval model whose predictions remain case-grounded from start to finish.