Princeton University · Research Talk

Auto
Research

Toward autonomous superintelligence: principled architectures and agentic RL that automate AI research.

Yifan Zhang Princeton University · 2026 · ~60 min · extended

Press → or Space to begin  ·  ? for help

Research agenda

Toward autonomous superintelligence

The route is AutoResearch: automating AI research itself, so improvements compound. That goal sets three pillars.

I · Architectures Long context & continual learning Structured state that holds long rollouts and keeps learning. TPA · GRAPE · DeepLoop · HLA · Fast-Weight Attention
II · Agentic RL Scaled to long-horizon tasks Objectives derived, not tuned; training systems that scale. RPG · SDPG · Molt · Kernel Mismatch
III · AutoResearch Automating AI research Agents and harnesses that coevolve with the model. Proposer-Verifier · Meta Prompting · MathCode · Web World Models

NeurIPS Spotlight · methods adopted by Frontier Labs

The thesis

Superintelligence as a loop, not a leap

Architectures long context · continual learning Agentic RL long-horizon training Research agents propose · verify · ship Better models and better methods improvements compound: the loop is AutoResearch
  • 01The goal is AutoResearch: automating AI research itself, so progress compounds.
  • 02Architectures and agentic RL are not separate interests; they are what the loop demands.

Selected Works

Selected Works

[Kernel Mismatch] Reliable RL Scaling Requires Accounting for Prefill-Decode Kernel Mismatch
Yifan Zhang et al.
Tech Report · 08/2026
[Molt] Scalable PyTorch-Native Agentic RL Training
Jian Hu, Huiying Li, Hao Zhang, Binfeng Xu, Yifan Zhang, et al.
NVIDIA · arXiv:2607.21653
[DeepLoop] Depth Scaling for Looped Transformers
Shuzhen Li, Yifan Zhang, Jiacheng Guo, Quanquan Gu, Mengdi Wang
arXiv:2607.13491
[SDPG] Self-Distilled Policy Gradient
Yifeng Liu*, Shiyuan Zhang*, Yifan Zhang*, Quanquan Gu
arXiv:2606.04036
[Falcon] Fast-Weight Attention for Continual Learning
Yifan Zhang et al.
Preprint
[FlashSampling] Fast & Memory-Efficient Exact Sampling
Tomas Ruiz*, Zhen Qin*, Yifan Zhang†, et al.
arXiv:2603.15854
[DDL] Deep Delta Learning
Yifan Zhang, Yifeng Liu, Mengdi Wang, Quanquan Gu
arXiv:2601.00417
[GRAPE] Group Representational Position Encoding
Yifan Zhang, Zixiang Chen, Yifeng Liu, et al.
ICLR 2026
[RPG] KL-Regularized Policy Gradient for LLM Reasoning
Yifan Zhang*, Yifeng Liu*, Huizhuo Yuan, et al.
ICLR 2026 · Tinker · DeepSeek-V3.2
[TPA] Tensor Product Attention Is All You Need
Yifan Zhang*, Yifeng Liu*, Huizhuo Yuan, et al.
NeurIPS 2025 Spotlight

* equal contribution · † corresponding author

Impact

Research that ships

1,594
Citations (Google Scholar)
19
Peer-reviewed papers (NeurIPS · ICLR · ICML · ACL · AAAI · TMLR)
15
Peer-reviewed first-author papers (incl. co-first)
2
Spotlights: TPA (NeurIPS 2025) & Beyond Squared Error (ICLR 2025)
5
GRAPE — OpenAI · SpaceXAI · Thinking Machines · Jane Street; RPG — Thinking Machines · DeepSeek
2.8k+
GitHub stars on open-source projects

✦ Most-cited704 Cumulative Reasoning · 99 Meta Prompting · 97 Iterative Question Composing · 51 AutoMathText · 47 Contrastive Learning Theory · 46 Tensor Product Attention · 40 Diagram of Thought · 37 General Preference Model

★ Most-starred584 MathCode · 458 TPA · 355 Deep Delta Learning · 308 Cumulative Reasoning · 302 Meta Prompting · 101 HLA · 95 GRAPE · 91 AutoMathText

The question

What would it take to automate AI research itself?

  • 01Models that hold long horizons: today's attention is memory-bound, and learning stops after pretraining.
  • 02RL that scales: today's policy gradients are a brittle patchwork of heuristics.
  • 03A loop that closes: model and harness must coevolve for improvements to compound.

My research makes each piece principled, and ships the result.

Pillar I

Architectures

Structure the state — expressive and efficient.

The problem

Attention is memory-bound

  • 01The KV cache grows linearly with context — the wall for long-context and inference.
  • 02Compute is quadratic in sequence length.
  • 03Today's fixes (MQA / GQA) trade quality for memory.

What if the state itself were low-rank?

NeurIPS 2025 · Spotlight

Tensor Product Attention

Factorize queries, keys, and values as contextual tensor products — compressing the KV cache by up to ~10×.

$\;Q_t=\sum_{r=1}^{R}\,a^{Q}_{t,r}\otimes b^{Q}_{t,r}\,,\qquad$ likewise for $K_t,\,V_t$ — rank $R \ll d$.
  • →Longer context at fixed memory; unifies MHA / MQA / GQA as special cases.
  • →Compatible with any positional encoding (RoPE, ALiBi, additive RPE, GRAPE); the T6 backbone trains better at equal budget.

"Tensor Product Attention Is All You Need."

Pillar I · Tensor Product Attention

The KV cache is the long-context wall

Standard KV cache (MHA) TPA: up to ~10× smaller, equal quality context length KV-cache memory
  • →Queries, keys, and values factor as contextual tensor products of rank R ≪ d.
  • →Unifies MHA / MQA / GQA as special cases; compatible with any positional encoding, e.g., RoPE, ALiBi, additive RPE, GRAPE.

NeurIPS 2025 Spotlight · 458 GitHub stars

Pillar I · TPA · FlashTPA Decoding

FlashTPA: decode on the factors

query factors AQ, BQ K-factor cache · rank Rₖ V-factor cache · rank Rᵥ einsum chain scores on the factors online softmax masked · stable LSE output the full Q, K, V tensors are never materialized
  • →Per-token attention FLOPs scale as $\Theta(M(R_Q R_K D + H R_Q R_K + H R_V E))$: with small ranks TPA saves prefill and decode compute, not just KV-cache memory.
  • →Benchmarked against FlashMHA, FlashGQA, FlashMQA, and FlashMLA: faster decoding at long sequence lengths (up to 512K tokens), with Triton decode and prefill kernels in the repo.

Algorithms 2 & 3 in the paper · arXiv:2501.06425 · CUDA kernel in progress

Pillar I · FlashSampling · arXiv:2603.15854

FlashSampling: sample without the logits

hidden state one row per token LM-head matmul, tile by tile, on chip tile 1 tile 2 tile 3 tile 4 + Gumbel noise · one maximizer kept per tile small reduction argmax over tile winners token ✓ the |V|-sized logits tensor never reaches HBM
  • →Exact, not approximate: tilewise Gumbel-max keeps one maximizer per row and per vocabulary tile; a small reduction over tiles finishes the categorical sample.
  • →Fused into the LM-head matmul: eliminates the logits memory overhead and cuts decoding time by up to 19%.

Ruiz*, Qin*, Zhang†, Shen, Zhong, Wang† · 02/2026

A family of methods

One principle, many forms

GRAPE Group-representational position encoding — symmetry-aware positions. ICLR 2026
Higher-order Linear Attention Linear-time state with higher-order token interactions. Preprint · HLA
Fast-Weight Attention Fast weights as memory for continual learning. Preprint · Falcon
Deep Delta Learning The delta rule, made deep — error-driven state updates. Preprint · DDL
DeepLoop Depth scaling for weight-tied looped Transformers: the aligned-visit exponent p = 1/2. arXiv:2607.13491

Throughline: structured state → linear cost, lasting memory.

Pillar I · GRAPE · ICLR 2026

Positions as group representations

Positions act via one-parameter subgroups, giving the General Relative Law: $G(t-s)=G(s)^{-1}G(t)$, so scores depend only on the offset $t-s$.
GRAPE-M · Multiplicative · SO(d) nω G(n) = exp(nωL) norm-preserving rotations,rank-2 skew generators L recovers RoPE exactly; extends to learned commuting subspaces and non-commuting mixtures GRAPE-A · Additive · GL(d+k) InA0I unipotent lift with A² = 0:position enters as anadditive logit bias recovers ALiBi and FoX exactly, with the exact relative law preserved

One frame, exact relative position laws, streaming cacheability · see also OpenAI's GPT-4 Turbo & GPT-5, xAI's Grok 4, Thinking Machines' Inkling, and Jane Street's blog

Pillar I · DeepLoop · the framework

Loop the blocks, scale the depth

K = 2 physical blocks, stored once Block 1 · φ₁ attn → ffn Block 2 · φ₂ attn → ffn loop × R rounds x₀ φ₁ φ₂ φ₁ φ₂ φ₁ φ₂ xN round r = 1round r = 2round r = 3 R = 3 rounds × K = 2 blocks ⇒ unrolled depth N = KR = 6, the same weights revisited in every round
  • →A looped Transformer stores K physical blocks once and revisits them for R rounds: effective depth N = KR with no new parameters.
  • →Because every visit reuses the same weights, residual updates align across rounds; that alignment is exactly what the α, β rule accounts for.

Pillar I · DeepLoop · arXiv:2607.13491

Loop-aware depth scaling that shows up in loss

$\mathbf{x}_{i+1}=\mathrm{Norm}(\alpha\,\mathbf{x}_i+f(\mathbf{x}_i)),\qquad \alpha=(2N)^{1/2},\quad \beta=(8N)^{-1/2},\qquad p=1/2$ is the aligned-visit threshold.
FineWeb-Edu val loss @ 100K steps · 124M backbone 2.862.822.782.74 R = 1R = 3R = 5R = 7 2.7700 2.7514 base (α = β = 1)DeepLoop

Gap widens with loop count; recovers DeepNorm's 1/4 when visits decorrelate; also lifts HRM voted accuracy on ARC-AGI.

Pillar I · fixed-state models

A state that stays fixed, and keeps learning

attention KV cache: grows with every token fixed state (HLA · Falcon): constant memory tokens processed
  • →HLA: higher-order token interactions at linear time; the state does more per byte.
  • →Falcon: fast weights as episodic memory, so learning continues after pretraining ends.

Pillar I · Higher-order Linear Attention · Zhang, Qin, Gu

HLA: higher-order attention, linear time

softmax attention: O(n²) pairwise scores, the matrix grows with the context replace with token stream prefix statistics Sₜ constant size · exact identities second order and beyond HLA: per-token linear time, no n × n matrix ever materialized
  • →Higher-order interactions from compact prefix sufficient statistics: closed-form streaming identities and a strictly causal masked variant with two extra summaries.
  • →Chunk-parallel training via associative scans reproduces the serial recurrence exactly; extends to third order and higher.

Preprint · 101 GitHub stars · attention-like data-dependent mixing at recurrent-model cost

Pillar I · HLA · the second-order mechanism

Two summaries are enough

$o_t=q_t^{\top}S^{K}_{t}\,C^{QV}_{t},\qquad S^{K}_{t}=\textstyle\sum_{i\le t}k_ik_i^{\top}\in\mathbb{R}^{d\times d},\qquad C^{QV}_{t}=\textstyle\sum_{i\le t}q_iv_i^{\top}\in\mathbb{R}^{d\times d_v}$
token t accumulate k k⊤O(d²) per token accumulate q v⊤O(d dₕ) per token bilinear read constant-size state chunk 1 chunk 2 chunk 3 chunk 4 ⊕ ⊕ ⊕ associative scan (semidirect product) chunk-parallel training = the serial loop, exactly
  • →Second-order tensor attention $Q(K^{\top}K)Q^{\top}$ touches the context only through second moments, so two streaming summaries are sufficient statistics.
  • →An associative semidirect-product operator turns the recurrence into a parallel scan; unnormalized and normalized variants share the same state.

Strict causality via a masked variant with two extra summaries · extends to third order and higher

Pillar I · Falcon · Preprint 03/2026

Falcon: fast weights as continual learning

incoming tokens write bounded fast-weight state Wₜ = update(Wₜ₋₁, xₜ) an online continual-learning rule read attention-like mixing constant memory, any length every token updates the state: at inference, the model is continually learning from its own context
  • →One lens: attention, recurrent fast-weight memories, and selective SSMs are write rules over a bounded memory.
  • →Compressing a growing context into a bounded state makes the write rule an online continual-learning rule, connecting architecture design to learning theory.

The continual-learning perspective behind the fixed-state family

Pillar I · Falcon · six write rules

The Falcon family

$S_t=(1-\eta_t\lambda_t)\,S_{t-1}+\eta_t\,x_t r_t^{\top},\qquad (x_t,y_t)=(\phi(k_{t-1}),\,v_t)$  under a read-after-write convention.
regressioninner-product NLMS ridge stepsenergy-normalized Falcon-1one scalar step ηₜ Falcon-2per-column steps Falcon-3sliding-window mini-batch Falcon-1Ascalar normalized write Falcon-2Aper-column writes Falcon-3Asliding-window write
  • →Read-after-write alignment: the prefix-aligned pair $(\phi(k_{t-1}), v_t)$ trains the fast memory on what was actually available at prediction time.
  • →Consistent semantics: $\beta_t$ plasticity gain, $\lambda_t$ ridge shrinkage, $\eta_t$ induced step size; smoothness-matched or energy-normalized.

Aligned, normalized updates keep language-model quality and improve arithmetic length extrapolation

Pillar I · Falcon · one gradient step per token

The write is a gradient step

$\nabla_S\ell_t(S)=x_t(S^{\top}x_t-y_t)^{\top}+\lambda_t S\;\;\Rightarrow\;\;S_t=(1-\eta_t\lambda_t)S_{t-1}+\eta_t\,x_t r_t^{\top},\qquad \eta_t=\dfrac{\beta_t}{\|x_t\|_2^2+\lambda_t+\varepsilon}$
predict with St−1 read at xₜ = φ(kt−1) residual rₜ error vs the value vₜ normalized write one NLMS ridge step → Sₜ next token, same cycle: the fast memory is trained online
  • →The write is literally one gradient-descent step on an instantaneous ridge objective; the residual $r_t=y_t-S_{t-1}^{\top}x_t$ is measured against the previous state's prediction, unlike standard DeltaNet.
  • →The loss is $L_t$-smooth with $L_t=\|x_t\|_2^2+\lambda_t$, so $\eta_t$ is smoothness-matched; with $\lambda_t=\varepsilon=0$ it reduces exactly to classical NLMS.

Pillar I · Falcon · experiments

Aligned writes extrapolate

Variable-digit addition · trained on 1–32 digits, evaluated on 33–48 (teacher-forced mean accuracy) Transformer (RoPE) Falcon-1A.3 Falcon-3A.3 65.8 85.9 87.2 0255075
  • →Language modeling holds: 124M–130M models on FineWeb-Edu at a matched ~50B-token budget stay competitive with Transformer, Mamba-2, DeltaNet, and Gated DeltaNet baselines.
  • →Length extrapolation improves: the shifted, normalized writes win where storage and carry propagation dominate, 87.2 vs 65.8 mean accuracy out of distribution.

A controlled diagnostic that isolates the memory-writing behavior of the recurrent state

Pillar I · Deep Delta Learning · arXiv:2601.00417

DDL: gating the residual spectrum

$\mathbf{X}_{l+1}=(\mathbf{I}-\beta_l\,\mathbf{k}_l\mathbf{k}_l^{\top})\,\mathbf{X}_l+\beta_l\,\mathbf{k}_l\mathbf{v}_l^{\top},\qquad \lambda(\mathbf{A})=\{1,\ldots,1,\;1-\beta\},\quad \beta\in[0,2]$
β → 0 · identityβ → 1 · projectionβ → 2 · reflection skip connection, spectrum {1}forgetting: erases along kHouseholder, eigenvalue −1 k⟂k⟂k⟂
  • →A rank-1 Householder-style update on the identity shortcut: erasing along k and writing v are coupled, with the gate as a synchronous dynamic step size.
  • →The network controls the spectrum of its layer transition operator, modeling non-monotonic dynamics while keeping gated-residual training stability.

arXiv:2601.00417 · 355 GitHub stars · Zhang, Liu, Wang, Gu

Pillar II

Learning

Derive the objective, then scale agentic RL to long horizons.

Pillar II · a unifying view

REINFORCE is all you need

$\nabla_\theta J=\mathbb{E}_{x\sim\mu}\big[\textstyle\sum_t A_t\,\nabla_\theta\log\pi_\theta(x_t\mid x_{<t})\big]$: pretraining is $\mu=$ corpus, $A_t\equiv 1$; RL is $\mu=\pi_\theta$, $A_t$ from reward.
one estimator ∇ log π · advantage Pretraining · off-policy REINFORCE behavior = the corpus · advantage A ≡ 1 Agentic RL · on-policy REINFORCE behavior = π rollouts · A from verifiable reward corpusreplayed rolloutsstale async rollouts, IS-correctedfresh rollouts high staleness on-policy the spectrum: off-policy staleness interpolates pretraining and RL scaling
  • →Pretraining is off-policy REINFORCE with A ≡ 1: the corpus is the behavior policy, and every next token is a positively rewarded action.
  • →Scaling REINFORCE across the spectrum, off-policy (corpus, replayed rollouts, IS-corrected) to on-policy (verifier-rewarded rollouts), unifies the training stack.

One objective family from pretraining to agentic RL

The problem

Policy gradients for LLMs are a zoo

  • 01RLHF and RLVR now drive frontier reasoning.
  • 02But GRPO / PPO are a patchwork of clipping and KL hacks — unstable, hard to reason about.

Start from the KL-regularized objective — and be exact.

ICLR 2026

KL-regularized policy gradient, done right

A unified design space for KL-regularized policy-gradient algorithms — with the correct gradient and stop-gradient treatment.

Impact Adopted in Thinking Machines' Tinker and DeepSeek-V3.2.
  • →More stable training; recovers and repairs popular methods as special cases.

Pillar II · RPG · ICLR 2026

One design space, every KL policy gradient

Forward KLReverse KL NormalizedUnnormalized k1 · k2 · k3 estimatorspolicy-gradient or fully differentiable k1 · k2 · k3 estimatorspolicy-gradient or fully differentiable k1 · k2 · k3 estimatorscorrect gradient, derived once k1 · k2 · k3 estimatorscorrect stop-gradient placement
  • →Every cell gets the correct gradient and stop-gradient treatment, derived from the objective.
  • →Recovers and repairs GRPO-style methods as special cases; measurably more stable training.

Adopted in Thinking Machines' Tinker and DeepSeek-V3.2

Alignment & RL

Beyond rewards and clips

Self-Distilled Policy Gradient The policy distills into its own target — stable, low-variance updates.
$\nabla_\theta\,\mathbb{E}_{a\sim p_t}\!\big[-\log p_t(a)\,\mathrm{SG}(\log \bar q_t/\bar p_t+\bar D_t)\big]$
SDPG
General Preference Model Beyond Bradley–Terry: a skew-symmetric operator captures cyclic, intransitive human preference. ICML 2025

Reward modeling and RL, on one principled footing.

Agentic RL at scale

The framework and the science

Molt An agentic-first, PyTorch-native RL framework: Ray, vLLM, and NVIDIA AutoModel. About 9.2K lines of RL code that scale to 1T-class MoE. NVIDIA · arXiv:2607.21653
Prefill-Decode Kernel Mismatch Same checkpoint, different policy: rollout and learner kernels define different effective policies. Four remedies, from importance ratios to exact rejection sampling. Tech Report · 08/2026

Scaling agentic RL takes both the infrastructure and the science.

Pillar II · Molt · NVIDIA · arXiv:2607.21653

Molt: three boxes, one async loop

Agent Env · ChatAgent · reward in Python vLLM rollout async engines · partial rollout Single-actor trainer AutoModel + FSDP2 · TP / EP / CP Ray queue weight sync, fully async
  • →About 9.2K lines of RL code; MoE-native to 1T-class, think DeepSeek-V3 at EP 256.
  • →reinforce · rloo · grpo · gae · on-policy distillation; router replay (R3) and IS correction for async rollout.

Pillar II · Pretraining RL Science · 08/2026

Same checkpoint, different policy

Rollout Learner parallel prefill recurrent decodeone token at a time actorlog-probabilities teacher-forcedparallel scoring same θ, different effective policy
  • →Chunkwise, scan, and recurrent kernels are algebraically equal but different finite-precision operators; DeltaNet's corrective write feeds the error into future updates.
  • →Four remedies: actor-emitted importance ratios, one canonical execution rule, higher-precision state, and exact modified rejection sampling over a fast proposal.

Pillar III

AutoResearch

Agents that automate research itself.

Reasoning

Reasoning as structure, not length

Structured inference Cumulative Reasoning & Diagram of Thought — reasoning as a DAG, not a linear chain. Proposer–verifier loops that build and check intermediate results.
Control Meta Prompting — structured prompting as a control layer for LLMs. Composable, reusable reasoning scaffolds for AI systems.

Better structure beats longer chains.

Pillar III · structured reasoning

A chain forgets; a DAG accumulates

chain of thought: one path, no reuse, one bad step poisons the rest ✓ verified ✓ verified cumulative reasoning: propose, verify, accumulate
  • →A proposer suggests steps; a verifier checks them; only verified results join the context.
  • →Cumulative Reasoning (TMLR, 704 citations) and Diagram of Thought formalize the DAG view.

Agentic systems

Agents that act on the world

Flagship MathCode — a terminal coding agent that formalizes natural-language math into Lean 4 and proves it.
  • →Lanser-CLI — RL from compiler & language-server feedback (cf. Claude Code v2.0.74).
  • →Web World Models — controllable, open-ended worlds for language agents: rules in web code, context from LLMs.

Agents that browse, compute, and prove.

Pillar III · MathCode

From natural language to verified proof

natural-languagemath statement autoformalizeLean 4 statement proveagent loop verified proof ✓ checked by the kernel compiler feedback
  • →A terminal coding agent where the compiler is the referee; 584 GitHub stars, the most-starred project.
  • →CriticLean (ACL 2026) adds critic-guided RL to the formalization step itself.

Pillar III · environments

Rewards with executable ground truth

Agent browse · compute · prove Lean 4 & compilers: proofs that check Language servers (Lanser-CLI): typed, analyzable feedback Web World Models: rules in code, context from LLMs act / verify

An environment whose reward can be gamed teaches the model to game it; executable ground truth removes the exploit.

Pillar III · closing the loop

Model and harness coevolve

Model policy · memory · skills Harness perceive · act · evaluate a stronger model demands a better harness a better harness exposes what to learn next ↓  recursive self-improvement
  • →We build our own harness; token-exact traces survive even opaque context compaction.
  • →This loop is what turns research agents into AutoResearch.

What's next

Toward autonomous superintelligence

  • 01AutoResearch: automate AI research itself, so improvements compound instead of being hand-built.
  • 02Million-token memory: architectures for long context and continual learning, fast-weight state at scale.
  • 03Agentic RL scaling: long-horizon tasks with verifiable feedback from compilers, language servers, and world models.
  • 04Model-harness coevolution: build the harness, close the loop, approach recursive self-improvement.

Principled methods, built to ship.

Thank you

Thank you.

With gratitude to my collaborators and advisors.

yifzhang@princeton.edu yfz.ai Google Scholar GitHub
← Back
1 / 42

Slide navigatorClick a slide to jump · Esc to close

Keyboard & controls

→Space↓
Next build / slide
←↑
Previous
HomeEnd
First / last slide
F
Toggle fullscreen
O
Slide overview
S
Fast mode — instant page flips, no builds
?H
This help
Esc
Close overlay

Press any key to close