Nonparametric Teaching of Attention Learners

Chen Zhang; Jianghui Wang; Bingyang Cheng; Zhongtao Chen; Wendong XU; Cong Wang; Marco Canini; Francesco Orabona

cs.LGFeb 24, 2026

Nonparametric Teaching of Attention Learners

Authors: Chen Zhang, Jianghui Wang, Bingyang Cheng, Zhongtao Chen, Wendong XU, Cong Wang, Marco Canini, Francesco Orabona,

Organizations: The University of Hong Kong · King Abdullah University of Science and Technology · Independent Researcher

Abstract

Attention learners, neural networks built on the attention mechanism, e.g., transformers, excel at learning the implicit relationships that relate sequences to their corresponding properties, e.g., mapping a given sequence of tokens to the probability of the next token. However, the learning process tends to be costly. To address this, we present a novel paradigm named Attention Neural Teaching (AtteNT) that reinterprets the learning process through a nonparametric teaching perspective. Specifically, the latter provides a theoretical framework for teaching mappings that are implicitly defined (i.e., nonparametric) via example selection. Such an implicit mapping is embodied through a dense set of sequence-property pairs, with the AtteNT teacher selecting a subset to accelerate convergence in attention learner training. By analytically investigating the role of attention on parameter-based gradient descent during training, and recasting the evolution of attention learners, shaped by parameter updates, through functional gradient descent in nonparametric teaching, we show for the first time that teaching attention learners is consistent with teaching importance-adaptive nonparametric learners. These new findings readily commit AtteNT to enhancing learning efficiency of attention learners. Specifically, we observe training time reductions of 13.01% for LLMs and 20.58% for ViTs, spanning both fine-tuning and training-from-scratch regimes. Crucially, these gains are achieved without compromising accuracy; in fact, performance is consistently preserved and often enhanced across a diverse set of downstream tasks.

Figures & tables

Figure 1: An illustration of the workflow for an attention neural network with an input sequence

\bm{S}

Model	AtteNT	Avg. Time(↓)	GSM8K(↑)	MATH(↑)	HumanEval(↑)	MBPP(↑)	MT-Bench(↑)
LLaMA 2-7B	w/o	246±1m	42.96±0.12	5.06±0.16	18.35±0.31	35.65±0.25	4.58±0.01
LLaMA 2-7B	w	213±2m	43.45±0.55	6.48±0.24	21.80±0.38	37.61±0.42	4.49±0.02
Mistral-7B	w/o	204±2m	69.13±0.22	20.06±0.20	43.42±0.14	58.52±0.13	5.03±0.05
Mistral-7B	w	180±2m	71.26±0.23	23.12±0.44	46.55±0.25	61.74±0.54	5.32±0.04
Gemma-7B	w/o	228±2m	75.23±0.45	30.52±0.48	53.83±0.27	65.69±0.29	5.42±0.04
Gemma-7B	w	201±2m	77.74±0.32	31.40±0.36	54.26±0.28	66.28±0.46	5.44±0.08

Table 1: AtteNT on NLG tasks. The results are averaged over three runs, with standard deviations included. The GSM8K and MATH datasets share a math fine-tuned model, while HumanEval and MBPP use a code fine-tuned model. MT-Bench utilizes a conversation fine-tuned model. The "Avg. time" represents the average fine-tuning time for the three models.

Model	AtteNT	Pretraining Time(↓)	ImageNetS50(↑)	NYUv2(S)(↑)	NYUv2(D)(↑)
Multi-Modal MAE	w/o	1234m	92.2	51.9	52.1
Multi-Modal MAE	w	980m (-20.58%)	92.3	52.6	57.2

Table 2: AtteNT across various CV downstream tasks. ImageNetS50 uses 50 categories from ImageNet for classification, evaluated by accuracy. NYUv2(S) is a semantic segmentation task with mIoU as the metric. NYUv2(D) involves depth estimation, evaluated using the

\delta_{1}

metric, which measures the percentage of pixels with an error ratio below 1.25 ( Doersch and Zisserman, 2017 ) .

Pre-training				Downstream
Ratio	Interval	Selection	Training time(↓)	ImageNetS50(↑)	NYUv2(S)(↑)	NYUv2(D)(↑)
-	-	-	1234m	92.2	51.9	52.1
Cosine	Incremental	Random	966m	88.6	45.3	49.6
Cosine	Incremental	Soft	995m	92.1	52.2	58.8
Cosine	Fixed	Soft	1301m	93.2	53.6	61.4
Incremental	Incremental	Soft	980m	92.3	52.6	57.2

Table 3: Ablation study of AtteNT pre-training configurations. Ratio controls how the fraction of selected samples increases over epochs. Interval denotes how often the subset is re-sampled. Selection specifies the sampling strategy: Random (no difficulty prior), Hard (selects only difficult samples), and Soft (Gumbel-Top-k difficulty-aware sampling). The configuration (Incremental, Incremental, Soft) in the red color row is adopted as our final AtteNT setting, as it simultaneously reduces pre-training time and improves performance on all downstream tasks.

Appendix figures & tables10 assets

Supplementary material from the paper’s appendix.

Appendix

Notation	Description
$\bm{S}_{S\times d}$	Matrix containing all feature vectors from the ordered sequence $(\bm{x}_{1},\dots,\bm{x}_{S})$ , with shape $S\times d$
$[x_{s,j}]_{d}^{\top}$	$d$ -dimensional feature vector for the $s$ -th element, with components $x_{s,j}$
$\bm{x}$	Short form for $[x_{j}]_{d}$
$\bm{S}_{(s,:)}$	The $s$ -th row of $\bm{S}$ , representing the feature vector for the $s$ -th element
$\bm{S}_{(:,i)}$	The $i$ -th column of $\bm{S}$ , which represents the $i$ -th feature across all elements)
$\bm{e}_{i}$	The $i$ -th basis vector, having a value of 1 at the $i$ -th position and 0 elsewhere

Appendix

Table 4: Summary of Key Notations.

Panel of Figure 2 — Figure 2: An illustration of the workflow for different multi-output attention learners, with input sequence $\bm{S}$ and $\bm{S}^{\prime}$ (in the case of cross-attention).

Figure 3: Graphical depiction of the ANTK computation process: \(K_{\theta}(\bm{S}_{S},\bm{S}^{\prime}_{S^{\prime}})=\left\langle\frac{\partial f_{\theta}(\bm{S})}{\partial\theta},\frac{\partial f_{\theta}(\bm{S}^{\prime})}{\partial\theta}\right\rangle=\bigg[\frac{\partial{f_{\theta}(\bm{S})}_{(i,:)}}{\partial\bm{W}^{V}_{(1)}}\frac{\partial{f_{\theta}(\bm{S}^{\prime})}_{(j,:)}}{\partial\bm{W}^{V}_{(1)}}+\dots+\frac{\partial{f_{\theta}(\bm{S})}_{(i,:)}}{\partial\bm{W}^{V}_{(d)}}\frac{\partial{f_{\theta}(\bm{S}^{\prime})}_{(j,:)}}{\partial\bm{W}^{V}_{(d)}}+\frac{\partial{f_{\theta}(\bm{S})}_{(i,:)}}{\partial\bm{W}^{Q}_{(1,1)}}\frac{\partial{f_{\theta}(\bm{S}^{\prime})}_{(j,:)}}{\partial\bm{W}^{Q}_{(1,1)}}+\dots+\frac{\partial{f_{\theta}(\bm{S})}_{(i,:)}}{\partial\bm{W}^{Q}_{(d,p)}}\frac{\partial{f_{\theta}(\bm{S}^{\prime})}_{(j,:)}}{\partial\bm{W}^{Q}_{(d,p)}}+\frac{\partial{f_{\theta}(\bm{S})}_{(i,:)}}{\partial\bm{W}^{K}_{(1,1)}}\frac{\partial{f_{\theta}(\bm{S}^{\prime})}_{(j,:)}}{\partial\bm{W}^{K}_{(1,1)}}+\dots+\frac{\partial{f_{\theta}(\bm{S})}_{(i,:)}}{\partial\bm{W}^{K}_{(d,p)}}\frac{\partial{f_{\theta}(\bm{S}^{\prime})}_{(j,:)}}{\partial\bm{W}^{K}_{(d,p)}}\bigg]_{S\times S^{\prime};i\in\mathbb{N}_{S},j\in\mathbb{N}_{S^{\prime}}}\) . — Figure 3: Graphical depiction of the ANTK computation process: $K_{\theta}(\bm{S}_{S},\bm{S}^{\prime}_{S^{\prime}})=\left\langle\frac{\partial f_{\theta}(\bm{S})}{\partial\theta},\frac{\partial f_{\theta}(\bm{S}^{\prime})}{\partial\theta}\right\rangle=\bigg[\frac{\partial{f_{\theta}(\bm{S})}_{(i,:)}}{\partial\bm{W}^{V}_{(1)}}\frac{\partial{f_{\theta}(\bm{S}^{\prime})}_{(j,:)}}{\partial\bm{W}^{V}_{(1)}}+\dots+\frac{\partial{f_{\theta}(\bm{S})}_{(i,:)}}{\partial\bm{W}^{V}_{(d)}}\frac{\partial{f_{\theta}(\bm{S}^{\prime})}_{(j,:)}}{\partial\bm{W}^{V}_{(d)}}+\frac{\partial{f_{\theta}(\bm{S})}_{(i,:)}}{\partial\bm{W}^{Q}_{(1,1)}}\frac{\partial{f_{\theta}(\bm{S}^{\prime})}_{(j,:)}}{\partial\bm{W}^{Q}_{(1,1)}}+\dots+\frac{\partial{f_{\theta}(\bm{S})}_{(i,:)}}{\partial\bm{W}^{Q}_{(d,p)}}\frac{\partial{f_{\theta}(\bm{S}^{\prime})}_{(j,:)}}{\partial\bm{W}^{Q}_{(d,p)}}+\frac{\partial{f_{\theta}(\bm{S})}_{(i,:)}}{\partial\bm{W}^{K}_{(1,1)}}\frac{\partial{f_{\theta}(\bm{S}^{\prime})}_{(j,:)}}{\partial\bm{W}^{K}_{(1,1)}}+\dots+\frac{\partial{f_{\theta}(\bm{S})}_{(i,:)}}{\partial\bm{W}^{K}_{(d,p)}}\frac{\partial{f_{\theta}(\bm{S}^{\prime})}_{(j,:)}}{\partial\bm{W}^{K}_{(d,p)}}\bigg]_{S\times S^{\prime};i\in\mathbb{N}_{S},j\in\mathbb{N}_{S^{\prime}}}$ .

Hyperparam	Baseline	AtteNT
Batch Size	2048	2048
Learning Rate	1e-4	1e-4
Min Learning Rate	1e-6	1e-6
Weight Decay	0.05	0.05
Adamw $\epsilon$	1e-8	1e-8
Adamw $\beta_{1}$	0.9	0.9

Appendix

Table 5: Hyperparameters for pre-training Multi-Modal MAE.

Hyperparam	ImageNetS50	NYUv2(S)	NYUv2(D)
Epoch	100	100	2000
Warm up Epoch	5	20	100
Batch Size	1024	1024	2048
Learning Rate	4e-3	1e-4	1e-4
Min Learning Rate	1e-6	1e-6	0
Weight Decay	0.05	0.05	1e-4

Appendix

Table 6: Hyperparameters for fine-tuning Multi-Modal MAE on various downtasks. The augmentation strategy LSJ is large scale jittering ( Ghiasi et al., 2021 ) . We use drop path ( Huang et al., 2016 ) in classification and semantic segmentation tasks.

Panel of Table 7 — Table 7: Sample Ratio Study of ViT models.

Figure 4: Downstream Task Performance vs Sample Ratio.

Methods	Time(↓)	ImageNetS50(↑)	NYUv2(S)(↑)	NYUv2(D)(↑)
AtteNT(Ours)	980m	92.3	52.6	57.2
Class Weight Sampling	1108m	90.4	48.2	52.0
Fixed Weight Sampling	1065m	89.6	49.7	54.6
GradNorm Sampling ( Chen et al., 2018 )	1112m	91.9	52.4	55.8

Appendix

Table 8: Performance Comparison of Different Methods.

Figure 5: Frobenius norm of the difference between the empirical NTK at different training steps and the canonical kernel.

Figure 6: Evolution of the empirical $10\times 10$ NTK matrix during training. Color represents value $K_{\theta^{t}}(\bm{S}_{i},\bm{S}_{j})$ . The matrix stabilizes visually after 200 epochs and shows negligible changes thereafter.

Training Data Selection

Training data selection studies how to choose, filter, weight, or schedule examples to improve learning while reducing data and compute costs.

See all 184 papers

Why Attend to Everything? Focus is the Key
Mar 12, 2026Hengshuai Yao, Xing Chen, Ahmed Murtadha +8Attention Mechanisms Efficient Attention
Sapient Intelligence · University of Alberta · Tsinghua University
Parallax: Parameterized Local Linear Attention for Language Modeling
May 27, 2026Yifei Zuo, Dhruv Pai, Zhichen Zeng +3Language Model Pretraining Self-Attention
Northwestern University · Tilde Research · University of Washington
Faster Query-Key Learning Sharpens Attention in Self-Attention Models
Aug 7, 2026Rahul Vashisht, Harish G. RamaswamySelf-Attention
Department of CSE, IIT Madras, India · Department of DSAI, IIT Madras, India.

Mar 12, 2026J/K move · Enter open · S save

Why Attend to Everything? Focus is the Key

Hengshuai Yao, Xing Chen, Ahmed Murtadha, Jin Li+7

Sapient Intelligence · University of Alberta · Tsinghua University

Attention Mechanisms Efficient Attention Attention Approximation Sparse Attention Long-Context Modeling

Open paper AskPDF arXiv

Standard attention scales quadratically with sequence length. Efficient attention methods reduce this O(n^2) cost, but when retrofitted into pretrained models, they often degrade perplexity, downstream accuracy, or both. We introduce Focus, a method that learns which token pairs matter. Focus adds a small set of learnable centroids--as few as 148K parameters per layer--that act as gates: only token pairs belonging to the same centroid group attend to each other over long ranges. Focus is composable: it can be added to any pretrained model by training only the centroids while keeping all original weights frozen. Experiments show that composing Focus onto pretrained models yields zero degradation on downstream benchmarks across model sizes from 124M to 70B parameters and five attention architectures. Surprisingly, sparse Focus attention outperforms full attention at 124M scale (30.3 vs. 31.4 perplexity) and matches full attention when trained from scratch at 7B scale (13.82 vs. 13.89). Focus is also fast: top-k group membership gives a 2x speedup with better quality than the original pretrained model. Using our FlashAttention decomposition, Focus achieves an 8.6x speedup at 1M tokens without custom kernels.

Abstract

Training Data Selection

Explore similar work

Why Attend to Everything? Focus is the Key

Parallax: Parameterized Local Linear Attention for Language Modeling

Faster Query-Key Learning Sharpens Attention in Self-Attention Models