Sequential Parallel Duality in Prefix Scannable Models
Authors: Morris Yau
Sharut Gupta
Valerie Engelmayer
Kazuki Irie
Stefanie Jegelka
Jacob Andreas
Venue: ICLR 2026 Poster (2026)
Keywords: state space models
linear RNNs
linear transformers
sequence modeling
Primary area: foundation or frontier models, including LLMs
TL;DR: Propose new architecture Transformer-PSM that generalizes state space models to use softmax attention

Abstract: Modern neural sequence models are designed to meet the dual mandate of parallelizable training and fast sequential inference. Recent developments have given rise to various models, such as Gated Linear Attention (GLA) and Mamba, that achieve such ``sequential-parallel duality.'' This raises a natural question: can we characterize the full class of neural sequence models that support near-constant-time parallel evaluation and linear-time, constant-space sequential inference? We begin by describing a broad class of such models -- state space models -- as those whose state updates can be computed using the classic parallel prefix scan algorithm with a custom associative aggregation operator. We then define a more general class, Prefix-Scannable Models (PSMs), by relaxing the state aggregation operator to allow arbitrary (potentially non-associative) functions such as softmax attention. This generalization unifies many existing architectures, including element-wise RNNs (e.g., Mamba) and linear transformers (e.g., GLA, Mamba2, mLSTM), while also introducing new models with softmax-like operators that achieve O(1) amortized compute per token and log(N) memory for sequence length N. We empirically evaluate such models on illustrative small-scale language modeling and canonical synthetic tasks, including state tracking and associative recall. Empirically, we find that PSMs retain the expressivity of transformer-based architectures while matching the inference efficiency of state space models -- in some cases exhibiting better length generalization than either.

Published as a conference paper at ICLR 2026
SEQUENTIAL-PARALIE DUALITY IN PREAX-SCANNABLE MODELS
Morris Yau" MIT CSAIL morrisyêmie edu
Sharut Gupta' MIT CSAIL sha marrttmitt. edu
Valerie Engelmayer TU Munich
engeeeeenggllen wengelmmyerr@uum
Kazuki Irie
Jacob Andreas MIT CSAIL jda@mit edu
Stefanie Jegelka
Harvard University kirielg. harvard.
TU Munich, MIT CSAIL stefje@mit.
ABSTRACT
Modern neural sequence models are designed to meet the dual mandate of parallelizable training and fast sequential inference. Recent developments have focused on various models, such as Gated Linear Attention (GLA) and Mamba, that achieve such "sequential-parallli duality." This raises natural question: can we characterize the full class neural sequence models that support near-constant-time parallel evaluation and linear-time, constant-space sequential inference? We begin by describing a broad class of such models, state space models, those whose state updates be computed using the classic parallel prefix scan algorithm with a custom associative aggregation operator. We then define more general class, Prefix-Scannable Models (PSMs), by relaxing the state aggregation operator to allow arbitrary (potentially non-associative) functions such as softmax attention. This generalization unifies many existing architectures, including element-wiss RNNs (e.g., Mamba) and linear transformers (e.g.. GL.A, Mamba2, mL.STM., while also introducing new models with softmax-like operators that achieve O(1) amortized compute per token and log( N) memory for sequence length N. We empirically evaluate such models on illustrative language modeling and canonical synthetic tasks, including state tracking and associative recall. Empirically, we find that PSMs retain the functional effectiveness transformer-base architectures while matching the inference efficiency of state space models and some cases exhibiting better length generalization than either.
INTRODUCTION
Transformers have revolutionized sequence processing by enabling parallelizable training over the sequence dimension (Vaswani et al., 2017) 2017)-unlike unlike classic recurrent neural networks (RNNs) (Elman, 1990; Jordan, 1986; Hochreiter & Schmidhuber, 1997), which require sequential training: and by handling arbitrary-length sequential dependencies with constant parameter count-unlike convolutional neural networks, which, while parallelizable over sequence elements, require more parameters to capture longer-range dependeneies (Gehring et al., 2017: Oord et al., 2016; Kalchbrenner et al., 2016; Dauphin et al., 2017). However, transformers suffer from two fundamental limitations: first, their computational and memory complexities scale quadratically with sequence length (Vaswani et al., 2017; Katharopoulos et al., 2020), which is particularly problematic during inference; second, they have limited expressivity, i.e., there are computations they struggle to learn to perform (Hahn, 2020; Bhattamishra et al., 2020a;b; Merrill & Sabharwal, 2023; Irie et al., 2023; Merrill et al., 2024; Grazzi et al., 2025; Strobl et al., 2024; Siems et al., 2025; Movahedi et al., 2025).
A body of research in neural sequence modeling has focused developing architectures that address the primary shortcomings of transformers. In particular, recent years have seen the introduction of diverse models that target the inference time complexity problem. In these models, the inference
Equal contribution.

---

Published as a conference paper at ICLR 2026
compute requirement is linear in time and constant in memory. just like in classic RNNs, while retaining transformer-like parallelizability during training. Such models include element-wer wise recurrent models, which are derived by simplifying either fully recurrent neural networks (Hochreiter & Schmidhuber, 1997) (e.g. Quasi RNNs (Bradbury et al., 2017) or SRU (Lei et al., 2018); see also (Qin al., 2023; Li et al., 2018; Balduzzi & Ghifary, 2016; Mozer, 1989) linear time-invariant dynamical systems (e.g. Mamba (Gu & Dao, 2024))-at at the cost of sacrificing expressivity (Merrill al., 2020; Grazzi al., 2025). Another model family has been derived from linear transformers (Katharopoulos et al., 2020) and fast weight programmers (Schmidhuber, 1992; Irie et al., 2021), including Gated Random Feature Attention (Peng et al., 2021). DeltaNet (Schlag et al., 2021; Yang et al., 2024b), RetNet (Sun et al., 2023), GLA (Yang et al., 2024a), mLSTM in xLSTM (Beck et al., 2024), Mamba2 (Dao & Gu, 2024), and versions of RWKV (Peng et al., 2025).
These models share the fundamental property of sequential- parallelizable over sequence elements, while inference is sequential and its inference time complexity is linear. This raises a natural question: What is the class of neural sequence models that can be evaluated in parallel in nearly constant depth, and sequentially in nearly constant space?
duality (SPD) is
In this work we aim to characterize the family of models exhibiting SPD. In particular, we show that these models are computable using the classic parallel prefix scan algorithm (Blelloch, 1990; Martin & Cundy, 2018) with a choice of associative aggregation operator that is specific to each model. We define a broader model class, which we call -eeen--Scannble Models (PSMs). by generalizing the aggregation operator used in prefix scan computation. By construction, this family subsumes all existing SPD-compatible models with associative state updates. More generally. it enables the design of novel models with non-associative aggregation rules, whose per-token inference cost remains amortized (1) with memory scaling (N))(NN) in sequence length N. An alternate view is that PSMs are strict generalization of RNNs: they move beyond affine state updates to support general Transformer-style self-attention over local chunks -iving rise to novel model belonging to the PSM family, which we call Transformer-PST
token mixing operations aperations-includding
We probe Transformer-PSM in our experiments using small but illustrative tasks: next-token prediction on WikiTexx--033 (Merity et al., 2017) and synthetic algorithmic tasks that test precise state tracking and retrieval (Merrill et al., 2024; Grazzi et al., 2025; Li et al., 2025; Arora et al., 2024). We find that Transformer- PSMs inherit certain advantages of both Transformers and State Space Models. They preserve the associative recall capability of Transformers, whilst exhibiting an impressive ability to track state. Furthermore, by varying the "chunk" size by which we break up sequence of tokens, can alter the asymptotics of a PSM from SSM-like to Transformer-like notion we make precise in our discussion on Sequential-Pan Paralllel Duality, which we empirically demonstrate WikiText--033. In summary,
1. We define the SPD family of sequence models and unify modern linear RNNs as those with state computable by the prefix scan algorithm with a custom choice of associative aggregator.
2. We derive strict generalization thereof, the Prefix Scannable Models (PSMs), that admit general state aggregation functions, such as softmax attention, whilst preserving parallel training in ((N) compute and Ollog N) memory bound at inference.
3. We instantiate Transformer-PSS PSM and evaluate its abilities for state tracking, associative recall, and language modeling, using canonical sequence modeling benchmarks.
2 SEQUENCE MODELS AND SEQUENTIAL- EEQUENTIAL-PARALLEI DUALITY
Here we formally define sequence models and sequential-paralle duality, and provide examples. For more details on conventions, we refer Appendix A. Throughout, let A be a finite alphabet of tokens and au:n-1 E A an input sequence of length n. Let M be latent space containing the state of a sequence model. For example, for an RNN, this is the space of the hidden state vector. First, for the sequential view, we define causal sequence models by introducing state dynamics and inference.
Definition 2.1 (State kernel). state kernel is map U: M x - M with an identity element E M. It induces a state sequence $ = e, (SL-1, a+) for t 2 0. We denote by m(n) the memory required to store sn-1-
2

---

is Published as a conference paper at ICLR 2026
Definition 2.2 (Inference module). An inference module is map F: M A RlAl producing a distribution D F(( az) over the next token.
Definition 2.3 (Sequence model). A pair (U, F) comprising state kernel and inference module called a causal sequence model (or simply, sequence model). The model's memory bound m(n) required to evaluate U and F once the state Sn-1 is available.
Second, to formalize parallel training, we define a parallel training circuit for sequence models.
Definition 2.4 (Parallel circuit family), A parallel circuit family for sequence model (U, F) is uni- form family of circuits { {Cn, I 1221 such that, for all E A and all < n, [Cn(a)], F ((s-11, az), where - is the state (Def. 2.1). The model's compute bound T(n) is the size of the eircuit Cn.
The circuit corresponds to the training graph: every token can be processed simultaneously provided sufficient parallel hardware. Together, the sequential and parallel views and their tradeoff's will characterize the Sequential-Paralee Duality (Def. 2.5).
Definition 2.5 EeqqeentialPaaralle Duality SPD (T)).. m(n))). A sequence model (U, is said to satisfy (T)()(m), m(n)) if the following two conditions hold:
1. Parallel training. There exists uniform circuit family {CrIn21 of depth O(1) and size (nn) that realises all token-wise predictions (Def. 2.4).
2. Sequential inference. Given St-1, the pair (st. ÿ:) computable by depth-O(1) circuit using at most m(n) working memory.
As illustrative examples, we discuss the following sequence models in light of SPD.
Vanilla Transformer: SPD-(n2, n). Training computes all n² attention scores in parallel with circuit depth O(1) and work T(x) ( At inference, each new token requires attending to and storing all n past keys/values, yielding (n)) = (n1) memory.
Fully recurrent RNN: SPD. A strict RNN (e.g. LSTM, GRU) updates its hidden state through chain length n. Because each step depends on the previous one, there is no sub-linear-depte eircuit that simultancously computes every output. Such networks therefore fall outside the SPD framework.
As a preview of our results: we will additionally derive the following characterization.
Prefix-Scannable and Related Models: SPD-(n, 1) and SPD-(n, log(n))) Modern RNN architec- tures that admit a Blelloch-style scan (discussed Sec. 3) for their state update have compute bound T(n) (n)). parallel depth O(log n), and memory bound m(n) O(log n) m(n) @(1), depending on whether the state size grows logarithmically or remains constant. We therefore write SPD-(n, log n.) or SPD-(n, 1). both of which strictly improve the Transformer's linear memory latency while retaining fully parallelisable training.
3 PREFIX-SCANNABE MODELS
Next, we define a broad family of models that obtain a sequential-paralle duality of SPD-(n, log(n)). This family consists of sequence models whose training graph can be expressed by Blelloch prefix scan (see the caption in Fig. 1) over chunk representations, followed by an independent chunk-local prediction head. The Blelloch scan takes sequence of tokens or chunks and an aggregation operator, and computes prefixes where the aggregator is applied over the first n tokens; it computes all prefixes in (n)) work and O(log n) parallel depth. We refer to Alg. in Sec. 3.3 for the full upsweep/ddwnsweep algorithms. We call these Prefix-Scannable Models (PSMs). To understand the topic further we first give a brief overview of the classic parallel prefix sean.
Bielloch Scan. Let M be set with binary operator Agg .M x M M and identity E M. Given ag...., an-1 E M. the exclusive prefix at index t is P. == an Agg aj Agg ... Agg a;-1 (with Po e). The Blelloch prefix-scan computes all P, in O n) parallel steps via perfect binary tree: (i) an upsweep reduces adjacent pairs bottom-up until the root aggregates the whole sequence; (ii) downsweep propagates prefixes top-down, using stored intermediate values that every leaf
3

---

Published as a conference paper at ICLR 2026
Inactive Token
Summary Token
Chunk
Identity
Output
x[e] x[0:3] x[2] x[0:3] x[4:5] x[6]
x[0:7]
Inpur
x[e]
x[e:3]
x[e:3] x[4]
x[4:5] x[6] x[e:7]
Output
x[e] x[3]
x
x[0:3]
x[0:6]
(a) Upsweep
(b) Downsweep
Figure An illustration of the Blelloch parallel scan used to compute prefix states Prefix- Scannable Models (PSMs). Here the input has 16 tokens grouped into 8 chunks {==00)...., [[]] (see (a) bottom), and the goal is to produce prefix states {e, x=(0, (0:1,,... x [0[06]]; where xi-:j] aggregates all tokens from chunks i to j. and is the identity. (a) In the upsweep, chunks are aggregated along a binary tree through series of chunk aggregation operations (solid arrows), producing intermediate values and some of the final prefix states (e.g.. 200:1], 2 (0:33). (b) In the downsweep, the missing prefix states are filled in by propagating values backward: ar [0:7] is reset to e, and copy (dotted arrows) and aggregation (solid arrows) operations complete the sequence. When each chunk is treated atomic element, this recovers the classic Blelloch scan.
receives Pt. When the binary operator Agg is associative, the final prefix array identical to what a left-to-right sequential loop would compute. If Agg not associative, the result is still well defined-thh tree fixes a unique parenthesisation (see discussion in Appendix Eu-but may differ from the purely left-nested order used by sequential recurrence. We refer to this upsweep-downsweep as the static (training) algorithm (Alg. 1), and to the left-to-right procedure as the online (inference) algorithm (Alg. 2), which reproduces the same tree parenthesisation. In the next section define Prefix Scannable Models (PSMs) by instantiating the static scan with a general choice of Agg.
3.1 MODEL DESCRIPTION
Definition (Prefi--Scannable Model). Fix a chunk length S n and partition sequence ag:n-1 into n/cc disjoint chunks C, (atey..., ((441)c-1). A Prefix-Scannable Model (PSM) is specified by three learnable modules with depth O(1):
Enc A M, Aggo M M M, Inf, :Mx A + A
and identity element E M.
Chunk encoding x, = Enc(C,) for i = O...., 1.
2. Prefix state {sihxe(r] BlellochScan ((z; (e(r)) Agg e).
3. Chunk prediction Dic:((1))--1 Inf Inf C.).
Note that, in terms of notation, we have 2 [0:ii] defined in Fig. 1. We discuss the asymptotics of the PSM model depending on both n and e in Appendix For now, we derive the following the leading order term and focus
immediate complexity corollary with asymptotics depending on discussing its connections to recently proposed efficient sequence models. Proposition 3.2 follows from properties of the parallel and streaming versions of the Blelloch scan.
Proposition 3.2. Every Prefix-Scannable Model is in the class SPD-(n, log n). That is, its training work is (n)) with parallel depth ((1), while online inference runs in O(1) amortised time and Ollog n.) memory per token.
4

---

Published as a conference paper at ICLR 2026
Table 1: Representative examples of recently proposed layer types that cast into the affine state-uddate
template (Eq. (B.1)). The same associative aggregator (E, f) + (E", f¹) + (E o E
f
shared by all, and therefore, they are all in SPD-(n, 1) by Theorem B.3.
Model family
E 8:-1
f:
Gate operator
Linear Attention (Katharopoulos al., 2020)
v.k! Brokk]
identity projector
DeltaNet (Schlag et al., 2021)
88-1
Gated DeltaNet (Yang et al., 2025)
01
(1
Bro
projector
RetNet (Sun al., 2023)
78t-1
vrkk
scalar gate y scalar gate fr scalar gate gr diagonal gate
mL.STM (Beek al., 2024)
f181-11
ivck
Gated RFA (Peng et al., 2021)
9:8t
(1 ge )
S4 / S6 (Gu al., 2022)
ax ① 8t-1
B o (w,1")
Mamba (Gu & Dao, 2024)
A(2))81-1
((x )ar
diagonal gate
GLA (Yang et al., 2024a)
la. 81-1
o.k.
diagonal gate
Proof Sketch. The static Blelloch depth (Alg. 1). The streaming evaluation replaces that scan by the online algorithm of Alg. 2, whose Theorem 3.5 and Corollary 3.6 show O(1) amortised work and O(log n) state. The chunk-local Inf adds constant overhead.
chunk encodings costs linear work and O n)
3.2 MODERN RNN LAYERS FIT ONE AFFINE SCAN
To relate PSMs to recent models, this section shows that broad family of recent fast-inference layers (Table are all PSM's. Their state kernel can be expressed as specializations of single associative affine state-update template. This enables SPD-(n, 1) complexity.
Definition 3.3 (Affine recurrence). Let (M, +, 0) be an additive group and >: M + M fixed bilinear action of monoid (R,, o, 1) on M. A layer is said to have an affine state update if its hidden state obeys
8, E 8:-1 + ft,
8-1 $-1=0 = 0,
(3.1)
where (E,, fi) E R x M are (learnable) functions of the current chunk x,. That is E, == Eg and f. for (Ix) for learnable functions Eg and fo.
The models in Table all satisfy this affine state update template and all share the following associative aggregator. For proof see Appendix B.
Lemma 3.4. (Associative Affine Aggregator) Define for (E., E RX
E2,fal EE f1) (Ego E,, f2 E2 fi); (T,0)-
Then (R x M, B, is monoid ① associative with identity e and
Er, ft) + ... @ (Eo, fo) -EE,, st),
where 8, is the state given by Eq. (B.1) and E, is an auxiliary variable.
Once written in the affine update form, their binary operator is associative, hence each layer is Seeti-Scaanaale Model with SPD-(n, 1) complexity. For formal theorem and proof see Theo- rem B.3. Importantly, we can instantiate Def. 3.1 with associative aggregators capturing learnable function families like linear dynamical systems and Gated Linear Attention. Further discussion and the corresponding theorems can be found in Appendix B.I. Next, we turn to general PSM's, which enables new (non-associative) aggregators, most notably softmax attention.
3.3 BEYOND AFFINE STATE RECURRENCE: PSMs WITH GENERAL AGGREGATION
The parallel prefix-scan computes per-position prefixes with O(n) work and Ologg n) depth when the binary operator is associative (Blelloch, 1990). We generalize this view to non-asson non-associative operators (e.g., softmax attention). For longer discussion with extensive proofs, see Appendix E.
5

---

Published as a conference paper at ICLR 2026
The key issue is parenthesisation: for non-associative Agg, different groupings of To Aggar 11 Agg ... Agg 21-1 produce different values. The Blelloch sean resolves this by fixing single full binary tree (upswecp/dowwsweee),, hence unique parenthesisation. Let
Agg: : M x M - M, identity e E M,
(3.2)
with no associativity assumption unless stated. Define TBielloch as the binary-tree parenthesisation induced by the static scan. The stalic Blelloch scan (Alg. 1) computes, for every i,
R (2o Agg #1 Agg Agg xt-) evaluated under TBielsch-
This matches the sequential left-to-right recurrence when Agg is associative; otherwise it is still well-defined value for the fixed tree. The work is (n),, and depth is nl)
Online binary counter (inference). The online variant (Alg. 2) maintains most one root per block size 2"; inserting performs the usual binary carry with Agg. The current prefix is the most significant bit (MSB) > least significant bit (LSB) fold of occupied roots. This reproduces exuctly TBlelloch for each / while using Ollog n) memory.
Algorithm 1: STATICBLELLOCHHSAN
Algorithm 2: BINARYCOUNTERUPDATE
Input: (fæz) }, Aggo e): Array of encoded
Input:
(root, Aggo e): Stream of encoded chunks Tn. ........ 1r-11 operator Agg with identity e
chunks q(æz) 2r--1 I| with r 2 (power of two) chunks; operator Agg with identity
Output: Prefix value pt for cach f (Blelloch parenthesisation)
Output: Exclusive prefixes written back into q
State:
Representation. Store the complete binary tree in the usual heap layout T(1.21 1|:
2 roet [k] stores the root of the current block of size 2 os empty initialise all to empty.
1. leaves T i] <- q for O,....,--
2. an internal node e has children 20 and 20+1.
3 for f + to do
carry + æ;;
h: 5-0 while root[k] empty do carry + Agg( root [k), carry)
Upsweep (reduction). for do parallel
down
T(u + Agg (20), [[2u+1]
root [k| empty
k- k+1
Downsweep (prefix propagate). Allocate P[]; set P[1] <- /:// root gets identity for v t 1 to n - do in parallel
root [k] + carry; // place merged tree
p e 12
P[2v] t Piu]:
for t [log, (++1) down to do
P[2v + 1] + Age( Pie], T(2u))
13
if root [k] * empty then Agg(p, root(k])
Write back. for i 0 to do in parallel L |
P(n +
15
emit p
Together, the static and online scans yield PSMs in SPD(n, lo::): linear work for training and loga- rithmic memory for streaming inference. (See Fig. for the chunked Transformer-De PSM inference.) We obtain the following correctness and complexity analysis. We defer proofs to Appendix E
Theorem 3.5.
Let ps be the value emitted at time t by Alg. 2. returned by the static Blelloch scan, regardless of whether Agg is associative.
Then pt equals the exclusive prefix
Corollary 3.6. After t+1 chunks Alg. 2 stores at most [logg(++1] (++1)] root values; hence the worst-case space usage is Oloog n).
Work. Inserting a new element touches exactly the trailing 1-bits of i; the expected number of such bits is 2, so the amortised number of Agg calls per element is constant.
Together, Theorem 3.5 and Corollary 3.6 show that the online binary-counter sean is an opti- mal-space, streamable realisation of the Blelloch parenthesisation, extending prefix-scan techniques to con-associative operators without increasing asymptotic cost time. This flexibility enables larger class of prefx--canndll binary operator that admits O (log n) space (11) time online evaluation via the mechanism above. We provide further analytical details of PSMs in Appendix C.
models: sequence models whose poketoken state update is any

---

Published as a conference paper at ICLR 2026
Inactive Token
Active Token
x[0:3]
x[e:1]
x[e:1]
x[e:1]
x[0:1]
x[2:3]
x[0]
x[1]
x[0]
x[1]
x[2]
x[9]
x[1]
x[2]
x
x[1] x[2]
x[3]
Figure 2: An illustration of the autoregressive state computation of "Transformer-PSM (Sec. 4) at inference time. Here the model uses a chunk size of 2. From left to right, a single new token is fed to the model at a time. Two first figures: when predicting tokens in chunk a [2]. the model only requires tokens from the prefix state æ00:1] and those within [2]]. Third figure: predicting tokens in chunk x [3] requires the prefix state x [0:11], and chunks x [2] and x []. Last figure: once tokens in prefix state z (0:3) is computed, which is later used predict tokens in x [4), and so on. Prefix state 8; corresponds to 8; ==0:i].
chunk x (3 are processed,
3.4 TRANSFORMER-PSI
In this section we instantiate Enc, Agg, Inf to concretely define the Transformer-PSl architecture that
empirics in Sec. 4 to validate theoretical predictions. The model is uniquely specified by the following three modules.
Encoder (Enc): This is simple embedding layer that transforms discrete vocabulary tokens into contimuous vectors, implemented as a standard nn.embedding layer.
Aggregation (Agga). A GPT-2 style Transformer (hidden dim d, H heads, I. layers) with bidi- rectional attention mask, GPT Rdx2e Rdx2e Given two chunk states 22 Rdxe, define token-concat [z, mty] Rdx2c and the right-half slice RH(Y)= Y[:, e: 2c] E Rdxe, We write
Aggo(zi, = RH(GPT) 2,,)) E TRdxe
Inference (Inf A GPT-2 style Transformer (hidden dim d, H heads, layers) with a causal mask, GPT Rdx22 + Rdx22 Given prefix state 8,-1 E dd and a token chunk Enc(C) E Rdxe
Inf (nf(s:-1, RH(GPT ([ Enc(C.) Rde
which we interpret as per-token logits for C,[1 (next-token prediction within chunk). Once the three modules defined, we train Transformer-PP PPM with Alg. and inference with Alg. 4.
Transformer- PSM Training (static scan over chunks)
Algorithm 3:
Algorithm 4:
Transformer-PSo PPMM
Inference
(binary-counter scan)
Input: Sequence of tokens dom. Enc, Agga- Inf chunk size Output: Predictions yo... r < na/c;
Input: Streamed tokens a, Enc, Aggo Inf chunk size
Output: Streaming predictions
State:
// number of
2 root [k] <- empty for all (ef. Alg. 2)
chunks
3 buf <- []: // collects current chunk
for i 4- t do in parallel
4 for each a, do append to buf;
// token index t= 0,1,
L x: End (e)d)))(-))--1)
{s:) +
if buu then // completed one chunk # <- Enc(buf);
STATICBLELLOCHSCANNEEE Aggo, e)
// Alg.
8 f-
BINARCCUNEENPPPAT(POO..CC Agg, e):
Agga.
// Alg. 2
for i +- to do in parallel
9 -e Infa(s,
Dre((41)e-11 +- Inf (10(1111)-1)
10
clear buf;

---

Published as a conference paper at ICLR 2026
1.0
0.8
0.8
Error Rate
0.6
GPT2 Transformer Mamba
Erar
T-PSM c84
0.4
Transforme-PSM
0.4
SWT d64
:32
0.2
0.2
0.0
0.0
18
66 90
114 138 162
Sequence Length
Sequence Length
Figure 3: Error rate on the state tracking S5 task. After training on sequences with lengths up 18, Transformer-PSM. generalizes more than 160 tokens, far beyond Transformer and Mamba.
Figure 4: Error rate on MQAR of Transformer- PSM (T-PSM), Sliding Window Transformer (SWT) and Mamba. Evaluated lengths are in- distribution.
4 EXPERIMENTAL RESULTS
The main goal of our experiments is to evaluate and explore the capabilities and properties of Transformer-PST (Sec. 3.4). For this, we conduct experiments on representative sequence learning tasks: synthetic algorithmic task requiring state-tracking (Sec. 4.1), synthetic task for associative recall, and language modeling (Sec. 4.3). Each experiment was conducted on a single NVIDIA V100- 32GB GPU. All experiments were implemented using the PyTorch framework (Paszke et al., 2019).
4.1 STATE TRACKING S5
The Sş state tracking problem (Merrill et al., 2024; Kim & Schuster, 2023; Li et al., 2025) is the formal version the "permute cups and balls" challenge, where a sequence permutations composed, with the objective of tracking the resulting permutation at each time step. Problems as diverse as tracking finite state automata and evaluation of boolean expressions can be reduced to this task. Naturally, the sequence of permutations lengthens, this task becomes increasingly difficult for a constant-depth model that has a constant budget for sequential computation. Indeed, the Ss state tracking task is NC² complete (Barrington, 1986). It is known to be difficult for both standard Transformers and linear RNNs such as Mamba (Merrill et al., 2024; Grazzi et al., 2025).
We train from scratch on sequences of length to 18 a curriculum and subsequently evaluate lengths up 180 test for length generalization. We generate 100,000 sequences per length and train for 20 epochs for each of three different models: (1) a standard GPT2 model with 12 layers, 12 heads, 768 hidden dimensions; (2) 370M-parameter Mamba model with 48 layers and a 1024-dimensional hidden state; (3) Transformer-PShe with (d 768, Hl 1, L 1) for Agg, (d 768, H 1, 1) layer for Inf, and chunk size 1. All models are trained with Adam with dropout 0. weight decay 0.01, learning rate 10-4
Fig. 3 shows the results. We find that whilst Mamba slightly outperforms GPT2, the new T-PSM has remarkably low етrоr rate even for sequences significantly longer than those observed during training, showing that these models exhibit strong length generalization for state tracking tasks.
4.2 MULTI QUERY ASSOCIATIVE RECALL (MQAR)
In Associative Recall, the task is to recall whatever value followed a key earlier in a given sequence. MQAR extends this task to multiple such key-value pairs to increase the memory demand (Arora et al., 2023). While constant state size recurrent models struggle with this task, a 2-layer transformer excels by solving it perfectly. To gauge where this spectrum models on MQAR for 64 epochs with vocabulary size 8192 and 8 key-value pairs. In the typical setting of this task, sequences
model falls, we train different
constructed in way that key is queried shortly after it appears for the first time; here we do not use such a bias and sample queries uniformly, which makes the task harder than the standard setting.
8

---

Published as a conference paper at ICLR 2026
0.06
Transformer- PSM
24.5
Evaluation Perplexity
0.05
Mamba
Time per taken (s)
GPT-2 Transformer
24.0
0.04
Transforer-PSM GPT-2 Transformer Mamba
23.5
0.03
23.0
0.02
22.5
0.01
0.00
32
64
128
256
10000
20000 30000 40000 Token index
Chunk size
Figure 5: Evaluation perplexity of Transformer- PSM with varying chunk sizes on WikiText-103
Figure 6: Inference time per token for Transformer-Po5 and GPT-2 Transformer
Here we instantiate Transformer-PSW with (d = 256, H = L, -= 2) Agg, (d = 256, H = = 1, 2) Inf. We also
learnable linear projection to compress the chunks instead of taking the right half. The chunk size is 32 or 64. For comparison, we include both Mamba and Sliding Window Transformer (SW (S) baselines (Beltagy et al., 2020; Zaheer al., 2020). The SWT is GPT2 model with (d = 256, H 1, L = 4), where we use sliding window size of 32 or 64.
Fig. 4 shows the results. Here all the evaluation lengths are the training distribution. We find that T-PSM with chunk size of 64 achieves the perfect accuracy like the full context transformer, while reducing its chunk size to yields performance degradation on long length (512). Mamba fails in our setting; unlike in prior work (Arora et al., 2024; Okpekpe & Orvieto, 2025), our setting is harder due to our uniform query sampling as discussed above.
4.3 LANGUAGE MODELING ON WIKITEXT-IOS WITH TRANSFORMER-PSM
Here evaluate perplexity the WikiText- dataset (Merity al., 2017). We benchmark Transformer-PSM. (d = 768, H = 12, L = 1) Agg, (d = 768, H = 12, L 11) Inf, by varying the self-attention chunk size from 32 256 tokens and measuring test perplexity against the vanilla GPT-2 (base) baseline with a context size of 512 trained from scratch. As shown in Fig. 5, as the chunk size grows, perplexity falls gracefully from 24.12 32 tokens 22.45 256 lakens-closed approaching vanilla GPT-2's perplexity of 22.28 demonstrating that larger chunks recover nearly full-context modeling power while preserving our model's ''' linear-time inference. For reference, we also include baseline for Mamba trained from scratch at 130m parameters, 768 hidden dimension, trained for 10 epochs with the same optimizer hyperparameters achieving ppl of 24.7.
Next, we measure per-token latency over 40,000 WikiText-2 tokens for our model versus 4-layer, 4-head, 256-dimensional GPT-2 baseline. We train Transformer-PSM (d= 768, Agg. (d 768, H 4, L = 2) Inf, thus keeping the parameter count identical to the baseline. As shown in Fig. 6, GPT-2's inference cost grows linearly with context length (O(n) per token) with KV cache, inflating latency from R 0.002s at the start 0.04s by token 40,000. In contrast, our Transformer- design reuses 64-token chunk summaries, leading to a 0 + n Iog(n/64)) inference cost (as discussed in Eq. (C2) in Appendix D), keeping per-token time below 0.008s. For reference, we also include inference time measurement for a Mamba model with layers, 256 hidden dimension, with an average inference time per token of A .006s.
4, L 2)
32
DISCUSSION AND CONCLUSION
Discussion. We give a concise conceptual view of parallelisable, inference-efficient sequence models via prefix scannability, unifying many closely related models developed under different names. Our results deepen the link between prefix-scan algorithms and efficient sequence models, extending the design space beyond prior work (Martin & Cundy, 2018).
This algorithmic lens offers a framework for analysing and designing future models. For exam- ple, concurrent work "log linear attention" (Guo al., 2025) also fits this view, proposing

---

Published as a conference paper at ICLR 2026
linear-attention mechanism with log memory, structured state, and efficient chunkwise-parallel primitive.
Conclusion. We formalise vequential-paralee duality: models that train in parallel yet decode sequentially. Recent efficient sequence models exhibit this duality and achieve linear-time inference. We characterise them instances of the classic parallel prefix-scan with model-specific operator, motivating and analysing the broader class of parallel scannable models (PSMs). In particular, we go beyond existing examples of PSMs by defining and empirically studying novel sequence model based on non-associative aggregators. Our experiments suggest that such model may have benefits in length generalization for some tasks, and opens avenues of exploring this design space light of specific applications. Overall, this provides an insightful unification of efficient sequence models, that cannot be found in any prior work.
ACKNOWLEDGEMENTS
This project was partially supported by NSF award CCF-21 12665, the MIT Quest for Intelligence, and the Alexander von Humboldt Foundation. S.G. is supported by the Math Works Engineering Fellowship.
10

---

Published as a conference paper at ICLR 2026
REPRODUCIBILITY STATEMENT
We provide proofs in Appendix B and relevant background for all theoretical results in Sec. 2 and Ap- pendix A. For experiments, we detail the training protocols in Sec. 4, and algorithm implementations in Sec. 3.4. All datasets are publicly available, and we follow established preprocessing procedures. We will release all code and scripts in public GitHub repository upon acceptance.
REFERENCES
Simran Arora, Sabri Eyuboglu, Aman Timalsina, Isys Johnson, Michael Poli, James Zou, Atri Rudra, and Christopher Ré. Zoology: Measuring and improving recall in efficient language models. Ini. Conf. on Learning Representations (ICLR), 2023.
Simran Arora, Sabri Eyuboglu, Michael Zhang, Aman Timalsina, Silas Alberti, Dylan Zinsley, James Zou, Atri Rudra, and Christopher Ré. Simple linear attention language models balance the tecall-throughput tradeoff. arXiv preprini 2rXiv:2402./8668 2024.
David Balduzzi and Muhammad Ghifary. Strongly-typed recurrent neural networks, In Proc, Int. Conf. on Machine Learning (ICML), New York City, NY, USA, June 2016.
David A. Mix Barrington. Bounded-width polynomial-size branching programs recognize exactly those languages in ncl. In Symposium on the Theory of Computing, 1986.
Maximilian Beck, Korbinian Pöppel, Markus Spanring, Andreas Auer, Oleksandra Prudnikova, Michael Kopp, Günter Klambauer, Johannes Brandstetter, and Sepp Hochreiter. xLSTM: Extended long short-term memory. In Proc. Advances in Neural Information Processing Systems (NeurIPS), Vancouver, Canada, December 2024.
Iz. Beltagy, Matthew Peters, and Arman Cohan. Longformer: The long-document transformer. Preprint N0MNv:20004.555000 2020.
Satwik Bhattamishra, Kabir Ahuja, and Navin Goyal. On the ability and limitations of transformers to recognize formal languages. In Proc. Conf. on Empirical Methods in Natural Language Processing (EMNLP), pp. 7096-7116 7116, Virtual only, November 2020a.
Satwik Bhattamishra, Arkil Patel, and Navin Goyal. On the computational power of transformers and its implications in sequence modeling. In Proc. Conf. on Computational Natural Language Learning (CoNLL). pp. 455-475. Virtual only, November 2020b.
Guy E Blelloch. Prefix sums and their applications. Technical report, School of Computer Science, Carnegie Mellon University Pittsburgh, PA, USA. 1990.
James Bradbury, Stephen Merity, Caiming Xiong, and Richard Socher. Quasi-recurrent neural networks. In Int. Conf. on Learning Representations (ICLR), Toulon, France, April 2017.
Kyunghyun Cho, Çağlar Gülçehre, Bart van Merriënboer, Dzmitry Bahdanau, Fethi Bougares Holger Schwenk, and Yoshua Bengio. Learning phrase representations using RNN encoder-decode for statistical machine translation. In Proc. Conf. on Empirical Methods in Natural Language Processing (EMNLP), pp. 1724-1734, Doha, Qatar, October 2014.
Tri Dao and Albert Gu. Transformers are SSMs: Generalized models and efficient algorithms through structured state space duality. In Proc. Int. Conf. on Machine Learning (ICML), Vienna, Austria, July 2024.
Yann N Dauphin, Angela Fan, Michael Auli, and David Grangier. Language modeling with gated convolutional networks. In Proc. Int. Conf. on Machine Learning (ICML) pp. 933-941, Sydney, Australia, August 2017.
Jeffrey I. Elman. Finding structure in time. Cognitive science, 14(2):179-213 1990.
Jonas Gehring, Michael Auli, David Grangier, Denis Yarats, and Yann N. Dauphin. Convolutional sequence to sequence learning. In Proc. Int. Conf. on Machine Learning (ICML), pp. 1243-1252, Sydney, Australia, August 2017.
11

---

Published as a conference paper at ICLR 2026
Riccardo Grazzi, Julien Siems, Jörg KH Franke, Arber Zela, Frank Hutter, and Massimiliano Pontil. Unlocking state-tracknn in linear mns through negative eigenvalues. In Int. Conf. on Learning Representations (ICLR), Vancouver, Canada, April 2025.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. In Conference on Language Modeling (COLM), 2024.
Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. In Int. Conf. on Learning Representations (ICLR), Virtual only, April 2022.
Han Guo, Songlin Yang, Tarushii Goel, Eric P Xing, Tri Dao, and Yoon Kim. Log-linear attention. Preprint arXiv:2506.04761, 2025.
Michael Hahn. Theoretical limitations of self-attention neural sequence models. Transactions of the Association for Computational Linguistics, 8:156-171, 1716-171,00 2020.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural compulation, 9(8): 1735-1780, 1997.
Kazuki Irie, Imanol Schlag, Róbert Csordás, and Jürgen Schmidhuber. Going beyond linear transform- ers with recurrent fast weight programmers. In Proc. Advances Neural Information Processing Systems (NeurIPS), Virtual only, December 2021.
Kazuki Irie, Róbert Csordás, and Jürgen Schmidhuber. Practical computational power of linear transformers and their recurrent and self-referential extensions. In Proc. Conf. on Empirical Methods in Natural Language Processing (EMNLP), Sentosa, Singapore, 2023.
Michael I Jordan. Attractor dynamics and parallelism in connectionist sequential machine. In Proc. Conf. of the Cognitive Science Society, pp. 331-546, Amherst, MA, USA, August 1986.
Nal Kalchbrenner, Lasse Espeholt, Karen Simonyan, Aaron van den Oord, Alex Graves, and Koray Kavukcuoglu. Neural machine translation in linear time. Preprint / 10099, 2016.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are RNNs: Fast autoregressive transformers with linear attention. In Proc. Int. Conf. Machine Learning (ICML), Virtual only, July 2020.
Najoung Kim and Sebastian Schuster. Entity tracking in language models. In Proc. Association for Computational Linguistics (ACI.), Toronto, Canada, July 2023.
Tao Lei, Yu Zhang, Sida I. Wang, Hui Dai, and Yoav Artzi. Simple recurrent units for highly parallelizable recurrence. In Proc. Conf. on Empirical Methods in Natural Language Processing (EMNLP), pp. 4470-4481, Brussels, Belgium, November 2018
Belinda Li,, Zifan Carl Guo, and Jacob Andreas. (How) do language models track state? Preprint arXiv:2503.02854. 2025.
Shuai Li, Wanqing Li, Chris Cook, Ce Zhu, and Yanbo Gao. Independently recurrent neural network (IndRNN): Building a longer and deeper RNN. In Proc. IEEE Conf. on Computer Vision and Pattern Recognition (CVPR), pp. 5457-5466 5466. Salt Lake City, UT, USA, June 2018.
Eric Martin and Chris Cundy. Parallelizing linear recurrent neural nets over sequence length. In Int.
Conf. on Learning Representations (ICLR), Vancouver, Canada, April 2018.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. In Int. Conf. on Learning Representations (ICLR), Toulon, France, April 2017.
William Merrill and Ashish Sabharwal. The parallelism tradeoff: Limitations of log-precision transformers. Transactions of the Association for Computational Linguistics (TACL., 11:531-545, 2023.
William Merrill, Gail Weiss. Yoav Goldberg, Roy Schwartz, Noah A. Smith, and Eran Yahav. A formal hierarchy of RNN architectures. In Proc. Association for Computational Linguistics (ACL), pp. 443.459, Virtual only, July 2020.
12

---

Published as a conference paper at ICLR 2026
William Merrill, Jackson Petty, and Ashish Sabharwal. The illusion of state in state-space models. In Proc. Int. Conf. on Machine Learning (ICML), Vienna, Austria, July 2024.
Sajad Movahedi, Felix Sarnthein, Nicola Muca Cirone, and Antonio Orvieto. Fixed-point RNNs: From diagonal to dense in a few iterations. Preprint arXiv:2503. 10799, 2025.
Michael C. Mozer. A focused backpropagation algorithm for temporal pattern recognition. Complex Systems, 3(4):349-38,, 1989.
Destiny Okpekpe and Antonio Orvieto. Revisiting associative recall modern recurrent models. In First Workshop Scalable Optimization for Efficient and Adaptive Foundation Models, 2025.
Aaron van den Oord, Sander Dieleman, Heiga Zen, Karen Simonyan, Oriol Vinyals, Alex Graves, Nal Kalchbrenner, Andrew Senior, and Koray Kavukcuoglu. WaveNet: A generative model for raw audio. Preprint arXiv: 100.0399,9 2016.
Adam Paszke et al. Pytorch: An imperative style, high-performance deep learning library. In Proc. Advances Neural Information Processing Systems (NeurIPS), pp. 8026-8037, Vancouver, Canada, December 2019.
Bo Peng, Ruichong Zhang. Daniel Goldstein, Eric Alcaide, Xingjian Du, Haowen Hou, Jiaju Lin, Jiaxing Liu, Janna Lu, William Merrill, et al. RWKV-7" goose" with expressive dynamic state evolution. Preprint arXiv:2503. 14456, 2025.
Hao Peng, Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah A Smith, and Lingpeng Kong. Random feature attention. In Int. Conf. Learning Representations (ICLR), Virtual only, 2021.
Zhen Qin, Songlin Yang, and Yiran Zhong. Hierarchically gated recurrent neural network for sequence modeling. In Proc. Advances in Neural Information Processing Systems (NeurIPS), New Orleans, LA, USA, December 2023.
Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. Linear Transformers are secretly fast weight programmers. In Proc. Int. Conf. on Machine Learning (ICML), Virtual only, July 2021.
Jürgen Schmidhuber. Learning to control weight memories: An alternative to dynamic recurrent networks. Neural Computation, 4(1):131-139, 1992.
Julien Siems, Timur Carstensen, Arber Zela, Frank Hutter, Massimiliano Pontil, and Riccardo Grazzi. Deltaproduct: Improving state-tracking in linear rans via householder products. Preprint arXiv:2502. 10297, 2025.
Lena Strobl, William Merrill, Gail Weiss, David Chiang, and Dana Angluin. What formal languages can transformers express? a survey. Transactions of the Association for Computational Linguistics (TACL), 12:543-561, 2024.
Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma. Yuqing Xia, Jilong Xue, Jianyong Wang, and Furu Wei. Retentive network: successor to transformer for large language models. Preprint arXiv:2307.08621, 2023.
Leslie G Valiant. Parallelism in comparison problems. SIAM Journal on Computing, 4(3):348-355, 1975.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Proc. Advances in Neural Information Processing Systems (NIPS), pp. 5998-599900 Long Beach, CA, USA, December 2017.
Songlin Yang, Bailin Wang, Yikang Shen, Rameswar Panda, and Yoon Kim. Gated linear attention transformers with hardware-efficient training. In Proc. Int. Conf. on Machine Learning (ICML), Vienna, Austria, July 2024a.
Songlin Yang, Bailin Wang, Yu Zhang, Yikang Shen, and Yoon Kim. Parallelizing linear transformers with the delta rule over sequence length. In Proc. Advances in Neural Information Processing Systems (NeurIPS). Vancouver, Canada, December 2024b.
13

---

Published as a conference paper at ICLR 2026
Songlin Yang, Jan Kautz, and Ali Hatamizadeh. Gated delta networks: Improving Mamba2 with delta rule. In Int. Conf. on Learning Representations (ICLR), Vancouver, Canada, April 2025.
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, al. Big bird: Transformers for longer sequences, Advances in neural information processing systems, 33: 13:7737317 2020.
14

---

Published as a conference paper at ICLR 2026
APPENDIX
Preliminaries and Depth Conventions
16
Additional Proofs
16
B.1 Modern RNN Layers Fit One Affine Scan
16
B.2 Examples of Prefix-Scannabl Sequence Models
17
Analytical Details of Chunking in Prefix Scannable Models
19
Computational complexity of PSMs
19
Beyond Affine State Recurrence, PSM's with General Aggregation: Further Details
20
15

---

Published as a conference paper at ICLR 2026
A PRELIMINARIES AND DEPTH CONVENTIONS
Depth conventions. Throughout this paper we restrict attention to causal sequence models whose training and inference graphs can be executed, in the Random-Acces, Machine (RAM) model with unbounded fan-out gates, at depth (measured by the longest path of synchronous operations)
Opolylogn), abbreviated (
Whether the hidden polylog factor is log (log n)² hinges the chosen primitive set-for instance, treating GEMM and so f ftmax as unit-time kernels versus expanding them into arithmetic gates. Classical fully recurrent networks such as LSTTss (Hochreiter & Schmidhuber, 1997) and GRUs (Cho et al., 2014), whose forward pass has depth (m) and therefore admits no sub-linear parallel schedule, fall outside this scope.
Our focus is on polynomial separations between the principal model families: standard Transformers, prefix-scannable models (as will define in Sec. 3), and linear recurrent RNNs. Unless logarithmic factor is essential to the argument, we suppress it with the tilde notation. To characterize our example models, we need to specify how their depth counted, i.c., the model of computation: How tall is their training circuit?
Transformer. A single self celf-attention head executes the composite map (Q. V) softmax QK /VT//VVVVV V. where d denotes the keylquery head dimension. In Random-Access Machine with unbounded fan-out gates, the pointwise linear projections w have depth but the n X n matrix multiply QK' and every row-wise softmax (vector sum + normalisation) require a parallel reduction of numbers. Using binary tree this costs (log n) depth." Hence, I-laaer Transformer has depth D(n)= = ( L. log n). If one treats the GEMM and softmax kernels as unit-time primitives, this is often reported as "constant depth," but strictly speaking it (1) (polylogarithmic).
Mamba, Gated Linear Attention, RWKV. The expensive step is parallel scan that produces the running state. Its depth is @logg.). and the pointwise gating around it adds O(1). Stacking Lagg such layers gives ((n) @ logn) n))
B ADDITIONAL PROOFS
B. MODERN RNN LAYERS FIT ONE AFFINE SCAN
To relate PSMs recent models, this section shows that broad family of recent fast-inference layers (Table 1) are all specializations of single affine state-update template, i.e. their state kernel can be expressed as an affine bilinear function. This enables SPD-(n, 1) complexity.
Definition B.1 (Affine recurrence). Let (M,+ +,0) 0) be an additive group and A: R x M - M fixed bilinear action of monoid (R, o, 1) on M. A layer said have an affine state update its hidden state obeys
8; E 8:-1 + St, 8-1 8-1=0 = 0,
(B.1)
where (Ee, fi) E x M are (learnable) functions of the current chunk x,. That is E, Eg (xx)) and ft = for (xx) ) for learnable functions Eg and fø.
The models in Table 1 all satisfy this affine state update template and all share the following aggregator, which is associative.
Lemma B.2. (One associative operator, Affine aggregator) Define for (E,, fi) E R M
(E2, (Ei, fi) (E20 E,, f2 Ea fi); e= (T,)). Then (R M, M, e) is monoid @ associative with identity e-and (E., f() @ where 8; is the state given by Eq. (B.1) and E, is an auxiliary variable.
(Eo. fo) (E., st),
"There suh-logarithmic circuits for exact matrix multiply (e.g. Valiant, 1975), but they very wide and rarely exploited MI. practice; log therefore matches realistic GPU/ /PU/ kernels.
16

---

Published as a conference paper at ICLR 2026
Proof. Straightforward verification using the action axioms; full details in Appendix B.
Once written in that form, their binary operator is associative, hence each layer is Trefix-Scannabee Model with SPD-(n, 1) complexity.
Theorem B.3. Every layer that satisfies Def. B.1 is Prefix-Scannabl Model with chunk size 1, encoder Enc, aggregator Agg, and inference module Inf defined as
2; Enc(C,) (Ei, fs)
for O,...,r--- with the aggregator Lemma B.2.
Agg(22, == (E2, f2) a fi)
We define (n-18 C.) to be the function that takes states and current token and oulputs predictions. Hence these models admits training work (n)), parallel depth (logg)) and online inference cost (11) time and (1)) memory per token: the layer is in SPD-(n, 1).
Proof. Apply the static Blelloch sean to the pairs (E,, f() using @ to obtain every prefix in O(n) work and (log n) depth. Lemma B.2 ensures the sean outputs the correct state 8,, which the inference head may consume chunk-wise. During streaming inference, the online left to right scan maintains the same prefixes with constant work and constant additional memory because + is associative. O
Table shows catalogue of affine layers. Note that the affine form absorbs normalisation variables common linear Transformers (e.g. running scalars/vectors 24 typically running sum of keys (Katharopoulos et al., 2020) related variables (Beck et al., 2024)) by enlarging the state vector and treating the auxiliary variable as just another coordinate updated affinely. The proof of Theorem B.3 requires change.
B.2 EXAMPLES PREFIX-SCANNABE SCANNABLE SEQUENCE MODELS
In the following, we present two families of models whose parallel circuits can be obtained as the computation of Blelloch parallel scan. fact, it suffices to show that for all family of architectures that are affine in their state, there exists associative operator that defines monoid over which the Blelloch parallel scan operates.
One type of prefix-scannable models are linear time invariant dynamical systems.
Definition B.4 (LTT Linear Dynamical System). A linear time invariant system defined by four matrices (A, B, C, D) dddd defining
81+1 As; + Bx,
(B.2)
yt Cs + Dx:
(B.3)
Here 80 = 0 is the initial state, and 8, is the state at time L E= a, E Rd is the input at time l.
Definition B.5 (Associative Operator for Affine State Monoid). For each timestep, let gi be an augmented pair g: == (E,, fi) = (A, Bæ;) where E, E idddd is a matrix and f. E IRd is vector. We define associative operator
E2, f2) @ fi.f1) (E Ezf1)
(B.4)
To demonstrate that a sequence model is Prefix Scannable, we must verify two properties. Firstly, that the operator O applied to all the 9; over all timesteps computes the state. Secondly that, @ is associative.
Lemma B.6. Let G2 be the augmented pair equal to @ applied to the sequence of augmented pairs g1...., gr. Then
G. = (Ei, (EI, (A',
Σ A t
(B.5)
k=0
Secondly for any inputs gi gi-9i-9 9k have (g. gj) @ gk 9: (g) @g)
17

---

Published as a conference paper at ICLR 2026
Proof. Proof by induction for the equality and straightforward computation for associativity.
We have base case.
EE f2) @ (E,, f1) (A², Bx² + ABx1) Apply definitions to see this is true for general t.
(B.6)
Proof of associativity. 91.92: 92, 9a we have (gs + +g2) + g1 = go
(g2 + g1)
ga (g2 gi) (A, Bwa) (A", Ba + ABrr) (A², BBz + ABwa Aa))
(B.7)
(A², Bx3 + ABx²) @ (A, Bx.) (A, Bos)) @ (A, Bx²) 1 (A, Bæ
(B.8)
gg g2) 1
(B.9)
Another type of prefix-scannable models are linear transformers and their gated variants.
Definition B.7. Gated Linear Attention (GLA) is defined with a states $1,.... 8T Raxd inputs 21, TT, gating function r Rd + I. keys k1, kr, kernel map & : IR + R
8: r(x)) 8 - (k)))
(B.10)
(B.11)
We observe that GLA has affine state recurrence.
Lemma B.8. Let E, E R be a scalar that can computed at. Let ft E IRpxd be a matrix that can be computed from Ty. Then the GLA state recurrence is affine function of the form
8: E,8t-1 + ft
(B.
(B.13)
In particular, let g. (Ec, fi) be augmented pair, and let @ be operator defined (E2, f2) @ (Ei, fi) E1, f2 Ezfi)
(B.14)
Then @ is associative, and 8, = g: 3 ... + g
Proof. Proof by induction for the equality and straightforward computation for associativity.
First we prove s; = 9: + + gt by induction. Consider the base case.
g2 + gl (r(#2), o ( @ (r(z)) (k;
(B.15)
= (r(#2) @ r(x)), 122 + r(x2) @ (kk )of)
(B.16)
(r(#2) 2 r(x1), 82)
(B.17)
Then assuming the identity holds timestep 1-1
gr + (r(xx:1) @ 2 r(a))) 82
(B.18)
(r(w.) @ 2 (a(a)))r((z))))) 2 8 + ((k:)oo)))
(B.19)
as desired.
Then we also check associativity.
gs (gz 91) gs @ (r(a2) Or(m1) + r(x) o ((k)ee
(B.20)
(r(z)) r(x2) r(x1), o(ks) uu + r(xs) o(k)) +r(x)) o r(z)) o okk
(B.21)
= (r(zz) O (2))) r(( o okk)) eg
g: (B.22)
(ga
(B.23)
18

---

Published as a conference paper at ICLR 2026
C ANALYTICAL DETAILS OF CHUNKING IN PREFIX SCANNABLE MODELS
Here we summarize the analytical details of PSMs (Def. 3.1). The model state of PSM after t chunks is the Blelloch prefix defined be 8, Agga Blelloch (a::a;), and the model outputs D) = Inf ( ( Ct). The same state sequence {st) can be produced online with O(log 1) memory by replacing the static scan with the binary-counter scan of Alg. 2. The corresponding (log t)- parallel depth training algorithm and Olog m)memoyy online decoding algorithm be found in Alg. and Alg. 4, respectively. Note that both parallel loops and the STATIC BLELLOCHSCAN have depth (log t), dominated by the static Blelloch scan, so the whole training pass admits efficient batch execution.
The model has the following properties:
Sequential-paralle duality. Alg. 3 and Alg. 4 produce identical state sequences {st} (Theorem 3.5), so a PSM trained with the static scan can evaluated online with logarithmic memory.
Model family. Choosing Aggo to be associative recovers known scan-friie friendly models as a strict subset of PSMs; non-associative choices (e.g. a Transformer block) enlarge the design space while retaining online decodability.
Complexities. Training: For sequences of length n, chunks of size c, we have O(n) work, (log(n/c)) depth. Online inference: O(c) amortised work per token and O log(n))) mem- ory after n/c chunks.
Further details about the computational complexity are detailed below in Appendix D.
D COMPUTATIONAL COMPLEXITY OF PSMs
Let
sequence length,
c chunk size (n (c-um_chunks),,
Lagg number of Transformer layers in Aggs-
Lint number of Transformer layers in Inf
d, and de hidden widths of the modules (held constant).
pass as single "time unit" and use the usual dense-attention cost (L for a length-/ Transformer block with L layers. Only the scaling with n, c, Lagg. Lint is retained; constant factors in d, d are suppressed.
Throughout we count one
Training (static Blelloch scan). The three parallel loops of Alg. 3 give
T Timin A(nnT cn.int),
Sman of en Lint en Lagg
(C1)
Depth is Oini + log(n/c) L.ge). because the static Blelloch scan dominates parallel runtime. Total nonparallel runtime is linear in sequence length O(enLagg + enLint). Additional factor of comes from é² dense attention for n//c chunks.
Online inference (binary-counter scan). Each new chunk incurs
1. one Inf, call = cost O c²), and
2. most log(n/c) calls Agg, per chunk amortised cost O
Hence, for the whole length-n stream, we make n calls to Inf and calls to Agg. The space at inference is store the ky-cache for the tokens in Inf and the log(n/c) chunks of e tokens in Agg
19

---

Published as a conference paper at ICLR 2026
Tinter of ne Linf + ne Lags)
Sinter O(c Linf clog(n/c) Logs).
(C2)
Per-token latency. Dividing (C2) by n gives
alLLint Linf + c Laag
work and (log n) space, confirming constant-timm amortised decoding under fixed c.
Remarks.
When = ((1) (token-wise chunks) both training and inference are linear in n with constant memory for Inf and logarithmic memory for Aggo-
For larger the quadratic self-attention of Inf over each chunk dominates work.
If Agg, is associative, we may swap the static and online without affecting costs; thus SSMs and gated linear attention inherit (CI) (C2) as special cases.
E BEYOND AFFINE STATE RECURRENCE, PSM's WITH GENERAL AGGREGATION: FURTHER DETAILS
The seen (a.k.a, parallel prefix) is fundamental to many parallel algorithms. When the binary operator is associative, the classic Blelloch scan (Blelloch, 1990) computes, in O(n) work and (Ioggn) depth, the same left-to-right prefix values that a sequential loop would produce. This section extends the view aon-associaeive operators such those expressible by softmax attention.
But, there is challenge with non-associativity; the numerical results of straightforward parallel and sequential versions would differ since parenthesisation differs, challenging duality. Parenthesisa- tion here means the explicit placement of parentheses that fixes which two elements are combined first when evaluating long chain of binary operations. For instance,
a Agg b Agg Agg d may be grouped as ((a Agg b) Agg c) Agg d or a Agg (b Agg AcAgg d)),
and when Agg is not associative the two expressions generally differ. The Blelloch algorithm removes this ambiguity by committing to a single, fixed parenthesisation: the full binary tree generated by its upsweep and downsweep. All variants we describe static and online-evaluaa evaluate exactly that same tree, guaranteeing identical results even for non-associative operators.
We first review the static tree formulation, then present online variant that realises exactly the same parenthesisation while using only O(log n)) memory. Throughout, let
Agg M + M, identity element e E M,
(AI)
be an arbitrary binary operator. No associativity assumption required unless stated.
First, we introduce the static Blelloch scan which is "parallel" training over sequence elements. Then we introduce the online binary counter scan which the "sequential" inference over sequence elements that computes prefixes with the same parenthesisation. The runtime required to run the static Blelloch sean is T(n) = (n)). whereas the amount of space required during the online binary counter scan is mnn) = (log(n)) Taken together this analysis gives us PSMs in SPD-(1, log(r)) Le linear compute during training and nearly linear space during inference.
Static Blelloch Scan (Alg. 1). Alg. I is agnostic to Eq. (A1), that is. it is valid for any operator. When Agg is not associative, however, the output for index no longer equals the sequential recurrence 8; Agg a,, $;-1). 1). Instead, it is the unique value obtained by applying Agg along the fixed binary- tree parenthesisation imposed by the algorithm. The next subsection shows how to realise the same parenthesisation online with logarithmic space.
Online Binary-Counter Scan (Alg. 2). The online variant processes the stream ag..... On left to right while maintaining binary counter of complete mini-trees. At time (0-indexed) the binary
20

---

Published as a conference paper at ICLR 2026
expansion of 1+1 reveals which block sizes 2 present. There is at most one mini-tre (its root value) per block size, hence at most (t++(]] roots memory. Inserting a new element is identical to adding to binary counter: each trailing 1 bit triggers merge with Agg and carry to the next bit. Aggregating the occupied roots from most- to least-significant bit (MSB - LSB) reproduces the value that the static Blelloch tree would hold after processing the same prefix.
We obtain the following correctness and complexity analysis.
Proposition E.1. After processing t+1 elements (t 0), every non-empty root [k] equals the aggregate of the most-recent 2 tokens Dt -1 Lt. and (1+1) is divisible by 2
Proof. By induction on f. The base case f=0 is immediate. For the inductive step, the carry chain merges two adjacent blocks of size 2 precisely when bit & flips from 1 to 0 in the binary counter. The merged value therefore covers the 2 most recent tokens and stored at position k+1, where divisibility holds. Untouched positions keep their invariant.
Theorem 3.5. Let p. be the value emitted at time t by Alg. 2. Then equals the exclusive prefix returned by the static Blelloch scan, regardless of whether Agg is associative.
Proof. By Proposition E. the occupied roots partition the first ++11 tokens into blocks whose sizes are decreasing powers of two when listed MSB - L.SB. This is exactly the leaf order the perfect binary tree used by the static algorithm. Each block's internal value constructed by the same sequence of merges as in that tree; aggregating the blocks left-to-right therefore reproduces the tree's evaluation order and thus its numeric result.
Corollary 3.6. After t+1 chunks Alg. 2 stores most [logz (t++)] root values; hence the worst-cass space usage is (log n))
Proof. The binary representation of ++1 contains most [log((++1]++1 bits, and there is at most one root per bit.
Work. Inserting a new element touches exactly the trailing 1-bits of t; the expected number such bits is 2, the amortised number of Agg calls per element is constant.
Together, Theorem 3.5 and Corollary 3.6 show that the online Sinary-counte sean is an opti- mal-space, streamable realisation of the Blelloch parenthesisation, extending prefix-scan tech- niques to non-associative operators without increasing asymptotic cost in time. This flexibility enables larger class of prefis-scammabe models: sequence models whose per-token state update is any binary operator that admits (log n) space O(1) time online evaluation via the mechanism above. We provide further analytical details of PSMs in Appendix C.
21


Reviews:
[{"review_id": "NgfHfKDafW", "reviewer": "Reviewer_8wg9", "summary": "This paper introduces an algorithmic framework, namely \"*Prefix-Scannable Models*\" (PSMs), to characterize and generalize sequence models that support both parallelizable training and efficient sequential inference. The authors term this property \"Sequential-Parallel Duality\" (SPD) and formalize it with the notation SPD(T(n), m(n)), capturing the training compute $T(n)$ and inference memory $m(n)$.\nThe authors first unify a large class of recent efficient models (e.g., Mamba, GLA, DeltaNet, etc.) under this definition. They show that these models are all instances of an associative prefix scan based on a shared \"affine state update\". Because their aggregation operator is associative, they can achieve SPD(n, 1) compared to the SPD($n^2, n$) of Transformers. Afterwards, to generalize the framework to non-associative aggregation operators (such as softmax attention), the classic Blelloch parallel prefix scan is used for training, which fixes a specific binary-tree parenthesization. The crucial insight is proposing an online \"binary counter\" algorithm for inference (Alg. 2) that exactly reproduces this same parenthesization, guaranteeing training-inference consistency. This generalization creates a new class of models, including the proposed *Transformer-PSM*, that achieve SPD(n, log n) with $O(1)$ amortized inference time. Experiments show that Transformer-PSM achieves very good length generalization on state-tracking tasks, and at the same time matching the performance of a full Transformer on associative recall and language modelling.", "questions": "1. How does Transformer-PSM compare to recent hybrids which alternate Transformer and SSM layers? Do such models fall in the SPD framework too? What are the hypothesized trade-offs in terms of expressivity and parallelization between these two different ways of combining Transformers and efficient recurrences?\n\n2. You successfully parallelize a non-linear operator by restricting its dependencies to a fixed binary-tree structure. Concurrent work like DEER and ELK provides an algorithmic alternative for parallelizing general non-linear RNNs via fixed-point iteration. Could you compare these two fundamental approaches?\n\n3. I would be interested to see more results on other state-tracking tasks such as the ones used in Grazzi et al. 2025. Can the authors provide some results on those tasks, e.g. parity, modular arithmetic w/ and w/o brackets? Additionally, can you also evaluate Mamba2 and DeltaProduct with negative eigenvalues in the state-transition matrix in your experiments?\n\n**-- References --**\n\nSiems et al. Deltaproduct: Improving state-tracking in linear rnns via householder products. In NeurIPS 2025.\n\nGrazzi et al. Unlocking state-tracking in linear rnns through negative eigenvalues. In ICLR 2025.\n\nGus et al. Log-linear attention. Preprint arXiv:2506.04761, 2025.\n\nAI21 Labs. Jamba: A Hybrid Transformer-Mamba Language Model. arXiv 2024\n\nGlorioso et al. Zamba: A Compact 7B SSM Hybrid Model. arXiv, 2024\n\nLim et al. Parallelizing non-linear sequential models over the sequence length. In ICLR 2024\n\nGonzalez et al. Towards Scalable and Stable Parallelization of Nonlinear RNNs. In NeurIPS 2024", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "1. **Methodological and technical contributions.** I think that the clean, algorithmic lens it provides for understanding and unifying the recent \"zoo\" of efficient sequence models is a significant conceptual contribution to the field. Formalizing the SPD(T, m) property, unifying Linear RNNs, SSMs and Transformers, and the transition from associative to non-associative operators is a non-trivial and useful generalization. Technically, using the online binary counter (Alg. 2) to exactly match the Blelloch scan's fixed parenthesization is another instance of applying a classic algorithm to solve a modern deep learning problem. \n\n2. **Empirical evaluations**. The proposed Transformer-PSM model shows very good length generalization result on the $S_5$ state-tracking task (Figure 3), up to sequences >6x longer than the longest training sequence. Furthermore, the MQAR result in Figure 4 show that the Tranformer-PSM can match the recall power of a full Transformer, while the WikiText-103 experiment (Figure 5) shows the trade-off between chunk size and perplexity on language modelling, validating the model's ability to achieve RNN-like efficiency and Transformer-like expressivity on these tasks.\n\n3. **Clarity and Structure**: The paper is well-written, and the use of formal definitions (Def 2.5 for SPD, Def 3.1 for PSM) and algorithms (Alg 1, 2, 3, 4) make the core ideas succinct and easy to follow.", "weaknesses": "1. **Limited scale**: The experiments are on small-to-medium datasets (WikiText-103) with relatively small models. While the synthetic results are interesting, it remains to be verified how the architectural benefits and trade-offs (especially $O(\\log n)$ memory) of Transformer-PSM scale to state-of-the-art compute with billions of parameters.\n\n2. **Missing baselines and other related work**: While the theoretical unification is broad, the baselines used in the empirical comparisons are narrow, focusing almost exclusively on GPT-2 and Mamba. I think the paper would greatly benefit if the authors would also compare to several highly relevant lines of concurrent work such as:\n   - *Transformer-SSM Hybrids*: Models like Jamba, Zamba, etc., which interleave Transformer and SSM layers.\n   - *Log-Linear Attention*: The paper cites \"Log-linear attention\" (Guo et al., 2025) as concurrent work that \"fits this view\" but provides no empirical comparison, even though it appears to be a direct competitor in the SPD(n, log n) class.\n   - *Parallel Non-linear RNNs*: Other approaches to parallelizing non-linear recurrences, such as those based on fixed-point iteration (e.g., DEER and ELK), offer a completely different algorithmic approach to the same problem (parallelizing non-linearity) and should be discussed or benchmarked.\n   - *Linear RNNs with enhanced expressivity*: Methods such as DeltaProduct, that have been shown to work great on state-tracking tasks too.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "rLiQZl6UBx", "reviewer": "Reviewer_LR1S", "summary": "This paper focuses on addressing the core tradeoff between parallelizable training and efficient sequential inference in neural sequence models. It formalizes the Sequential-Parallel Duality (SPD), a property where models support near-constant-depth parallel training and linear-time/constant-space sequential inference, and introduces a generalized model class called Prefix-Scannable Models (PSMs) to unify and extend existing efficient architectures.", "questions": "1. The paper uses the Blelloch scan’s fixed binary-tree parenthesisation for non-associative aggregators (e.g., softmax attention). Is there a theoretical or empirical justification for this specific parenthesisation over alternative groupings (e.g., left-to-right sequential, random binary trees)? Could different parenthesisations impact task performance or complexity?\n\n2. For the Multi Query Associative Recall (MQAR) task, Transformer-PSM with chunk size 64 achieves perfect accuracy, but chunk size 32 degrades performance on long sequences (512 tokens). What causes this sensitivity to chunk size, and is there a principled way to select chunk size for different tasks without grid search?\n\n3. The paper compares Transformer-PSM to GPT-2, Mamba, and Sliding Window Transformer (SWT). Why were other efficient models (e.g., RetNet, GLA, xLSTM’s mLSTM) not included in the comparisons, especially for language modeling or associative recall? How would Transformer-PSM perform against these models on the same benchmarks?\n\n4. The paper claims O(log n) inference memory for PSMs, but chunk size (c) interacts with this scaling (Appendix C). For large chunk sizes (e.g., c=256), does the memory scaling still hold, or does the chunk-local KV cache in the Inf module dominate memory usage?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The paper demonstrates originality through three key contributions. First, it introduces a unified theoretical framework (Prefix-Scannable Models, PSMs) that formalizes sequential-parallel duality (SPD), a core property of modern efficient sequence models (e.g., Mamba, GLA, RetNet). This unification resolves the fragmentation of prior work by showing that diverse architectures are special cases of PSMs with associative aggregation operators. Second, it generalizes PSMs to non-associative aggregation functions (e.g., softmax attention), expanding the model design space beyond existing affine/linear constraints while retaining efficient inference. Third, the instantiation of Transformer-PSM is a creative combination: it merges Transformer-style self-attention (for expressive power) with PSM’s linear-time inference (for efficiency), addressing a long-standing tradeoff between performance and scalability.\n\n2. It provides strict formal definitions (e.g., SPD, state kernel, PSM) and rigorous proofs (e.g., Proposition 3.2, Theorem 3.5) for complexity guarantees (O(n) training compute, O(log n) inference memory). Appendixes supplement detailed derivations (e.g., affine recurrence proofs, computational complexity analysis) to support core claims. Empirically, the experiment design is comprehensive and controlled: it targets three canonical tasks (state tracking S₅, multi-query associative recall, WikiText-103 language modeling) and compares against strong baselines (GPT-2, Mamba, Sliding Window Transformer). Results are interpretable (e.g., Fig. 3 shows superior length generalization; Fig. 6 validates constant inference latency) and reproducible (detailed training protocols, public code plan).\n\n3. Theoretically, it deepens the connection between parallel prefix-scan algorithms and neural sequence modeling, providing a general framework for analyzing and designing efficient models—this will guide future research on SPD-compatible architectures (e.g., log-linear attention, as noted in the discussion). Practically, it addresses critical limitations of existing models: Transformer-PSM matches Transformer’s expressive power (e.g., perfect associative recall) while achieving Mamba-like linear inference efficiency (e.g., constant per-token latency for long sequences). This balance is transformative for real-world applications (e.g., long-document processing, streaming NLP) where both training parallelism and inference scalability are essential. Additionally, the paper’s length generalization results (e.g., S₅ task scaling to 160+ tokens) open new avenues for modeling algorithmic tasks that require precise state tracking.", "weaknesses": "1. The paper’s core contribution of generalizing PSMs to non-associative aggregation operators is underdeveloped, as it only instantiates `Aggθ` with Transformer-style softmax attention (Transformer-PSM). The design space of non-associative operators (e.g., additive attention variants, kernelized attention, or task-specific non-associative functions) remains unexplored. This raises critical questions: Are PSM’s advantages (e.g., length generalization on `S5` tasks) inherent to the prefix-scan framework, or specific to softmax attention? Without testing alternative non-associative `Aggθ` implementations, the paper fails to validate that PSMs truly expand the efficient sequence model design space beyond softmax-based variants.\n\n2. The paper demonstrates that chunk size (`c`) significantly impacts performance (e.g., Transformer-PSM with `c=64` achieves perfect accuracy on MQAR, while `c=32` degrades on long sequences; Fig. 4) and perplexity (Fig. 5). The paper does not quantify the tradeoff between chunk size, computation cost (e.g., `O(c²)` for chunk-local self-attention), and memory usage (e.g., `O(c log(n/c))`). How does `c` scale with sequence length (e.g., optimal `c` for 1k vs. 10k token sequences)?\n\n\n3. Missing Ablation Studies on Core PSM Components. \nThe PSM framework is defined by three modules: `Enc` (chunk encoder), `Aggθ` (aggregator), and `Infφ` (inference head). However, the paper does not perform component-wise ablation studies to isolate which modules drive PSM’s performance advantages. Does the simple embedding `Enc` (used in Transformer-PSM) limit performance, or is a more complex encoder (e.g., convolutional, hierarchical) unnecessary?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ivdxNUhVrN", "reviewer": "Reviewer_Koqn", "summary": "The paper formalizes Prefix-Scannable Models (PSMs), a class of sequence models whose training graph is a Blelloch prefix scan over chunk encodings with a chunk-local predictor head. This framework yields SPD-$(n, \\\\log n)$ complexity, entailing $\\\\Theta(n)$ work with logarithmic parallel depth during training. For inference, an online binary-counter algorithm reproduces the exact same computation with $O(1)$ amortized compute per token and $O(c \\\\log (n/c))$ memory (which is $O(\\\\log n)$ when $c$ is fixed). An implementation, Transformer-PSM, demonstrates on WikiText-103 that perplexity improves as $c$ increases, while per-token latency remains bounded in contrast to a full-context Transformer. The theory unifies associative scan-based models as a special case and correctly extends the framework to non-associative aggregation operators.", "questions": "* On WikiText-103, Fig. 5 reports perplexity vs. chunk size $c$, and Fig. 6 shows per-token time vs. token index for one $c$ (on WikiText-2). To make the accuracy–efficiency trade-off actionable with minimal overhead, please add a small table over $c \\\\in \\\\{32, 64, 128, 256\\\\}$ reporting (i) the existing WikiText-103 perplexity, (ii) mean per-token inference latency on one long input (e.g., $L=40{,}000$ on WikiText-2; discard the first $c$ tokens for warm-up, then average per-token time over $t \\\\in [c, L]$), and (iii) peak inference memory for that run; optionally include inference throughput (tokens/s).\n* Could the authors include system-level training metrics (tokens per second, total wall-clock time, peak memory) for the Transformer-PSM and baselines on at least one of the presented tasks?\n* For the MQAR experiment, please consider adding a sliding-window Transformer baseline with a matched compute budget. A summary table of all baseline configurations would also improve reproducibility.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* Clear formalization of PSMs and their sequential-parallel duality, including precise analysis of training and inference complexity.\n* A key theoretical result is the correctness of the online binary-counter evaluation relative to the static training-time tree, which resolves the parenthesization ambiguity for non-associative operators.\n* Experimental results show that perplexity improves with a larger chunk size $c$, while inference latency remains weakly dependent on context length, unlike a standard Transformer.", "weaknesses": "* The evaluation of efficiency lacks system-level measurements. Practical metrics such as training throughput (e.g., tokens/sec), wall-clock time, and peak memory usage are not reported, which limits the assessment of the model’s practical viability.\n* The trade-off between accuracy and latency is not fully explored. For large $c$, latency is dominated by the chunk-local predictor. The paper notes the $O(c)$ amortized compute and the quadratic ($O(c^2)$) attention cost within the Inf module, but does not map the accuracy–latency frontier across different values of $c$. This analysis is critical for deployment.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "0JM4ZnrY6m", "reviewer": "Reviewer_k4pr", "summary": "This paper proposes Prefix-Scannable Models (PSMs), a unified framework for sequence models that enables both parallel training and efficient sequential inference. By introducing a fixed parenthesization (bracket) structure in the prefix-scan algorithm, PSM extends beyond associative operations, allowing non-associative functions like softmax attention to be trained and decoded efficiently.", "questions": "- How robust is the fixed parenthesization choice? Would learning or adapting the tree structure dynamically lead to better trade-offs between accuracy and efficiency?\n\n- Could this framework be extended to bidirectional or hierarchical scanning, potentially improving global context modeling?\n\n- Could this theory inspire a hardware-aware implementation, where the prefix-scan tree maps directly onto GPU or TPU parallel primitives?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "– PSM generalizes a wide range of recent models under a single mathematical formulation. It reframes these architectures as special cases of prefix-scannable systems, giving conceptual clarity to an area previously defined by ad hoc designs.\n\n– By fixing a binary-tree structure for prefix computation, the authors bypass the associativity constraint that limited prior parallel scan models. This makes it possible to handle attention-like, non-linear interactions while retaining log-depth parallelism.\n\n– Experiments demonstrate that Transformer-PSM achieves near-Transformer accuracy (only slightly higher perplexity) while matching Mamba’s inference speed. This combination of expressivity and efficiency is rare and valuable for long-sequence modeling.", "weaknesses": "– Results are limited to moderate-scale models on synthetic and mid-size text datasets; it’s unclear whether the same advantages hold at billion-parameter scale or in real industrial workloads.\n\n– The fixed-tree scanning and binary-counter inference mechanisms, while elegant conceptually, may be hard to integrate into existing frameworks without specialized kernels. I'd like to hear the author's explanation about this aspect.\n\n– Model quality strongly depends on the chosen chunk length, adding a new hyperparameter that requires tuning for each task.\n\n– The paper focuses heavily on theory and performance metrics but provides little intuition or visualization about how information actually propagates through the prefix structure. A detailed explanation or visualization would be helpful. Although Figures 2 and 3 help illustrate the overall architecture and performance trends, the paper still offers limited qualitative insight into the internal behavior of PSMs. No visual analysis of prefix states, attention dynamics, or information propagation is provided.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)