--- abstract: | Transformers process tokens in parallel but are temporally shallow: at position $t$, each layer attends to key--value pairs computed based on the previous layer, yielding a depth capped by the number of layers. Recurrent models offer unbounded temporal depth but suffer from optimization instability and historically underutilize modern accelerators. We introduce the *Recurrent Transformer*, a simple architectural change where *each layer* attends to key--value pairs computed off its own activations, yielding layerwise recurrent memory while preserving standard autoregressive decoding cost. We show that the architecture can emulate both (i) a conventional Transformer and (ii) token-to-token recurrent updates under mild assumptions, while avoiding optimization instability. Naively, prefill/training appears bandwidth-bound with effective arithmetic intensity near $1$ because keys and values are revealed sequentially; we give an exact tiling-based algorithm that preserves the mathematical computation while reducing HBM traffic from $\Theta(N^2)$ to $\Theta(N\log N)$, increasing effective arithmetic intensity to $\Theta(N/\log N)$ for sequence length $N$. On 150M and 300M parameter C4 pretraining, `\methodname `{=latex}s improve cross-entropy over a parameter-matched Transformer baseline and achieve the improvement with fewer layers (fixed parameters), suggesting that recurrence can trade depth for width, thus reducing KV cache memory footprint and inference latency. Code is available at bibliography: - references.bib title: | The Recurrent Transformer:\ Greater Effective Depth and Efficient Decoding --- ```{=latex} \newcommand{\costin}[1]{{\color{magenta} Costin:#1}} ``` ```{=latex} \newcommand{\depen}[1]{{\color{blue} Depen:#1}} ``` ```{=latex} \newcommand{\alex}[1]{{\color{orange} Alex:#1}} ``` ```{=latex} \newcommand{\figleft}{{\em (Left)}} ``` ```{=latex} \newcommand{\figcenter}{{\em (Center)}} ``` ```{=latex} \newcommand{\figright}{{\em (Right)}} ``` ```{=latex} \newcommand{\figtop}{{\em (Top)}} ``` ```{=latex} \newcommand{\figbottom}{{\em (Bottom)}} ``` ```{=latex} \newcommand{\captiona}{{\em (a)}} ``` ```{=latex} \newcommand{\captionb}{{\em (b)}} ``` ```{=latex} \newcommand{\captionc}{{\em (c)}} ``` ```{=latex} \newcommand{\captiond}{{\em (d)}} ``` ```{=latex} \newcommand{\newterm}[1]{{\bf #1}} ``` ```{=latex} \def\figref#1{figure~\ref{#1}} ``` ```{=latex} \def\Figref#1{Figure~\ref{#1}} ``` ```{=latex} \def\twofigref#1#2{figures \ref{#1} and \ref{#2}} ``` ```{=latex} \def\quadfigref#1#2#3#4{figures \ref{#1}, \ref{#2}, \ref{#3} and \ref{#4}} ``` ```{=latex} \def\secref#1{section~\ref{#1}} ``` ```{=latex} \def\Secref#1{Section~\ref{#1}} ``` ```{=latex} \def\twosecrefs#1#2{sections \ref{#1} and \ref{#2}} ``` ```{=latex} \def\secrefs#1#2#3{sections \ref{#1}, \ref{#2} and \ref{#3}} ``` ```{=latex} \def\eqref#1{equation~\ref{#1}} ``` ```{=latex} \def\Eqref#1{Equation~\ref{#1}} ``` ```{=latex} \def\plaineqref#1{\ref{#1}} ``` ```{=latex} \def\chapref#1{chapter~\ref{#1}} ``` ```{=latex} \def\Chapref#1{Chapter~\ref{#1}} ``` ```{=latex} \def\rangechapref#1#2{chapters\ref{#1}--\ref{#2}} ``` ```{=latex} \def\algref#1{algorithm~\ref{#1}} ``` ```{=latex} \def\Algref#1{Algorithm~\ref{#1}} ``` ```{=latex} \def\twoalgref#1#2{algorithms \ref{#1} and \ref{#2}} ``` ```{=latex} \def\Twoalgref#1#2{Algorithms \ref{#1} and \ref{#2}} ``` ```{=latex} \def\partref#1{part~\ref{#1}} ``` ```{=latex} \def\Partref#1{Part~\ref{#1}} ``` ```{=latex} \def\twopartref#1#2{parts \ref{#1} and \ref{#2}} ``` ```{=latex} \def\ceil#1{\lceil #1 \rceil} ``` ```{=latex} \def\floor#1{\lfloor #1 \rfloor} ``` ```{=latex} \def\1{\bm{1}} ``` ```{=latex} \newcommand{\train}{\mathcal{D}} ``` ```{=latex} \newcommand{\valid}{\mathcal{D_{\mathrm{valid}}}} ``` ```{=latex} \newcommand{\test}{\mathcal{D_{\mathrm{test}}}} ``` ```{=latex} \def\eps{{\epsilon}} ``` ```{=latex} \def\reta{{\textnormal{$\eta$}}} ``` ```{=latex} \def\ra{{\textnormal{a}}} ``` ```{=latex} \def\rb{{\textnormal{b}}} ``` ```{=latex} \def\rc{{\textnormal{c}}} ``` ```{=latex} \def\rd{{\textnormal{d}}} ``` ```{=latex} \def\re{{\textnormal{e}}} ``` ```{=latex} \def\rf{{\textnormal{f}}} ``` ```{=latex} \def\rg{{\textnormal{g}}} ``` ```{=latex} \def\rh{{\textnormal{h}}} ``` ```{=latex} \def\ri{{\textnormal{i}}} ``` ```{=latex} \def\rj{{\textnormal{j}}} ``` ```{=latex} \def\rk{{\textnormal{k}}} ``` ```{=latex} \def\rl{{\textnormal{l}}} ``` ```{=latex} \def\rn{{\textnormal{n}}} ``` ```{=latex} \def\ro{{\textnormal{o}}} ``` ```{=latex} \def\rp{{\textnormal{p}}} ``` ```{=latex} \def\rq{{\textnormal{q}}} ``` ```{=latex} \def\rr{{\textnormal{r}}} ``` ```{=latex} \def\rs{{\textnormal{s}}} ``` ```{=latex} \def\rt{{\textnormal{t}}} ``` ```{=latex} \def\ru{{\textnormal{u}}} ``` ```{=latex} \def\rv{{\textnormal{v}}} ``` ```{=latex} \def\rw{{\textnormal{w}}} ``` ```{=latex} \def\rx{{\textnormal{x}}} ``` ```{=latex} \def\ry{{\textnormal{y}}} ``` ```{=latex} \def\rz{{\textnormal{z}}} ``` ```{=latex} \def\rvepsilon{{\mathbf{\epsilon}}} ``` ```{=latex} \def\rvtheta{{\mathbf{\theta}}} ``` ```{=latex} \def\rva{{\mathbf{a}}} ``` ```{=latex} \def\rvb{{\mathbf{b}}} ``` ```{=latex} \def\rvc{{\mathbf{c}}} ``` ```{=latex} \def\rvd{{\mathbf{d}}} ``` ```{=latex} \def\rve{{\mathbf{e}}} ``` ```{=latex} \def\rvf{{\mathbf{f}}} ``` ```{=latex} \def\rvg{{\mathbf{g}}} ``` ```{=latex} \def\rvh{{\mathbf{h}}} ``` ```{=latex} \def\rvu{{\mathbf{i}}} ``` ```{=latex} \def\rvj{{\mathbf{j}}} ``` ```{=latex} \def\rvk{{\mathbf{k}}} ``` ```{=latex} \def\rvl{{\mathbf{l}}} ``` ```{=latex} \def\rvm{{\mathbf{m}}} ``` ```{=latex} \def\rvn{{\mathbf{n}}} ``` ```{=latex} \def\rvo{{\mathbf{o}}} ``` ```{=latex} \def\rvp{{\mathbf{p}}} ``` ```{=latex} \def\rvq{{\mathbf{q}}} ``` ```{=latex} \def\rvr{{\mathbf{r}}} ``` ```{=latex} \def\rvs{{\mathbf{s}}} ``` ```{=latex} \def\rvt{{\mathbf{t}}} ``` ```{=latex} \def\rvu{{\mathbf{u}}} ``` ```{=latex} \def\rvv{{\mathbf{v}}} ``` ```{=latex} \def\rvw{{\mathbf{w}}} ``` ```{=latex} \def\rvx{{\mathbf{x}}} ``` ```{=latex} \def\rvy{{\mathbf{y}}} ``` ```{=latex} \def\rvz{{\mathbf{z}}} ``` ```{=latex} \def\erva{{\textnormal{a}}} ``` ```{=latex} \def\ervb{{\textnormal{b}}} ``` ```{=latex} \def\ervc{{\textnormal{c}}} ``` ```{=latex} \def\ervd{{\textnormal{d}}} ``` ```{=latex} \def\erve{{\textnormal{e}}} ``` ```{=latex} \def\ervf{{\textnormal{f}}} ``` ```{=latex} \def\ervg{{\textnormal{g}}} ``` ```{=latex} \def\ervh{{\textnormal{h}}} ``` ```{=latex} \def\ervi{{\textnormal{i}}} ``` ```{=latex} \def\ervj{{\textnormal{j}}} ``` ```{=latex} \def\ervk{{\textnormal{k}}} ``` ```{=latex} \def\ervl{{\textnormal{l}}} ``` ```{=latex} \def\ervm{{\textnormal{m}}} ``` ```{=latex} \def\ervn{{\textnormal{n}}} ``` ```{=latex} \def\ervo{{\textnormal{o}}} ``` ```{=latex} \def\ervp{{\textnormal{p}}} ``` ```{=latex} \def\ervq{{\textnormal{q}}} ``` ```{=latex} \def\ervr{{\textnormal{r}}} ``` ```{=latex} \def\ervs{{\textnormal{s}}} ``` ```{=latex} \def\ervt{{\textnormal{t}}} ``` ```{=latex} \def\ervu{{\textnormal{u}}} ``` ```{=latex} \def\ervv{{\textnormal{v}}} ``` ```{=latex} \def\ervw{{\textnormal{w}}} ``` ```{=latex} \def\ervx{{\textnormal{x}}} ``` ```{=latex} \def\ervy{{\textnormal{y}}} ``` ```{=latex} \def\ervz{{\textnormal{z}}} ``` ```{=latex} \def\rmA{{\mathbf{A}}} ``` ```{=latex} \def\rmB{{\mathbf{B}}} ``` ```{=latex} \def\rmC{{\mathbf{C}}} ``` ```{=latex} \def\rmD{{\mathbf{D}}} ``` ```{=latex} \def\rmE{{\mathbf{E}}} ``` ```{=latex} \def\rmF{{\mathbf{F}}} ``` ```{=latex} \def\rmG{{\mathbf{G}}} ``` ```{=latex} \def\rmH{{\mathbf{H}}} ``` ```{=latex} \def\rmI{{\mathbf{I}}} ``` ```{=latex} \def\rmJ{{\mathbf{J}}} ``` ```{=latex} \def\rmK{{\mathbf{K}}} ``` ```{=latex} \def\rmL{{\mathbf{L}}} ``` ```{=latex} \def\rmM{{\mathbf{M}}} ``` ```{=latex} \def\rmN{{\mathbf{N}}} ``` ```{=latex} \def\rmO{{\mathbf{O}}} ``` ```{=latex} \def\rmP{{\mathbf{P}}} ``` ```{=latex} \def\rmQ{{\mathbf{Q}}} ``` ```{=latex} \def\rmR{{\mathbf{R}}} ``` ```{=latex} \def\rmS{{\mathbf{S}}} ``` ```{=latex} \def\rmT{{\mathbf{T}}} ``` ```{=latex} \def\rmU{{\mathbf{U}}} ``` ```{=latex} \def\rmV{{\mathbf{V}}} ``` ```{=latex} \def\rmW{{\mathbf{W}}} ``` ```{=latex} \def\rmX{{\mathbf{X}}} ``` ```{=latex} \def\rmY{{\mathbf{Y}}} ``` ```{=latex} \def\rmZ{{\mathbf{Z}}} ``` ```{=latex} \def\ermA{{\textnormal{A}}} ``` ```{=latex} \def\ermB{{\textnormal{B}}} ``` ```{=latex} \def\ermC{{\textnormal{C}}} ``` ```{=latex} \def\ermD{{\textnormal{D}}} ``` ```{=latex} \def\ermE{{\textnormal{E}}} ``` ```{=latex} \def\ermF{{\textnormal{F}}} ``` ```{=latex} \def\ermG{{\textnormal{G}}} ``` ```{=latex} \def\ermH{{\textnormal{H}}} ``` ```{=latex} \def\ermI{{\textnormal{I}}} ``` ```{=latex} \def\ermJ{{\textnormal{J}}} ``` ```{=latex} \def\ermK{{\textnormal{K}}} ``` ```{=latex} \def\ermL{{\textnormal{L}}} ``` ```{=latex} \def\ermM{{\textnormal{M}}} ``` ```{=latex} \def\ermN{{\textnormal{N}}} ``` ```{=latex} \def\ermO{{\textnormal{O}}} ``` ```{=latex} \def\ermP{{\textnormal{P}}} ``` ```{=latex} \def\ermQ{{\textnormal{Q}}} ``` ```{=latex} \def\ermR{{\textnormal{R}}} ``` ```{=latex} \def\ermS{{\textnormal{S}}} ``` ```{=latex} \def\ermT{{\textnormal{T}}} ``` ```{=latex} \def\ermU{{\textnormal{U}}} ``` ```{=latex} \def\ermV{{\textnormal{V}}} ``` ```{=latex} \def\ermW{{\textnormal{W}}} ``` ```{=latex} \def\ermX{{\textnormal{X}}} ``` ```{=latex} \def\ermY{{\textnormal{Y}}} ``` ```{=latex} \def\ermZ{{\textnormal{Z}}} ``` ```{=latex} \def\vzero{{\bm{0}}} ``` ```{=latex} \def\vone{{\bm{1}}} ``` ```{=latex} \def\vmu{{\bm{\mu}}} ``` ```{=latex} \def\vtheta{{\bm{\theta}}} ``` ```{=latex} \def\va{{\bm{a}}} ``` ```{=latex} \def\vb{{\bm{b}}} ``` ```{=latex} \def\vc{{\bm{c}}} ``` ```{=latex} \def\vd{{\bm{d}}} ``` ```{=latex} \def\ve{{\bm{e}}} ``` ```{=latex} \def\vf{{\bm{f}}} ``` ```{=latex} \def\vg{{\bm{g}}} ``` ```{=latex} \def\vh{{\bm{h}}} ``` ```{=latex} \def\vi{{\bm{i}}} ``` ```{=latex} \def\vj{{\bm{j}}} ``` ```{=latex} \def\vk{{\bm{k}}} ``` ```{=latex} \def\vl{{\bm{l}}} ``` ```{=latex} \def\vm{{\bm{m}}} ``` ```{=latex} \def\vn{{\bm{n}}} ``` ```{=latex} \def\vo{{\bm{o}}} ``` ```{=latex} \def\vp{{\bm{p}}} ``` ```{=latex} \def\vq{{\bm{q}}} ``` ```{=latex} \def\vr{{\bm{r}}} ``` ```{=latex} \def\vs{{\bm{s}}} ``` ```{=latex} \def\vt{{\bm{t}}} ``` ```{=latex} \def\vu{{\bm{u}}} ``` ```{=latex} \def\vv{{\bm{v}}} ``` ```{=latex} \def\vw{{\bm{w}}} ``` ```{=latex} \def\vx{{\bm{x}}} ``` ```{=latex} \def\vy{{\bm{y}}} ``` ```{=latex} \def\vz{{\bm{z}}} ``` ```{=latex} \def\evalpha{{\alpha}} ``` ```{=latex} \def\evbeta{{\beta}} ``` ```{=latex} \def\evepsilon{{\epsilon}} ``` ```{=latex} \def\evlambda{{\lambda}} ``` ```{=latex} \def\evomega{{\omega}} ``` ```{=latex} \def\evmu{{\mu}} ``` ```{=latex} \def\evpsi{{\psi}} ``` ```{=latex} \def\evsigma{{\sigma}} ``` ```{=latex} \def\evtheta{{\theta}} ``` ```{=latex} \def\eva{{a}} ``` ```{=latex} \def\evb{{b}} ``` ```{=latex} \def\evc{{c}} ``` ```{=latex} \def\evd{{d}} ``` ```{=latex} \def\eve{{e}} ``` ```{=latex} \def\evf{{f}} ``` ```{=latex} \def\evg{{g}} ``` ```{=latex} \def\evh{{h}} ``` ```{=latex} \def\evi{{i}} ``` ```{=latex} \def\evj{{j}} ``` ```{=latex} \def\evk{{k}} ``` ```{=latex} \def\evl{{l}} ``` ```{=latex} \def\evm{{m}} ``` ```{=latex} \def\evn{{n}} ``` ```{=latex} \def\evo{{o}} ``` ```{=latex} \def\evp{{p}} ``` ```{=latex} \def\evq{{q}} ``` ```{=latex} \def\evr{{r}} ``` ```{=latex} \def\evs{{s}} ``` ```{=latex} \def\evt{{t}} ``` ```{=latex} \def\evu{{u}} ``` ```{=latex} \def\evv{{v}} ``` ```{=latex} \def\evw{{w}} ``` ```{=latex} \def\evx{{x}} ``` ```{=latex} \def\evy{{y}} ``` ```{=latex} \def\evz{{z}} ``` ```{=latex} \def\mA{{\bm{A}}} ``` ```{=latex} \def\mB{{\bm{B}}} ``` ```{=latex} \def\mC{{\bm{C}}} ``` ```{=latex} \def\mD{{\bm{D}}} ``` ```{=latex} \def\mE{{\bm{E}}} ``` ```{=latex} \def\mF{{\bm{F}}} ``` ```{=latex} \def\mG{{\bm{G}}} ``` ```{=latex} \def\mH{{\bm{H}}} ``` ```{=latex} \def\mI{{\bm{I}}} ``` ```{=latex} \def\mJ{{\bm{J}}} ``` ```{=latex} \def\mK{{\bm{K}}} ``` ```{=latex} \def\mL{{\bm{L}}} ``` ```{=latex} \def\mM{{\bm{M}}} ``` ```{=latex} \def\mN{{\bm{N}}} ``` ```{=latex} \def\mO{{\bm{O}}} ``` ```{=latex} \def\mP{{\bm{P}}} ``` ```{=latex} \def\mQ{{\bm{Q}}} ``` ```{=latex} \def\mR{{\bm{R}}} ``` ```{=latex} \def\mS{{\bm{S}}} ``` ```{=latex} \def\mT{{\bm{T}}} ``` ```{=latex} \def\mU{{\bm{U}}} ``` ```{=latex} \def\mV{{\bm{V}}} ``` ```{=latex} \def\mW{{\bm{W}}} ``` ```{=latex} \def\mX{{\bm{X}}} ``` ```{=latex} \def\mY{{\bm{Y}}} ``` ```{=latex} \def\mZ{{\bm{Z}}} ``` ```{=latex} \def\mBeta{{\bm{\beta}}} ``` ```{=latex} \def\mPhi{{\bm{\Phi}}} ``` ```{=latex} \def\mLambda{{\bm{\Lambda}}} ``` ```{=latex} \def\mSigma{{\bm{\Sigma}}} ``` ```{=latex} \newcommand{\tens}[1]{\bm{\mathsfit{#1}}} ``` ```{=latex} \def\tA{{\tens{A}}} ``` ```{=latex} \def\tB{{\tens{B}}} ``` ```{=latex} \def\tC{{\tens{C}}} ``` ```{=latex} \def\tD{{\tens{D}}} ``` ```{=latex} \def\tE{{\tens{E}}} ``` ```{=latex} \def\tF{{\tens{F}}} ``` ```{=latex} \def\tG{{\tens{G}}} ``` ```{=latex} \def\tH{{\tens{H}}} ``` ```{=latex} \def\tI{{\tens{I}}} ``` ```{=latex} \def\tJ{{\tens{J}}} ``` ```{=latex} \def\tK{{\tens{K}}} ``` ```{=latex} \def\tL{{\tens{L}}} ``` ```{=latex} \def\tM{{\tens{M}}} ``` ```{=latex} \def\tN{{\tens{N}}} ``` ```{=latex} \def\tO{{\tens{O}}} ``` ```{=latex} \def\tP{{\tens{P}}} ``` ```{=latex} \def\tQ{{\tens{Q}}} ``` ```{=latex} \def\tR{{\tens{R}}} ``` ```{=latex} \def\tS{{\tens{S}}} ``` ```{=latex} \def\tT{{\tens{T}}} ``` ```{=latex} \def\tU{{\tens{U}}} ``` ```{=latex} \def\tV{{\tens{V}}} ``` ```{=latex} \def\tW{{\tens{W}}} ``` ```{=latex} \def\tX{{\tens{X}}} ``` ```{=latex} \def\tY{{\tens{Y}}} ``` ```{=latex} \def\tZ{{\tens{Z}}} ``` ```{=latex} \def\gA{{\mathcal{A}}} ``` ```{=latex} \def\gB{{\mathcal{B}}} ``` ```{=latex} \def\gC{{\mathcal{C}}} ``` ```{=latex} \def\gD{{\mathcal{D}}} ``` ```{=latex} \def\gE{{\mathcal{E}}} ``` ```{=latex} \def\gF{{\mathcal{F}}} ``` ```{=latex} \def\gG{{\mathcal{G}}} ``` ```{=latex} \def\gH{{\mathcal{H}}} ``` ```{=latex} \def\gI{{\mathcal{I}}} ``` ```{=latex} \def\gJ{{\mathcal{J}}} ``` ```{=latex} \def\gK{{\mathcal{K}}} ``` ```{=latex} \def\gL{{\mathcal{L}}} ``` ```{=latex} \def\gM{{\mathcal{M}}} ``` ```{=latex} \def\gN{{\mathcal{N}}} ``` ```{=latex} \def\gO{{\mathcal{O}}} ``` ```{=latex} \def\gP{{\mathcal{P}}} ``` ```{=latex} \def\gQ{{\mathcal{Q}}} ``` ```{=latex} \def\gR{{\mathcal{R}}} ``` ```{=latex} \def\gS{{\mathcal{S}}} ``` ```{=latex} \def\gT{{\mathcal{T}}} ``` ```{=latex} \def\gU{{\mathcal{U}}} ``` ```{=latex} \def\gV{{\mathcal{V}}} ``` ```{=latex} \def\gW{{\mathcal{W}}} ``` ```{=latex} \def\gX{{\mathcal{X}}} ``` ```{=latex} \def\gY{{\mathcal{Y}}} ``` ```{=latex} \def\gZ{{\mathcal{Z}}} ``` ```{=latex} \def\sA{{\mathbb{A}}} ``` ```{=latex} \def\sB{{\mathbb{B}}} ``` ```{=latex} \def\sC{{\mathbb{C}}} ``` ```{=latex} \def\sD{{\mathbb{D}}} ``` ```{=latex} \def\sF{{\mathbb{F}}} ``` ```{=latex} \def\sG{{\mathbb{G}}} ``` ```{=latex} \def\sH{{\mathbb{H}}} ``` ```{=latex} \def\sI{{\mathbb{I}}} ``` ```{=latex} \def\sJ{{\mathbb{J}}} ``` ```{=latex} \def\sK{{\mathbb{K}}} ``` ```{=latex} \def\sL{{\mathbb{L}}} ``` ```{=latex} \def\sM{{\mathbb{M}}} ``` ```{=latex} \def\sN{{\mathbb{N}}} ``` ```{=latex} \def\sO{{\mathbb{O}}} ``` ```{=latex} \def\sP{{\mathbb{P}}} ``` ```{=latex} \def\sQ{{\mathbb{Q}}} ``` ```{=latex} \def\sR{{\mathbb{R}}} ``` ```{=latex} \def\sS{{\mathbb{S}}} ``` ```{=latex} \def\sT{{\mathbb{T}}} ``` ```{=latex} \def\sU{{\mathbb{U}}} ``` ```{=latex} \def\sV{{\mathbb{V}}} ``` ```{=latex} \def\sW{{\mathbb{W}}} ``` ```{=latex} \def\sX{{\mathbb{X}}} ``` ```{=latex} \def\sY{{\mathbb{Y}}} ``` ```{=latex} \def\sZ{{\mathbb{Z}}} ``` ```{=latex} \def\emLambda{{\Lambda}} ``` ```{=latex} \def\emA{{A}} ``` ```{=latex} \def\emB{{B}} ``` ```{=latex} \def\emC{{C}} ``` ```{=latex} \def\emD{{D}} ``` ```{=latex} \def\emE{{E}} ``` ```{=latex} \def\emF{{F}} ``` ```{=latex} \def\emG{{G}} ``` ```{=latex} \def\emH{{H}} ``` ```{=latex} \def\emI{{I}} ``` ```{=latex} \def\emJ{{J}} ``` ```{=latex} \def\emK{{K}} ``` ```{=latex} \def\emL{{L}} ``` ```{=latex} \def\emM{{M}} ``` ```{=latex} \def\emN{{N}} ``` ```{=latex} \def\emO{{O}} ``` ```{=latex} \def\emP{{P}} ``` ```{=latex} \def\emQ{{Q}} ``` ```{=latex} \def\emR{{R}} ``` ```{=latex} \def\emS{{S}} ``` ```{=latex} \def\emT{{T}} ``` ```{=latex} \def\emU{{U}} ``` ```{=latex} \def\emV{{V}} ``` ```{=latex} \def\emW{{W}} ``` ```{=latex} \def\emX{{X}} ``` ```{=latex} \def\emY{{Y}} ``` ```{=latex} \def\emZ{{Z}} ``` ```{=latex} \def\emSigma{{\Sigma}} ``` ```{=latex} \newcommand{\etens}[1]{\mathsfit{#1}} ``` ```{=latex} \def\etLambda{{\etens{\Lambda}}} ``` ```{=latex} \def\etA{{\etens{A}}} ``` ```{=latex} \def\etB{{\etens{B}}} ``` ```{=latex} \def\etC{{\etens{C}}} ``` ```{=latex} \def\etD{{\etens{D}}} ``` ```{=latex} \def\etE{{\etens{E}}} ``` ```{=latex} \def\etF{{\etens{F}}} ``` ```{=latex} \def\etG{{\etens{G}}} ``` ```{=latex} \def\etH{{\etens{H}}} ``` ```{=latex} \def\etI{{\etens{I}}} ``` ```{=latex} \def\etJ{{\etens{J}}} ``` ```{=latex} \def\etK{{\etens{K}}} ``` ```{=latex} \def\etL{{\etens{L}}} ``` ```{=latex} \def\etM{{\etens{M}}} ``` ```{=latex} \def\etN{{\etens{N}}} ``` ```{=latex} \def\etO{{\etens{O}}} ``` ```{=latex} \def\etP{{\etens{P}}} ``` ```{=latex} \def\etQ{{\etens{Q}}} ``` ```{=latex} \def\etR{{\etens{R}}} ``` ```{=latex} \def\etS{{\etens{S}}} ``` ```{=latex} \def\etT{{\etens{T}}} ``` ```{=latex} \def\etU{{\etens{U}}} ``` ```{=latex} \def\etV{{\etens{V}}} ``` ```{=latex} \def\etW{{\etens{W}}} ``` ```{=latex} \def\etX{{\etens{X}}} ``` ```{=latex} \def\etY{{\etens{Y}}} ``` ```{=latex} \def\etZ{{\etens{Z}}} ``` ```{=latex} \newcommand{\pdata}{p_{\rm{data}}} ``` ```{=latex} \newcommand{\ptrain}{\hat{p}_{\rm{data}}} ``` ```{=latex} \newcommand{\Ptrain}{\hat{P}_{\rm{data}}} ``` ```{=latex} \newcommand{\pmodel}{p_{\rm{model}}} ``` ```{=latex} \newcommand{\Pmodel}{P_{\rm{model}}} ``` ```{=latex} \newcommand{\ptildemodel}{\tilde{p}_{\rm{model}}} ``` ```{=latex} \newcommand{\pencode}{p_{\rm{encoder}}} ``` ```{=latex} \newcommand{\pdecode}{p_{\rm{decoder}}} ``` ```{=latex} \newcommand{\precons}{p_{\rm{reconstruct}}} ``` ```{=latex} \newcommand{\laplace}{\mathrm{Laplace}} ``` ```{=latex} \newcommand{\E}{\mathbb{E}} ``` ```{=latex} \newcommand{\Ls}{\mathcal{L}} ``` ```{=latex} \newcommand{\R}{\mathbb{R}} ``` ```{=latex} \newcommand{\N}{\mathbb{N}} ``` ```{=latex} \newcommand{\emp}{\tilde{p}} ``` ```{=latex} \newcommand{\lr}{\alpha} ``` ```{=latex} \newcommand{\reg}{\lambda} ``` ```{=latex} \newcommand{\rect}{\mathrm{rectifier}} ``` ```{=latex} \newcommand{\softmax}{\mathrm{softmax}} ``` ```{=latex} \newcommand{\sigmoid}{\sigma} ``` ```{=latex} \newcommand{\softplus}{\zeta} ``` ```{=latex} \newcommand{\KL}{D_{\mathrm{KL}}} ``` ```{=latex} \newcommand{\Var}{\mathrm{Var}} ``` ```{=latex} \newcommand{\standarderror}{\mathrm{SE}} ``` ```{=latex} \newcommand{\Cov}{\mathrm{Cov}} ``` ```{=latex} \newcommand{\normlzero}{L^0} ``` ```{=latex} \newcommand{\normlone}{L^1} ``` ```{=latex} \newcommand{\normltwo}{L^2} ``` ```{=latex} \newcommand{\normlp}{L^p} ``` ```{=latex} \newcommand{\normmax}{L^\infty} ``` ```{=latex} \newcommand{\parents}{Pa} ``` ```{=latex} \newcommand{\LN}{\mathrm{LN}} ``` ```{=latex} \newcommand{\RMS}{\mathrm{RMS}} ``` ```{=latex} \newcommand{\MLP}{\mathrm{MLP}} ``` ```{=latex} \newcommand{\Attn}{\mathrm{Attn}} ``` ```{=latex} \newcommand{\lnorm}[2]{\Vert #1 \Vert_{#2}} ``` ```{=latex} \newcommand{\ltwonorm}[1]{\lnorm{#1}{2}} ``` ```{=latex} \newcommand{\innerp}[2]{\langle{#1, #2}\rangle} ``` ```{=latex} \newcommand{\pluseq}{\mathrel{\raisebox{0.19ex}{$\scriptstyle+$}}=} ``` ```{=latex} \newcommand{\defeq}{\triangleq} ``` ```{=latex} \DeclareMathOperator{\sign}{sign} ``` ```{=latex} \DeclareMathOperator{\Tr}{Tr} ``` ```{=latex} \newcommand{\balpha}{\boldsymbol \alpha} ``` ```{=latex} \newcommand{\bbeta}{\boldsymbol \beta} ``` ```{=latex} \newcommand{\bgamma}{\boldsymbol \gamma} ``` ```{=latex} \renewcommand{\hat}{\widehat} ``` ```{=latex} \renewcommand{\tilde}{ } ``` ```{=latex} \renewcommand{\>}{{\rightarrow}} ``` ```{=latex} \renewcommand{\=}{\stackrel{\triangle}{=}} ``` ```{=latex} \newcommand{\half}{\textstyle{\frac{1}{2}}} ``` ```{=latex} \newcommand{\grad}{\nabla} ``` ```{=latex} \newcommand{\conv}{\operatorname{conv}} ``` ```{=latex} \newcommand{\block}{\mathrm{block}} ``` ```{=latex} \newcommand{\mixer}{\mathrm{mixer}} ``` ```{=latex} \newcommand{\sampler}{\mathrm{sampler}} ``` ```{=latex} \newcommand{\agg}{\mathrm{agg}} ``` ```{=latex} \newcommand{\frmstate}{\mathrm{read}} ``` ```{=latex} \newcommand{\cont}{\mathrm{cont}} ``` ```{=latex} \newcommand{\poly}{\operatorname{poly}} ``` ```{=latex} \newcommand{\rank}{\operatorname{rank}} ``` ```{=latex} \newcommand{\support}{\operatorname{support}} ``` ```{=latex} \newcommand{\bX}{{\mathbf X}} ``` ```{=latex} \newcommand{\z}{{\mathbf z}} ``` ```{=latex} \newcommand{\x}{{\mathbf x}} ``` ```{=latex} \newcommand{\g}{{\mathbf g}} ``` ```{=latex} \newcommand{\ba}{{\mathbf a}} ``` ```{=latex} \newcommand{\methodname}{Recurrent Transformer} ``` ```{=latex} \newcommand{\methodnames}{Recurrent Transformers} ``` ```{=latex} \newcommand{\methodnameabv}{\textsc{RT}} ``` ```{=latex} \newcommand{\qkRMS}{\mathrm{\textcolor{magenta}{RMS}}} ``` ```{=latex} \newcommand{\temp}{\textcolor{red}{\mathrm{temp}}} ``` ```{=latex} \newcommand{\zmark}[1]{\textcolor{blue!55!black}{#1}} ``` ```{=latex} \maketitle ``` Introduction ============ Transformers [@vaswani2017attention] are highly effective sequence models, but their computation across positions is structurally shallow: within each layer, position $t$ attends to key--value pairs computed from the previous layer embeddings, allowing essentially at most one interaction per layer between any pair of positions. A growing body of theory studies the fundamental limitations implied by bounded depth in attention models, including circuit-complexity characterizations of what low-depth Transformers can and cannot represent [@merrill2022saturated; @liu2022shortcuts]. These perspectives motivate architectures that achieve greater effective depth. We introduce the `\newterm{\methodname}`{=latex} (`\methodnameabv`{=latex}), a simple modification of how key--value pairs are computed that makes each layer temporally recurrent. In a standard Transformer, at layer $\ell$, the key--value pair at position $t$ is computed from the layer-$(\ell-1)$ representation at that position and can then be attended to by later positions $t' > t$. In the `\methodname`{=latex}, by contrast, the key--value pair at position $t$ in layer $\ell$ is computed from that position's output at layer $\ell$, rather than from its layer-$(\ell-1)$ representation. Consequently, a later position $t < t'$ at layer $\ell$ attends to a representation at $t$ that already reflects layer $\ell$ attention and MLP computation. Importantly, `\methodname{}`{=latex} performs this recurrence separately within each layer, so each layer maintains its own key--value memory. This differs from the Feedback Transformer [@fan2020feedback], which uses a shared memory across layers, and this layerwise separation is a key reason why our architecture can be implemented efficiently. We motivate `\methodname{}`{=latex}'s design through lenses of representation, optimization and computational efficiency: ```{=latex} \begin{figure*}[t] \includegraphics[width=\textwidth]{diagrams/rt_flow.pdf} \caption{One layer of the \methodname{} mapping input embeddings $\vx_1\ldots\vx_4$ to output embeddings $\vz_1\ldots \vz_N$. Notice how the \emph{persistent} key--value pairs are a function of the layer's output and are used for all subsequent attention computations. The \emph{temporary} key--value pairs are only used at the time they are computed and then discarded. They are only used to avoid ill-defined attention since, for example, $\va_2$ cannot attend to $(\vk_2, \vv_2)$ as that indirectly depends on it. This is in contrast to a vanilla Transformer that uses these same key--value pairs for all subsequent attention computation as well.} \label{fig:rt-arch} \end{figure*} ``` #### (i) Representational perspective. `\methodnames{}`{=latex} retains per-token key--value memory just like a Transformer, but increase the space of computations that can be expressed within a single layer by allowing later positions to attend to representations that have already undergone attention and MLP processing. Under mild assumptions, `\methodnames{}`{=latex} can emulate standard Transformer behavior; conversely, by restricting attention to the previous position, they can implement token-to-token recurrent computation. This positions `\methodname{}`{=latex} between fully parallel attention and purely recurrent state-space computation, while avoiding a capped-memory bottleneck. #### (ii) Training Stability. Viewing the model as a directed computation graph over positions, classical RNNs transmit information from position $i$ to $j$ only through the length-$(j-i)$ chain of intermediate states. The potentially large length of such paths gives rise to vanishing and exploding gradient phenomena [@bengio1994long; @pascanu2013difficulty], making it hard to ensure information flow between distant positions. `\methodname{}`{=latex} alleviates this by introducing many additional multi-hop paths, corresponding to repeated attend+MLP applications across positions within a layer, while still permitting direct one-hop attention interactions between any two positions. In practice, we find that this architecture, together with appropriate normalization before key--value computation and standard depth-wise residual scaling [@bordelon2023depthwise; @yang2023tensor], trains stably. We expand on this view, and on why exploding gradients are not expected to be an issue, in Section `\ref{sec:paths}`{=latex}. #### (iii) Training-time efficiency. A naive implementation of `\methodname{}`{=latex} training/prefill is sequential in position and appears bandwidth-bound: keys and values are revealed one position at a time, and each query must aggregate over a linearly-growing prefix, leading to a very low effective arithmetic intensity -- $\Theta(1)$ -- under the Roofline model [@williams2009roofline]. We give an *exact* tiling algorithm that preserves the mathematical attention computation while reorganizing memory movement, reducing high-bandwidth memory (HBM) traffic from $\Theta(N^2)$ to $\Theta(N\log N)$ and raising effective arithmetic intensity to $\Theta(N/\log N)$. Our key observation is that, during training/prefill, the full sequence of queries is available in advance even though persistent key--value pairs are revealed causally. This makes it possible to reorganize the computation into a tiled schedule, in the spirit of Flash Inference [@flashInference], that reuses each revealed key--value tile across many future queries before it is evicted from fast memory. The final algorithm interleaves attention blocks and MLP computation while employing the same methodology as [@rabe2021blockAttention; @dao2022flashattention] to accumulate attention contribution. #### (iv) Depth to inference efficiency. Crucially, the additional effective temporal depth can translate into a better depth--width tradeoff: at fixed parameter count, achieving the same quality with fewer layers reduces the amount of stored key--value state and the corresponding decode-time memory traffic. Our experiments support this regime, with shallower `\methodname{}`{=latex} models outperforming deeper Transformer baselines. #### Contributions. - In Section `\ref{sec:arch}`{=latex}, we propose the `\methodname{}`{=latex} (`\methodnameabv`{=latex}), a layerwise recurrent attention architecture that computes each layer's key--value pairs from that layer's outputs rather than from the previous layer's representations. - In Section `\ref{sec:repr}`{=latex}, we provide representational arguments showing `\methodname{}`{=latex} can emulate standard self-attention behavior and can implement token-to-token recurrent computation via attention concentration under mild assumptions. - In Section `\ref{sec:paths}`{=latex}, we provide a path-based analysis of training stability in `\methodname{}`{=latex}, showing how the architecture combines additional multi-hop computation with direct one-hop attention paths, and giving theoretical evidence in a simplified setting that neither exploding gradients nor vanishing gradients are expected under appropriate scaling. - In Section `\ref{sec:tiling}`{=latex}, we provide an *exact*, IO-aware tiling algorithm for prefill/training that preserves the mathematical attention computation while reducing memory traffic from $\Theta(N^2)$ to $\Theta(N\log N)$ and increasing effective arithmetic intensity from $\Theta(1)$ to $\Theta(N/\log N)$. - In Section `\ref{sec:comp_challenges}`{=latex}, we outline various computational challenges and design choices required to make `\methodname{}`{=latex} training more efficient and practical. - In Section `\ref{sec:expts}`{=latex}, we present empirical results on 300M-parameter C4 pretraining showing improved cross-entropy over parameter-matched Transformer baselines and favorable depth--width tradeoffs at fixed parameter count (as shown in Figure `\ref{fig:c4-300m}`{=latex}). In particular, `\methodname{}`{=latex} with $6$ layers performs comparably to $12$ layers (fixed parameters), reducing KV cache size by approximately 30% and lowering decode-time memory traffic, thereby improving inference efficiency. Additional results for the 150M-parameter model are provided in Appendix `\ref{app:150m-pretrain}`{=latex}. ![C4 pretraining: loss curves for 300m parameter model trained on C4 dataset.](plots/best_runs_val_300m_512.png){#fig:c4-300m width="\\linewidth"} ```{=latex} \hfill ``` ```{=latex} \captionof{table}{C4 pretraining loss for 300M parameter model.} ``` `\label{tab:c4-512-300m}`{=latex} Model Layers Width Val CE $\downarrow$ ------------------------- -------- -------- --------------------- Transformer 6 $2048$ $2.917$ Transformer 12 $1408$ $2.896$ Transformer 24 $1024$ $2.892$ `\methodname{}`{=latex} 12 $1408$ $2.867$ `\methodname{}`{=latex} 6 $2048$ $2.86$ Architectural overview and notation {#sec:arch} =================================== #### Architectural overview. Relative to a standard causal Transformer, the defining change in `\methodname{}`{=latex} is where the key--value pairs exposed to future positions come from. In a standard Transformer, the key--value pair at position $i$ is computed from the layer input at that position. In `\methodname{}`{=latex}, by contrast, the *persistent* key--value pair at position $i$ is computed from that position's layer output. Consequently, later positions attend to earlier positions whose representations have already undergone same-layer attention and MLP computation, making each layer recurrent along the temporal axis. This creates a circularity at the current position: because the layer output at position $i$ also attends to the current position, the persistent pair $(\vk_i,\vv_i)$ cannot itself be used while computing that output. To resolve this, `\methodname{}`{=latex} distinguishes between two kinds of key--value pairs. A *temporary* pair, computed from the current layer input, is used only when evaluating attention at the current position. A *persistent* pair, computed from the resulting layer output, is then stored and made available to all later positions. #### Notation. We present the single-head formulation; multihead attention applies the same construction independently per head and then uses the usual output projection. We assume a sequence length of $N$ and use $L$ for the number of stacked layers. Let $D$ be the embedding dimension and consider a single layer with inputs $\vx_1,\ldots,\vx_N\in\R^D$. Let $\MLP:\R^D\to\R^D$ denote the MLP block and let $\RMS:\R^D\to\R^D$ denote Root Mean Square normalization [@RMSNorm]. While in practice we use learnable parameters, as far as presentation and analysis is concerned, we take $\RMS(x)=\sqrt{D} \cdot \vx / \ltwonorm{\vx}$. We use (magenta) $\qkRMS$ to distinguish query/key normalization [@dehghani2023scalingQKNorm]. The attention operator $\Attn:(\R^D\times\R^D)^{*}\times \R^D\to\R^D$ maps a sequence of key--value pairs $(\vk_1,\vv_1),\ldots,(\vk_\ell,\vv_\ell)$ and a query $\vq$ to $$\begin{aligned} \Attn\big((\vk_1,\vv_1),\ldots,(\vk_\ell,\vv_\ell),\vq\big) = \sum_{i=1}^\ell {\vv_i \cdot \frac{\exp(\innerp{\vk_i}{\vq})} {\sum_{j=1}^\ell{\exp(\innerp{\vk_j}{\vq}})} }\end{aligned}$$ We use projection matrices $Q,K,V\in\R^{D\times D}$ to compute queries, keys and values based off an embedding. Following standard Transformer parameterizations [@bordelon2023depthwise; @yang2023tensor], we use pre-LN [@xiong2020layerPreLN] and assume attention and MLP residual updates are initialized/parameterized with an appropriate $1/\sqrt{L}$ scale so chaining maps of the form $\vx\mapsto \vx+ \frac{1}{\sqrt{L}} \{\Attn,\MLP\}(\RMS(\vx))$ is well-behaved. The Transformer layer {#sec:baseline-layer} --------------------- We first recall a standard *causal* decoder-only Transformer layer [@vaswani2017attention]. Given inputs $\vx_1,\ldots,\vx_N\in\R^D$, position $i$ forms its query, key, and value from the current layer input: $$\begin{aligned} \vq_i &= \qkRMS[Q\,\RMS(\vx_i)], \\ \vk_i &= \qkRMS[K\,\RMS(\vx_i)], \\ \vv_i &= V\,\RMS(\vx_i).\end{aligned}$$ The attention output at position $i$ is then computed by attending over the prefix of key--value pairs available up to that position: $$\begin{aligned} \va_i &= \Attn\big((\vk_1,\vv_1),\ldots,(\vk_i,\vv_i),\vq_i\big).\end{aligned}$$ Finally, the layer output is obtained by adding the attention and MLP residual branches: $$\begin{aligned} \vy_i &= \vx_i + \frac{1}{\sqrt{L}}\left(\va_i + \MLP[\RMS(\vx_i+\frac{1}{\sqrt{L}}\va_i)]\right).\end{aligned}$$ The key structural point is that, in a standard Transformer, the key--value pair stored at position $i$ is computed from the layer input at the same position. The `\methodname{}`{=latex} layer {#sec:rt-layer} --------------------------------- `\methodname{}`{=latex} layers (illustrated in Figure `\ref{fig:rt-arch}`{=latex}) differ from standard Transformer layers only in how the key--value pairs exposed to future positions are formed. At position $i$, `\methodname{}`{=latex} first forms the query together with a *temporary* key--value pair from the current layer input: $$\begin{aligned} \vq_i &= \qkRMS[Q\,\RMS(\vx_i)], \\ \vk_i^{\temp} &= \qkRMS[K\,\RMS(\vx_i)], \\ \vv_i^{\temp} &= V\,\RMS(\vx_i).\end{aligned}$$ These definitions are identical to the Transformer's query, key, and value projections at position $i$. The attention output at position $i$ is then computed using the persistent key--value pairs from earlier positions together with the temporary pair at the current position: $$\begin{aligned} \va_i &= \Attn\big( (\vk_1,\vv_1),\ldots,(\vk_{i-1},\vv_{i-1}), (\vk_i^{\temp},\vv_i^{\temp}), \vq_i \big).\end{aligned}$$ We next form the layer output representation $$\begin{aligned} \vz_i &= \vx_i + \frac{1}{\sqrt{L}}\left(\va_i + \MLP[\RMS(\vx_i+\frac{1}{\sqrt{L}}\va_i)]\right),\end{aligned}$$ which is both the representation passed to the next layer and the source from which the persistent key--value pair at position $i$ is computed. We define that persistent pair by projecting from this output: $$\begin{aligned} \label{eq:persistentK} \vk_i &= \qkRMS[K\,\RMS(\vz_i)], \\ \label{eq:persistentV} \vv_i &= V\,\RMS(\vz_i).\end{aligned}$$ Thus, $\vk_i^{\temp}, \vv_i^{\temp}$ is used only to compute attention at position $i$; it is not exposed to future positions. The persistent pair, by contrast, is defined only after $\vz_i$ has been formed and is then stored for use by all later positions. Thus, unlike in a standard Transformer, future positions attend not to a pair computed from the layer input at position $i$, but to one computed from the already-updated representation $\vz_i$. We reuse the same projection matrices $K$ and $V$ for both the temporary and persistent key--value pairs. Consequently, `\methodname{}`{=latex} does not introduce additional key/value projection parameters relative to a Transformer; this reuse also preserves a shared semantics between the temporary and persistent key--value representations. Closest Related Work {#sec:closest-related-work} -------------------- The closest representational relatives are Feedback Transformer variants. Feedback Transformer [@fan2020feedback] uses a cross-layer feedback memory shared across depth, essentially having just one list of key-value pairs computed based on the whole model's output rather than independently at each layer. Staircase Attention [@ju2021staircase] generalizes Feedback Transformers, studying recurrent processing and caching variants with weight sharing -- still at a model rather than layerwise level. This separation matters not only representationally but also computationally: within an `\methodnameabv{}`{=latex} layer, all queries are available early, which is the enabling condition behind our efficient training methodology (Section `\ref{sec:tiling}`{=latex}). TransformerFAM [@hwang2024transformerfam] is closer in that it also operates independently at each layer and allows later positions to access more processed representations. However, it does so through a bounded memory that is read from and written to via attention. By contrast, `\methodname{}`{=latex} retains per-token persistent key--value memory rather than compressing past information into a fixed-size state. This difference is important both for avoiding a bounded-memory bottleneck and for the representational results of Section `\ref{sec:repr-transformers}`{=latex}. Representational Perspective {#sec:repr} ============================ In this section, we theoretically show that `\methodname{}`{=latex} can emulate both a Transformer and RNN under mild assumptions. This shows that it subsumes both RNNs and Transformers, at least in the representation power. Representing Transformers {#sec:repr-transformers} ------------------------- Intuitively, `\methodnames{}`{=latex} can recover the behavior of a standard Transformer of lower width by ensuring that the persistent key--value pairs computed from $\vz_i$ track those that would have been computed from $\vx_i$ via $K$ and $V$ projections. We concretize this statement below: ```{=latex} \begin{thm_samy}[informal]Any width-$d'$ Transformer can be approximately simulated by a width-$d=3d'$ \methodname{}: the simulated Transformer activations can be embedded into disjoint feature groups of the \methodnameabv's embeddings. The \methodnameabv{} layer can be parameterized so that (i) attention scores are preserved and (ii) the layer output exactly tracks the Transformer layer output. \label{thm:infrml-gen} \end{thm_samy} ``` At a high level, the construction relies on representing smaller Transformer states inside a larger embedding of `\methodnameabv{}`{=latex} by dedicating disjoint feature blocks to different roles. One block stores a protected copy of the layer input $\vx_i$ so that when `\methodnameabv{}`{=latex} later computes persistent keys and values from the layer output $\vz_i$, it can still recover exactly the same key--value pairs the Transformer would have computed from $\vx_i$. A separate block is used to hold the attention contribution $\va_i$, so that when adding it to $\vx_i$ prior to applying the MLP, the contents of $\vx_i$ are protected from being lost. In this way, later tokens see identical attention scores, while the layer output matches the Transformer's layer output in the designated block. The width overhead of a factor of $3$, rather than $2$, is a subtle technical requirement for stacking multiple layers; a single layer can be replicated with an overhead of $2$. The complete construction and formal proof are provided in Appendix `\ref{app:transformer-sim}`{=latex}. Representing token-to-token recurrence {#sec:repr-rnn} -------------------------------------- If, using positional embeddings or biases, attention concentrates locally to the previous position, `\methodnames{}`{=latex} implement an RNN-like update. Formally, if $\va_i$ is dominated by the previous persistent value $\vv_{i-1}$, i.e. $\innerp{\vq_i}{\vk_{i-1}} \gg \innerp{\vq_i}{\vk^{\temp}_i}$ and $\innerp{\vq_i}{\vk_{i-1}} \gg \innerp{\vq_i}{\vk_j}$ for any $j < i - 1$, then we get that: $$\begin{aligned} \vz_i &\approx \vx_i + \vv_{i-1} + \MLP[\RMS(\vx_i+\vv_{i-1})] = V\,\RMS(\vz_{i-1}) + \vx_i + \MLP[\RMS(\vx_i+V\,\RMS(\vz_{i-1}))]\end{aligned}$$ Under the additional simplifying assumption that $V$ is the identity, this becomes a particular state recurrence with a skip connection: $$\begin{aligned} \vz_i = \RMS(\vz_{i-1}) + \vx_i + \MLP[\RMS(\vx_i+\RMS(\vz_{i-1}))]\end{aligned}$$ We do not claim to reproduce gated RNN/LSTMs, nor that training would yield to learning such structures. We stress that representationally, `\methodname{}`{=latex} is rich enough to express explicit iterative computation within a layer while also retaining full-prefix per-token memory (which was required to simulate Transformers in the previous section). Crucially, once an architecture can represent such iterative computation, a natural question is whether the classic learnability issues of RNNs [@bengio1994long; @pascanu2013difficulty] impede training. Section `\ref{sec:paths}`{=latex} explains why `\methodnames{}`{=latex} multi-hop dynamics can still train stably. Why temporal depth matters -------------------------- Transformers are shallow-through-time: deeper iterative computation along the sequence must be simulated primarily by stacking layers. Theory on low-depth attention models and finite-automata tracking problems suggests that bounded depth can have concrete consequences, with shallow Transformers being representationally insufficient to simulate certain automata [@liu2022shortcuts] and more generally bound to TC0 - a class of shallow circuits [@merrill2022saturated]. `\methodnames{}`{=latex} expose additional temporal depth within each layer. This is complementary to the depth obtained from stacking layers, thus pointing to the potential of achieving matching Transformers' effective depth while using fewer layers. We corroborate this hypothesis empirically in Section `\ref{sec:expts}`{=latex}. Training Stability of `\methodname{}`{=latex} {#sec:paths} ============================================= In this section, we explain how `\methodname{}`{=latex} manages to avoid degenerate dynamics such as gradient vanishing or exploding through depth. We formalize our arguments by viewing the model as a directed computation graph over positions: there is an edge $i\to j$ when the computation at position $j$ *directly* depends on quantities computed at position $i$. In a classical RNN, information (and gradients) from position $i$ to $j$ must traverse the full chain $i\to i\!+\!1\to\cdots\to j$, and repeated composition along long chains leads to vanishing/exploding gradient phenomena [@bengio1994long; @pascanu2013difficulty]. This chain topology forces all influence from position $i$ to $j$ through $(j-i)$ successive state transitions. Stabilizing training typically requires these transitions to be close to contractive, but then the influence of $\vx_i$ on $\vx_j$ shrinks rapidly with $(j-i)$, making distant information difficult to transmit. While in RNNs this issue can be alleviated through careful initialization schemes [@orvieto2023resurrecting], our method takes advantage of the fact that there are both direct hops, as well as additional multi-hop paths between layers. As in a standard Transformer, token $j$ can directly attend to any earlier token $it$. Chaining these write--read operations yields multi-hop influence paths whose length scales with the distance between positions, enabling within-layer iterative computation (as in Section `\ref{sec:repr-rnn}`{=latex}) while preserving the direct one-hop attention routes of a Transformer. #### Dampening long paths without eliminating long-range access. Multi-hop paths are only useful if they do not explode. In practice, two levers dominate. First, standard depth-wise scaling conventions for residual branches keep per-layer updates in a stable range [@bordelon2023depthwise; @yang2023tensor]. Second, normalization preceding computation of persistent keys/values (the $\RMS(\vz_i)$ inside Equations `\ref{eq:persistentK}`{=latex}-`\ref{eq:persistentV}`{=latex}) controls magnitudes even though $\vz_i$ is a sum of multiple components. Empirically, these choices place long multi-hop influences on the vanishing end: longer chains have smaller effect. Unlike a pure RNN, this does not remove long-range access because direct attention edges remain available even when long chains are damped. In the theorem below, for a very simplified setup without normalization, we show that we do not get exploding gradients at initialization. Normalization helps in stability, that is, it allows stable training of `\methodname{}`{=latex} with higher learning rates. This is demonstrated empirically in Appendix `\ref{app:layernorm_stable}`{=latex}. ```{=latex} \begin{theorem} \label{thm:train_stable} Consider a simplified 1-layer uniform-attention only \methodnameabv{} layer with inputs given by $x_1,...,x_n$ and outputs denoted as $z_1,...,z_n$, where \[ z_k = x_k + \frac{\alpha}{k} \left( V x_k + V \sum_{j=1}^{k-1} z_j \right) \] where $\alpha$ is a scalar denoting the scaling of the residual and $V$ is the value matrix. Then, for $k \geq 2$, \[ \frac{\partial z_k}{\partial x_1} = \frac{1}{k!} \sum_{r=1}^k {k \brack r}\,\alpha^r V^r \] where ${k \brack r}$ denotes the total number of permutations of k elements having exactly $r$ cycles. \end{theorem} ``` As the total number of permutations is $k!$, the theorem above shows that as long as the maximum eigenvalue of $\alpha V$ is smaller than $1$, we do not get an exploding gradient from $z_j$ to $x_1$. Thus, for orthonormal initialization, for any $\alpha < 1$, we expect to be in this regime. Moreover, since the overall gradient is summed over paths of various lengths (given by $r$ in the above expression), we can see that we have non-vanishing gradient even when the maximum eigenvalue of $\alpha V$ is smaller than $1$. In particular, since ${k \brack 1} = (k-1)!$, the term in the above expression corresponding to $r=1$ is $\frac{\alpha}{k} \cdot V$, which is precisely the gradient a vanilla transformer would yield. Proof for this theorem can be found in Appendix `\ref{app:train_stable}`{=latex}. ![ We use the tiling of [@flashInference] to increase arithmetic intensity during the forward pass since $(\vk_t, \vv_t)$ only become available after the attention output $\va_t$ is computed - this in turn happens once position $t$ has attended to all previous KVs. ](diagrams/RT_fast_tiling_8.png){#fig:tiling width="0.6\\columnwidth"} Exact Tiling for Training and Prefill {#sec:tiling} ===================================== #### What makes naive evaluation slow. During training/prefill, `\methodnames{}`{=latex} are fundamentally sequential in position; to compute the persistent pair $(\vk_t,\vv_t)$ we must first compute $\vz_t$, and $\vz_t$ depends on $\va_t$, which aggregates over all previous persistent key--value pairs: $\{(\vk_i, \vv_i)\}_{i