Exact Dynamics and Finite-Sample Trajectory Recovery of Linear Recursive Feature Machines

Andrew Cheng; Bobak T. Kiani; Yue M. Lu; Adityanarayanan Radhakrishnan

cs.LGOct 6, 2026

Exact Dynamics and Finite-Sample Trajectory Recovery of Linear Recursive Feature Machines

Authors: Andrew Cheng, Bobak T. Kiani, Yue M. Lu, Adityanarayanan Radhakrishnan

Organizations: Harvard University · Bowdoin College · Massachusetts Institute of Technology

Abstract

Recursive feature machines (RFMs) learn representations of data by alternating between fitting a predictor to a dataset and updating features of that predictor using the average gradient outer product (AGOP). Connections between AGOPs and feature learning in neural networks motivate linear RFMs as a simple setting for analyzing how representations evolve during training. Here, we study the dynamics and statistics of linear RFM in noisy multi-output regression with isotropic sub-Gaussian input data and targets generated by a low-rank teacher matrix of dimension $d$ . We extend the known connection between linear RFM and iteratively reweighted least squares from the interpolating setting to ridge-regularized multi-output regression with noise. We show that the learned feature matrix remains close to its infinite-data ideal counterpart at every iteration. Namely, for $n$ samples, we show the error in the feature matrix decays as $O(\sqrt{d/n})$ with high probability. Experiments on real-world text and single-cell gene-expression data illustrate the features learned by this simple linear model.

Figures & tables

Dataset	$n$	$d$	Streaming stage	Later RFM updates
OpenWebText	$10^{9}$	$20{,}000$	2.01 hours, 16 CPU cores	10 updates in 1.5 min on A100
HLCA	$10^{9}$	$20{,}000$	8.51 hours, 16 CPU cores	10 updates in 1.5 min on A100

Table 1: Total runtime split between streaming and RFM implementation.

Figure 1: OpenWebText feature evolution under regularized lin-RFM. A (top-left): Stable-rank trajectory of the floor-subtracted feature matrix $\widetilde{B}_{t}=B_{t}-\epsilon^{1/4}I_{d}$ , highlighting the initialization and the learned checkpoints $t=1$ and $t=5$ . B (bottom-left): Maximum operator-norm discrepancy over the evaluated iterations between finite-sample feature matrices and a high-sample reference trajectory, plotted against sample size $n$ and compared with an $n^{-1/2}$ reference rate. C (right): Selected nearest-neighbor neighborhoods for query words at $t=1$ and $t=5$ , together with randomly selected words at $t=0$ , illustrating the task-dependent structure that emerges across iterations.

Figure 2: HLCA feature evolution under regularized lin-RFM. A (top-left): Stable-rank trajectory of the floor-subtracted feature matrix $\widetilde{B}_{t}=B_{t}-\epsilon^{1/4}I_{d}$ , highlighting the initialization and the learned checkpoints $t=1$ and $t=5$ . B (bottom-left): Maximum operator-norm discrepancy over the evaluated iterations between finite-sample feature matrices and a high-sample reference trajectory, plotted against sample size $n$ and compared with an $n^{-1/2}$ reference rate. C (right): Selected nearest-neighbor neighborhoods for query genes at $t=1$ and $t=5$ , illustrating the task-dependent structure learned across iterations.

Figure 3: Trajectory recovery and convergence of regularized lin-RFM. All panels use the model

y=Ex+\xi

with Gaussian inputs and independent Gaussian noise, reusing the same dataset at every update. Top left: Maximum predictor discrepancy

\max_{0\leq t\leq 10}\|F_{t}-\bar{F}_{t}\|_{\mathrm{op}}

against sample size

n

. Points show medians over 20 independent datasets per sample size, with a 10th–90th percentile band. The observed decay is consistent with the

n^{-1/2}

rate in Corollary 6 . Top right: Distance

\|F_{t}-F_{500}\|_{\mathrm{op}}

to the final predictor from a 500-iteration run on the same dataset, for

n=100,300,1000

. Curves show medians over 16 datasets per sample size and exhibit approximately linear convergence before leveling off at numerical precision. Both top panels use isotropic inputs,

d=p=50

, teacher rank

3

(\lambda,\epsilon)=(0.3,0.04)

, and noise standard deviation

0.5

. Bottom: Estimated contraction factors

\hat{\rho}

against

\kappa(\Sigma)

for four covariance spectra: linearly spaced, geometrically spaced, one weak direction, and half weak directions. Each covariance is normalized to

\operatorname{tr}(\Sigma)=d

. Rates are estimated from the late-iteration decay of the Frobenius-norm distance to a numerical endpoint; markers show the estimates and smooth curves are descriptive sigmoid fits. Values closer to one indicate slower convergence. The separation between curves shows that equal condition numbers can yield different rates in these experiments. Here

d=p=80

n=4000

, teacher rank

5

(\lambda,\epsilon)=(0.01,0.001)

, and noise standard deviation

0.05

. The covariance constructions are given in Appendix B.3.

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 4: Left: Stable-rank trajectory of

\widetilde{B}_{t}=B_{t}-\epsilon^{1/4}I

, with checkpoints

A,B,C

. Right: nearest-neighbor neighborhoods at those checkpoints.

Figure 5: TinyStories feature tracking. We first compute a reference trajectory

B_{t}^{(\mathrm{Ref})}

using

n_{\mathrm{ref}}=10^{7}

TinyStories context-to-center samples. For each smaller sample size

n

, we draw random subsamples from these

10^{7}

examples without replacement, run regularized Lin-RFM for

T=10

iterations, and report

\sup_{0\leq t\leq T}\|B_{t}^{(n)}-B_{t}^{(\mathrm{Ref})}\|_{\mathrm{op}}

. The empirical log-log slope is close to the reference

n^{-1/2}

rate, even though the TinyStories design is sparse, highly anisotropic, and outside the assumptions of Theorem 5 .

Figure 6: HuffPost News test MSE. 80/20 train-test split with

n_{\mathrm{train}}=167,621

articles and

n_{\mathrm{test}}=41,906

. We perform Algorithm 2 on

(X_{\mathrm{train}},Y_{\mathrm{train}})

and evaluate on

(X_{\mathrm{test}},Y_{\mathrm{test}})

. We observe that our algorithm learns to generalize with respect to the MSE loss in Figure 6 . The ridge baseline is the initial predictor at

t=0

; subsequent lin-RFM iterations reduce test MSE relative to this initial ridge fit.

Figure 7: Left: Stable-rank trajectory of

\widetilde{B}_{t}=B_{t}-\epsilon^{1/4}I

., with checkpoints

A,B,C

. Right: nearest-neighbor neighborhoods induced by

B_{t}

with cosine similarity scores at those checkpoints (see Appendix B.2 for details).

Figure 8: HuffPost News feature tracking. We compute a large-sample reference trajectory

B_{t}^{(\mathrm{Ref})}

using the full HuffPost News dataset

(n=209,527)

. For each sample size

n

, we draw subsamples without replacement, run regularized Lin-RFM for

T=10

iterations, and report

\sup_{0\leq t\leq T}\|B_{t}^{(n)}-B_{t}^{(\mathrm{Ref})}\|_{\mathrm{op}}

. The empirical log-log slope is

-0.493

, close to the reference

n^{-1/2}

rate, even though the HuffPost design is sparse, anisotropic, and outside the assumptions of Theorem 5 .

Query	$t=1$	$t=5$
HBA1	HBA2 (0.837) HBQ1 (0.822) TRDC (0.816) LINC01505 (0.815) LIM2 (0.807)	LINC02446 (0.960) RGL4 (0.959) TESPA1 (0.956) ZNF683 (0.956) KLRC2 (0.953)
CFTR	TMEM92 (0.829) C16orf89 (0.804) SLC4A4 (0.793) PON3 (0.785) SHISA2 (0.771)	FREM2 (0.968) PDP2 (0.964) SLC6A20 (0.957) FOXP4 (0.957) RAB40C (0.952)
AVIL	GARIN1A (0.854) LINC02021 (0.831) ALG1L2 (0.819) CNN3-DT (0.791) DISP2 (0.791)	EBF4 (0.929) TMEM8B (0.919) FRAS1 (0.911) DLG2 (0.903) LINC01473 (0.899)
FOXJ1	LRRC71 (0.997) SPAG8 (0.995) DCDC2B (0.994) PRR29 (0.994) GPR162 (0.994)	FAM216B (0.995) MS4A8 (0.995) LDLRAD1 (0.994) CFAP43 (0.994) RSPH9 (0.993)
CHGA	ASCL1 (0.998) SCG3 (0.996) SST (0.993) CALCA (0.992) KIF1A (0.990)	GFRA3 (0.996) SCG5 (0.996) CHGB (0.996) ASCL1 (0.995) SCG3 (0.995)
SFTPA1	LAMP3 (0.697) ABCA3 (0.642) NAPSA (0.640) SFTPD (0.624) PEBP4 (0.591)	SFTPA2 (0.989) NAPSA (0.987) SFTA2 (0.983) SFTPB (0.982) LPCAT1 (0.980)

Appendix

Table 2: HLCA TF-IDF leave-one-out RFM gene-neighbor evolution. For each query gene, we report the top-5 nearest neighbors at

t=1

and

t=5

, with cosine similarities in parentheses. Neighbors are computed from the

B_{t}

-induced gene embeddings with embedding dimension

k=256

. Details in Appendix C.2 .

Query	$\mathbf{t=1}$	$\mathbf{t=5}$
crushed	lightly (0.796) toss (0.780) dip (0.777) tin (0.776) toast (0.773)	chopped (0.731) lemon (0.725) powder (0.709) dried (0.707) flour (0.699)
buys	gamble (0.822) buck (0.778) needless (0.761) theirs (0.759) geterrorstate (0.753)	sells (0.685) buying (0.662) buyer (0.650) seller (0.616) sell (0.602)
identifies	relates (0.786) acknowledges (0.780) recognizes (0.773) emphasize (0.763) illustrates (0.757)	identities (0.586) stereotypes (0.567) unidentified (0.567) identifying (0.565) identity (0.547)
bounds	alignment (0.773) secondly (0.770) needless (0.765) settling (0.764) illustrates (0.753)	ball (0.592) foul (0.584) defenders (0.573) basket (0.562) turnover (0.559)
innocence	awaiting (0.796) prosecuted (0.765) falsely (0.761) kidnapping (0.751) courtroom (0.750)	conviction (0.751) prosecution (0.738) jury (0.735) defendant (0.725) convictions (0.692)
rouge	needless (0.765) raymond (0.758) geterrorstate (0.758) apiculture (0.756) glen (0.756)	suburb (0.497) unarmed (0.464) qu (0.463) midwest (0.455) rapids (0.438)

Appendix

Table 3: Top-50 OpenWebText nearest-neighbor changes with cosine similarities. Shown are the top-5 neighbors at

t=1

and

t=5

for the top-ranked query words. We picked this list based off the top-50 Jaccard disimilarity scores between the

t=1

and

t=5

nearest neighbor sets.

Explore similar work

K-Inverse-RFM: A Modified RFM that Bridges the Gap to Neural Networks for Data-Corrupted Mathematical Tasks
Jul 1, 2026Gil PasternakImbalanced Classification Neural Network
UNIVERSITY OF CALIFORNIA SAN DIEGO
Average Gradient Outer Product in kernel regression provably recovers the central subspace for multi-index models
May 14, 2026Libin Zhu, Damek Davis, Dmitriy Drusvyatskiy +1Kernel Ridge Regression Kernel Method
Department of Mathematics, University of Washington, Seattle, WA 98195 · Wharton Department of Statistics and Data Science, University of Pennsylvania, Philadelphia, PA 19104, USA · Department of Mathematics, U. Washington, Seattle, WA 98195 +1
Scaling Laws from Sequential Feature Recovery: A Solvable Hierarchical Model
May 14, 2026Arie Wortsman-Zurich, Hugo Tabanelli, Yatin Dandi +2Feature Learning Scaling Laws
Département d’Informatique, Ecole Normale Supérieure, PSL & CNRS · Information Learning and Physics Laboratory, École Polytechnique Fédérale de Lausanne (EPFL) · Statistical Physics of Computation Laboratory, École Polytechnique Fédérale de Lausanne (EPFL)

Jul 1, 2026, cs.LGJ/K move · Enter open · S save

K-Inverse-RFM: A Modified RFM that Bridges the Gap to Neural Networks for Data-Corrupted Mathematical Tasks

Gil Pasternak

UNIVERSITY OF CALIFORNIA SAN DIEGO

Imbalanced Classification Neural Network Feed-Forward

Open paper AskPDF arXiv

Recursive Feature Machines (RFMs) are a class of kernel machines that utilize the Average Gradient Outer Product (AGOP) as a mechanism for feature learning. They have been shown to effectively replicate the learning dynamics and feature representations of Feedforward Neural Networks (FNNs) across various settings. However, despite comparable capacity for feature learning and the similarities in the features they acquire, RFMs exhibit significantly lower performance than neural networks in certain data-corrupted scenarios. In this work, we investigate these limitations in mathematical problems. As a solution, we introduce a remarkably effective transformation applied to the training labels which promotes learning in noisy, complexly represented, and class-imbalanced data. This simple yet powerful adjustment enables RFMs to close the performance gap with FNNs and, in some cases, even surpass them.

Abstract

Explore similar work

K-Inverse-RFM: A Modified RFM that Bridges the Gap to Neural Networks for Data-Corrupted Mathematical Tasks

Average Gradient Outer Product in kernel regression provably recovers the central subspace for multi-index models

Scaling Laws from Sequential Feature Recovery: A Solvable Hierarchical Model