Retrieval-Centric Deep Learning in Growing Nonparametric Neural Networks

Maximilian Schlegel; Rajai Nasser; Seijin Kobayashi; Yanick Schimpf; Oliver Sieberling; Robert Obryk; Kazuki Irie; João Sacramento

cs.LGOct 2, 2026

Retrieval-Centric Deep Learning in Growing Nonparametric Neural Networks

Authors: Maximilian Schlegel, Rajai Nasser, Seijin Kobayashi, Yanick Schimpf, Oliver Sieberling, Robert Obryk, Kazuki Irie, João Sacramento,

Organizations: Google, Paradigms of Intelligence Team, Zurich, Switzerland · MIT CSAIL, Cambridge, MA, USA · Department of Computer Science and Wu Tsai Institute, Yale University, New Haven, CT, USA

Abstract

We investigate a general-purpose layer for deep learning that, instead of compressing arbitrary-size training data into fixed-size weight matrices, stores a new pair of key-value representations for every data point during training, and retrieves and recombines these representations through an attention mechanism at inference time - resulting in a growing neural net (NN). While Irie et al. (arXiv:2202.05798) have put forward this perspective from the classic duality expressing any linear layer in a deep NN trained by gradient descent as linear attention (LA) over the training data points, replacing LA by more powerful attention functions, as they suggest, turns out to be non-trivial: we show that naively applying learning rules from the LA case to advanced kernels does not lead to principled optimization. Here we fill this gap and develop functional gradient-based learning rules for kernelized attention layers, based on radial basis function (RBF) and softmax-like kernels - establishing the principled "retrieval-centric deep learning" (RCDL) paradigm. Empirically, we demonstrate the promising performance and learning-efficiency of RCDL on image classification and synthetic teacher-student learning tasks. Moreover, we show that replacing LA in the dual form of NNs by advanced LA variants, namely MesaNet/DeltaNet, yields a formal connection to recently proposed optimizers for conventional fixed-size NNs, offering a novel perspective on deep learning optimization.

Figures & tables

Model	Parametric neural networks	Growing nonparametric networks
Building block	Linear layer	Attention layer
Learnable state	Fixed-size weight matrix	Growing database of keys & values
(for $t$ train examples)	${\bm{W}}\in\mathbb{R}^{d_{\text{out}}\times d_{\text{in}}}$	$\mathcal{D}=\{{\bm{K}}\in\mathbb{R}^{d_{\text{in}}\times t},{\bm{V}}\in\mathbb{R}^{d_{\text{out}}\times t}\}$
Inference	${\bm{y}}={\bm{W}}{\bm{x}}$	${\bm{y}}=s_{\mathcal{D}}({\bm{x}})=\text{Attention}({\bm{K}},{\bm{V}},{\bm{x}})$
( ${\bm{x}}\in\mathbb{R}^{d_{\text{in}}}$ , ${\bm{y}}\in\mathbb{R}^{d_{\text{out}}}$ )
Illustration

Table 1: Comparison of parametric vs. retrieval-centric deep learning paradigms.

Figure 1: Teacher-student regression loss (mean

\pm

s.e.m.) comparing: (blue, Weighted Softmax) our growing weighted softmax model; (red, Parametric Softmax) a very wide MLP baseline with width matching the final number of training samples using the softmax activation; and (gray, Reference MLP) an overparameterized student with matched architecture providing a lower bound.

Panel of Figure 2 — Figure 2: (Left Plots) Test set accuracy on CIFAR-10/100: Growing Weighted Softmax & RBF, while based on plain SGD, significantly outperform the SGD-trained MLP and Adam/RMSProp and are competitive with MLPs trained with Muon. The Mesa/FOOF optimizer achieves high accuracy on CIFAR-10 and performs best on CIFAR-100. (Middle Right) Training loss for Growing Weighted Softmax & RBF decreases much faster than that of the other models. (Right) Test set accuracy on CIFAR-10 after a single pass on the training data. Growing Weighted Softmax & RBF are the most efficient learning algorithms. To place our results, we include the test accuracy of two classical (not single pass) methods: k-nearest-neighbor (kNN) and SVM (starred; details in Appendix F.10 ).

Figure 5: CIFAR-10 performance vs final capacity for various growth thresholds

\tau

. The model retains high test acc. while exhibiting graceful performance degradation when compressing further.

Appendix figures & tables19 assets

Supplementary material from the paper’s appendix.

Appendix

Hyperparameter	Value
Model	Single Weighted Softmax layer
Data normalization	RMSNorm (on input)
Value learning rate $\beta$	$[0.1,0.3,1.0]$
Weight learning rate $\eta$	$[10.0,30.0,100.0]$
Softmax temperature $T$	$[0.1,1.0,3.0,10.0]$
Momentum	$[0.6,0.9]$

Appendix

Table 2: Hyperparameters Nonparametric Growing Weighted Softmax for Teacher-Student experiments.

Hyperparameter	Value
Dataset	CIFAR-10, CIFAR-100
Input	$32\times 32\times 3$ images flattened to $3072$
Model	Residual MLP
Residual stream size ( $d_{\text{e}}$ )	$3072$
# residual blocks ( $L$ )	$[0,1,2]$
Normalization in blocks	RMSNorm (pre-norm, per block)

Appendix

Table 3: Hyperparameters Growing Weighted Softmax for experiments on CIFAR-10 and CIFAR-100.

Hyperparameter	Value
Model	$L=3$ residual blocks + readout layer, no nonlinearity
Embedding dimension $d_{e}$	$3072$ (flattened image, no input embedding)
Normalization	RMSNorm (no learned scale) before each layer and the readout
Maximal width per layer $N$	$50016$
Growth threshold $\tau$	$[35,40,45,50,55,60]$
Value learning rate $\beta$	$[10,30,60]$

Appendix

Table 4: Hyperparameters for the Hybrid Growing Weighted Softmax on CIFAR-10 (Figure 5 ).

Hyperparameter	Value
Dataset	CIFAR-10
Input	$32\times 32\times 3$ images flattened to $3072$
Model	MLP
Normalization in blocks	RMSNorm (pre-norm, per block)
$W_{\text{Key}}$ shape	(2560000, 3072)
$W_{\text{Value}}$ shape	(2560000, 10)

Appendix

Table 5: Hyperparameters Parametric Softmax for experiments on CIFAR-10.

Hyperparameter	Value
Model	Single Parametric Softmax layer
Hidden size	$10240000$
Data normalization	RMSNorm (on input)
Learning rate $\eta$	$[0.0003,0.001,0.003,0.01,0.03]$
Temperature $T$	$[0.001,0.01,0.1,0.3,1.0,3.0,10.0,30.0,100.0]$
Optimizer	Adam

Appendix

Table 6: Hyperparameters Parametric Softmax for Teacher-Student experiments.

Hyperparameter	Value
Dataset	CIFAR-10, CIFAR-100
Input	$32\times 32\times 3$ images flattened to $3072$
Model	Residual MLP
Residual stream size ( $d_{e}$ )	$3072$
# residual blocks ( $L$ )	$[0,1,2]$
Normalization in blocks	RMSNorm (pre-norm, per block)

Appendix

Table 7: Hyperparameters RBF for experiments on CIFAR-10 and CIFAR-100.

	Mesa newton-approximation steps
Dataset	5	10	Inf.
CIFAR-10	0.6914	0.7130	0.7124
CIFAR-100	0.4149	0.4323	0.4319

Appendix

Table 8: Best Average Accuracy on CIFAR-10 and CIFAR-100 across different numbers of steps for the newton iteration algorithm for approximating the inverse. We find that generalization performance on CIFAR saturates with

10

iterations.

Hyperparameter	Value
Dataset	CIFAR-10, CIFAR-100
Input	$32\times 32\times 3$ images flattened to $3072$
Model	Residual MLP
Residual stream size ( $d_{e}$ )	$3072$
# residual blocks ( $L$ )	$[2,4]$
Normalization in blocks	RMSNorm (pre-norm, per block)

Appendix

Table 9: Hyperparameters Mesa for experiments on CIFAR-10 and CIFAR-100.

Hyperparameter	Value
Dataset	CIFAR-10, CIFAR-100
Input	$32\times 32\times 3$ images flattened to $3072$
Model	Residual MLP
Residual stream size ( $d_{e}$ )	$3072$
# residual ( $L$ )	$[2,4]$
Normalization in blocks	RMSNorm (pre-norm, per block)

Appendix

Table 10: Hyperparameters Delta for experiments on CIFAR-10 and CIFAR-100.

Figure 6: Training and test accuracy for a 10-layer residual MLP trained with SGD under varying DeltaNet decay rate

\lambda

. While

\lambda=0

achieves near-perfect training accuracy quickly, test accuracy plateaus at around 67%. Modest DeltaNet-like decay rates significantly improve generalization, as indicated by the reduced train-test gap, which demonstrates the regularizing effect of the update rule. Curves are smoothed using a running average.

Figure 7: Training and test accuracy for a 10-layer residual MLP trained with SGD under varying multiplicative decay rate

\beta

. Similarly to DeltaNet-like decay the training accuracy increases much slower, however, generalization improves only marginally. Curves are smoothed using a running average.

Hyperparameter	Value
Dataset	CIFAR-10
Input	$32\times 32\times 3$ images flattened to $3072$
Model	Residual MLP
Input projection	Linear $(3072\rightarrow 512)$
Residual stream size ( $d_{e}$ )	$512$
# residual blocks	$10$

Appendix

Table 11: Hyperparameters used for DeltaNet-GD experiments on extended training on Cifar-10 with deep nets.

Hyperparameter	Value
Dataset	CIFAR-10, CIFAR-100
Input	$32\times 32\times 3$ images flattened to $3072$
Model	Residual MLP
Residual stream size ( $d_{e}$ )	$3072$
# residual ( $L$ )	$[2,4]$
Normalization in blocks	RMSNorm (pre-norm, per block)

Appendix

Table 12: Hyperparameters MLP Baselines for experiments on CIFAR-10 and CIFAR-100.

Hyperparameter	Value
Model	Residual MLP
Residual stream size ( $d_{e}$ )	$[256,512,1024,4096,8192]$
# residual ( $L$ )	$2$
Normalization	None
MLP activation	ReLU
Optimizer	SGD

Appendix

Table 13: Hyperparameters MLP Baseline for Teacher-Student setup.

Hyperparameter	Value
Dataset	CIFAR-10
Input	$32\times 32\times 3$ images flattened to $3072$
Model	RBF kernel
Classification	$\arg\max$ on predicted output
Optimizer	Exact solver
Bandwidth $\gamma_{\text{mult}}$	$[0.25,0.5,1.0,2.0,3.0,5.0]$

Appendix

Table 14: Hyperparameters RBF kernel regression for experiments on CIFAR-10.

Hyperparameter	Value
Dataset	CIFAR-10
Input	$32\times 32\times 3$ images, $4\times 4$ patches ( $64$ patch tokens + CLS token)
Model	ViT, $2$ transformer blocks
Embedding dimension $d_{e}$	$128$
Attention heads	$8$
Positional embedding	learned

Appendix

Table 15: Hyperparameters Standard ViT for single-epoch experiments on CIFAR-10.

Hyperparameter	Value
Dataset	CIFAR-10
Input	$32\times 32\times 3$ images, $4\times 4$ patches ( $64$ patch tokens + CLS token)
Model	ViT, $2$ transformer blocks, FFN replaced by growing GWS
Embedding dimension $d_{e}$	$128$
Attention heads	$8$
Positional embedding	learned

Appendix

Table 16: Hyperparameters Weighted Softmax ViT for single-epoch experiments on CIFAR-10.

Panel of Figure 8 — Figure 8: Training loss for Weighted Softmax & RBF decreases much faster than that of the other models.

Figure 9: Test accuracy of non-parametric methods (Weighted Softmax, Gauge-Invariant WSM & RBF) as a function of the number of samples seen during training.

Deep Learning Optimization

Deep learning optimization develops algorithms and training strategies that help neural networks converge, generalize, and use compute efficiently.

See all 206 papers

May 19, 2026

Replacement Learning: Training Neural Networks with Fewer Parameters

End-to-end training with full-depth backpropagation remains the dominant paradigm for optimizing deep neural networks, but its efficiency deteriorates as models grow deeper. Since every block must be executed and differentiated under a single global objective, full-depth BP introduces substantial parameter redundancy, activation-memory cost, and training latency, especially when neighboring layers exhibit highly correlated learning patterns. Directly skipping or removing layers can reduce cost, but often weakens representation capacity or requires architecture-specific reuse designs. In this paper, we propose Replacement Learning (RepL), a training-time paradigm that reduces full-depth redundancy by replacing selected blocks rather than simply discarding them. For each removed block, RepL inserts a lightweight computing layer that synthesizes a surrogate operator from the parameters of its adjacent preceding and succeeding blocks through a learnable transformation, and applies the synthesized operator to the preceding activation. In this way, RepL preserves local contextual continuity while avoiding unnecessary full-layer computation. We instantiate RepL for CNNs and ViTs with tailored parameter-fusion blocks that handle convolutional channels, feature resolutions, and transformer submodules. Extensive experiments on CIFAR-10, SVHN, STL-10, ImageNet, COCO, and CityScapes show that RepL reduces trainable parameters, GPU memory usage, and training time while matching or surpassing standard end-to-end training across classification, detection, and segmentation. Additional results on WikiText-2, transfer learning, inference throughput, checkpointing, stochastic depth, and INT8 quantization further demonstrate its generality and compatibility.

Yuming Zhang, Peizhe Wang, Tianyang Han +5

The University of Hong Kong, Hong Kong, China · Southeast University, Nanjing, China

Deep Learning Optimization Neural Network Compression

Ask PDF

Oct 16, 2026

"Deep lifts": Neural networks can generalize even with the dumbest optimizer

Given a labeled dataset

\{(x_i, y_i)\}_{i = 1}^n

, consider drawing a parameter vector

\hatθ

uniformly at random from those that interpolate the data within a width-

W

depth-

D

ReLU architecture with continuous weights bounded within

[-M, M]

. We study the generalization properties of

\hatθ

in the case where the dataset can, in fact, be explained by a smaller

w

-by-

d

ReLU network. Despite the obliviousness of the selection rule for

\hatθ

, it adaptively generalizes at the rate

\tilde O\left(\frac{wdWD}{n}\right)

, rather than being governed by the VC dimension of the full

W

-by-

D

hypothesis class. The reason is due to a particular inductive bias in ReLU networks: the

w

-by-

d

parameter witnesses an exponential volume of

W

-by-

D

interpolating parameters, which dwarfs the volume of poorly generalizing parameters (conditioned on interpolation). We explicitly construct this set of near-equivalent

W

-by-

D

parameters by a proof technique called deep lifting, and then tie the achievable volume to stability properties of the

w

-by-

d

witness. Examples show that these techniques allow us to convert approximation-theoretic results into generalization guarantees for nonparametric classes. We further show that generalization of uniform-random interpolators can survive label corruptions.

Sunil Madhow, Yu-Xiang Wang

UC San Diego

Neural Network Generalization Generalization Bounds Overparameterized Neural Network ReLU Neural Networks Benign Overfitting

Ask PDF

Oct 9, 2026

Layer-wise Input Recoverability in Deep Neural Networks

The inner workings of deep neural networks are still not well understood. The information bottleneck framework, which characterizes deep networks by the mutual information between their hidden representations and the network's input, has been used to identify compression of input information with respect to an input image and output labels. We instead study how much of a layer's own input is recoverable from the layer's output in convolutional neural networks (CNNs). A convolutional layer applies its filters to one input patch at a time, so we evaluate recoverability on sampled input patches with a reconstruction-based probe and not on the entire feature map. The probe's size thus depends on channels and kernel size and not on feature map size, which lets our method scale to deep CNNs with wide feature maps. For each layer, we compare pretrained, random, and reconstruction-optimized encoders, using linear and linear-plus-residual-MLP decoders. We measure recoverability as

R^2

between input patches and their reconstructions on a held-out set. Across nine ImageNet CNNs, recoverability generally decreases with depth but is more consistently predicted by the compression of the layer. Learning-attributed recoverability, the difference between pretrained and random encoders, is not uniform across architectures. We introduce single-property random controls, which transfer one pretrained property onto a random encoder, and find that each control changes recoverability with a consistent sign across models: copying pretrained singular vectors tends to raise recoverability, while copying the pretrained singular-value spectrum or activation sparsity lowers it.

Maxwell Spiridonov, Alireza Rahmansetayesh, Pierre Baldi

Department of Computer Science, University of California, Irvine

Convolutional Neural Networks Neural Network Interpretability Representation Learning

Ask PDF

Abstract

Deep Learning Optimization

Explore similar work

Replacement Learning: Training Neural Networks with Fewer Parameters

"Deep lifts": Neural networks can generalize even with the dumbest optimizer

Layer-wise Input Recoverability in Deep Neural Networks