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

Replacement Learning: Training Neural Networks with Fewer Parameters
May 19, 2026Yuming Zhang, Peizhe Wang, Tianyang Han +5Deep Learning Optimization Neural Network Compression
The University of Hong Kong, Hong Kong, China · Southeast University, Nanjing, China
"Deep lifts": Neural networks can generalize even with the dumbest optimizer
Oct 16, 2026Sunil Madhow, Yu-Xiang WangNeural Network Generalization Generalization Bounds
UC San Diego
Layer-wise Input Recoverability in Deep Neural Networks
Oct 9, 2026Maxwell Spiridonov, Alireza Rahmansetayesh, Pierre BaldiConvolutional Neural Networks Neural Network Interpretability
Department of Computer Science, University of California, Irvine

May 19, 2026J/K move · Enter open · S save

Replacement Learning: Training Neural Networks with Fewer Parameters

Yuming Zhang, Peizhe Wang, Tianyang Han, Hengyu Shi+4

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

Deep Learning Optimization Neural Network Compression

Open paper AskPDF arXiv

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.

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