Rational Transductors
SAI paper + code review · Referee report
Summary
The paper proposes Rational Transductors, a hybrid sequence architecture that augments a Transformer backbone with a sidecar Weighted Finite Automaton (WFA) whose linear matrix-valued recurrence is injected into every Transformer layer (Deep Rational Injection). The main conceptual move is to keep the recurrence strictly input-driven and linear so that it (i) admits an parallel associative scan, (ii) sits algebraically inside the theory of rational power series, and (iii) is amenable to a spectral parameterization (Scaled Cayley) that mechanistically avoids vanishing/exploding gradients. The paper argues this closes the "Regular Gap" between attention (limited to AC/TC) and finite-automata state tracking (regular languages, cyclic groups, Boolean formula evaluation), while — unlike SSMs stacked in the residual stream — preserving strict parallelism because the recurrence never depends on non-linear layer outputs. Theoretical contributions include a decomposition characterization (), circuit-complexity placement in PNC[1], a Krohn-Rhodes-style representational-completeness argument, and learning-theoretic guarantees (random-feature universality, contractive length-independent generalization, Hankel-Rademacher complexity). Empirically the paper claims perfect length generalization on Modulo-5 counting (up to training length), 1000-digit long-integer addition, and near-machine-precision Base-2 evaluation, alongside a latency benchmark supporting the claim.
Overall the framing is elegant and the conceptual bridge between attention (aperiodic / FO) and WFA (group / MSO) is genuinely valuable; the Krohn-Rhodes lens is unusually clean for a deep-learning paper. However, several theoretical statements are stronger than what the proofs deliver (Thm 22 says "strictly subsumes" but proves equality; Thm 14's "irreducibility" is really an upper bound on a specific construction; the "virtual automaton" from Sec. C.4 conflicts with Thm 14; the " solely" reading of Thm 32 leaves an undefined term). Empirically, all "gap-solving" experiments run in the strict-orthogonal regime, so Thm 7's fading-memory bound is never actually exercised; no SSM/Mamba baseline is included despite the theory placing them in the same complexity class; and there is no large-scale-backbone evidence for the "no constraints on the Transformer" claim.
Strengths
- Conceptual contribution. Framing the sidecar as a WFA and lining up its algebraic operations (direct sum parallel heads, tensor product attention interactions, composition depth) is a clean and productive lens. Anchoring the recurrence in rational power series and Krohn-Rhodes gives the architecture a much sharper theoretical identity than most "hybrid Transformer/RNN" papers.
- Parallelism preserved by design. The commitment to an input-driven, purely linear recurrence — separated from Transformer non-linearities — is a principled architectural choice that genuinely delivers parallel training, unlike stacked SSMs which pay . This tradeoff is stated explicitly (Sec. B.2), not hidden.
- Circuit-complexity placement. Situating the model in PNC[1] between the AC[0]/TC[0] Transformer barrier and the P-complete RNN sequential bottleneck is a useful and defensible characterization; Corollary 25 (maximality of associative linear recurrence extensions) is a nice framing of "what more can attention be augmented with without giving up parallel scan?"
- Spectral parameterization. The Unified Scaled Cayley parameterization simultaneously supports the conservation (, orthogonal) and decay () regimes, with mechanistic gradient-norm preservation rather than heuristic clipping. Combined with near-identity initialization, this is a genuinely nice piece of design.
- Length-generalization signal. On Modulo-5, generalizing from to with essentially perfect accuracy — while a same-sized Transformer collapses to chance — is a striking demonstration, even accepting the caveats below.
Weaknesses
- Theorem 22 statement contradicts its proof. The theorem claims "strictly subsumes" the W-MSO-definable class, but the proof opens with "We establish the equality by mutual inclusion" and demonstrates both directions. Under the theorem's hard-attention hypothesis the correct claim is equality (mirroring Droste-Gastin). The strictness claim needs a soft-attention hypothesis, which is invoked separately in the subsequent remark.
- Non-Linear Irreducibility (Thm 14) proves the wrong direction. The theorem is stated as a lower bound ( for any single linear WFA), but the proof only exhibits an upper bound on a specific tensor-product construction. A rigorous irreducibility argument needs a Hankel-rank lower bound. Compounding this, the dimension thresholds in a single proof jump between , , and ; at the inequality holds with equality, not strictly.
- "Virtual Tensorization" conflicts with Theorem 14. Sec. C.4 argues the sidecar plus attention "approximates the decision boundaries of a much larger automaton of dimension " without serialization. Attention produces a bilinear score — a scalar readout linear in the tensor product — not a recurrent state that evolves under . Either this static-feature reading is enough (in which case Thm 14's "irreducibility" is not something the sidecar suffers from), or the sidecar cannot in fact reach the tensor state (in which case the Sec. C.4 payoff is overstated). The two claims cannot both be maximally strong.
- TC[0]/Parity conflation. Parity is exactly in TC[0] (indeed AC[0][2]). The claim that TC[0] models "can theoretically approximate" parity understates a canonical fact, and in Thm 16's proof the TC[0] upper bound is invoked as if it forbade recognition. The real barrier is learnability and uniform length generalization (which C-RASP addresses), not TC[0] expressibility.
- Krohn-Rhodes proof skips the feedback wiring. The proof correctly notes permutation matrices realize simple groups and FFNs realize arbitrary Boolean functions, but does not show how an FFN's output becomes the input-dependent transition matrix of the next rational head. Given the sidecar's core commitment to input-driven transitions, this coupling is exactly the mechanism that needs justification, and it is skipped; the resulting construction is also not the parallel sidecar but a non-linear stacked variant.
- Random-features theory (D.1) develops but does not empirically test. Theorems 31-32, Prop 33, and Cor 34 build a substantial random-features / rational-kernel story that "justifies" near-identity initialization, but no experiment freezes the random transitions and trains only a linear readout. The Compactness Gap ( learned vs. random) is never empirically demonstrated. Also, Thm 31 requires , exponential in horizon — the section header advertises "Efficiency", but only exponential-width universality is proved.
- All gap-solving experiments run in ; Thm 7 never bites. Thm 7's Time-Invariant Error Bounding requires strict contraction ; all reported experiments use the orthogonal Conservation Regime. The paper's claimed length-generalization mechanism (via Thm 7) is therefore not what the experiments validate — they validate algebraic exactness of orthogonal parameterization. The fading-memory regime that Thm 7 covers is never probed.
- Missing SSM/Mamba baseline. The related-work section explicitly places S4, Mamba, DeltaNet, and Kimi Linear as sharing time-varying linear recurrence and PNC[1] capacity. The "minimal theoretical proxy" defense justifies not benchmarking against Mamba, but the reader cannot then tell whether "solving the Regular Gap" is unique to Rational Transductors or shared by any modern linear-recurrence architecture with a suitable unitary initialization.
- Empirical scope narrow relative to the "no constraints on backbone" claim. The paper explicitly claims the rational head "is a modular addition that can be integrated into any existing Transformer", yet every experiment uses a 2-layer, backbone on synthetic tasks. Whether Deep Rational Injection remains stable and beneficial when added at scale is not tested.
- Deep Rational Injection vs. simple concatenation is not ablated. The Deep Injection design is defended by a narrative signal-attenuation argument (Sec. 2) but never compared head-to-head against the naive concatenate-at-input baseline.
- Context-free "bridging" clashes with the paper's own bounded-precision framing. Sec. C.2's claim that exact-arithmetic evaluation "theoretically bridges" the Context-Free Gap relies on unbounded real arithmetic (values growing like ), which contradicts the bounded-precision regime underpinning every other complexity claim (and Prop 19 explicitly denies context-free recognition).
- Theorem 13 sample-complexity reduction asserted, not derived. The parameter-count ratio is correct, but the leap to "sample complexity … by a factor of " skips the connection to the paper's own Hankel-Rademacher machinery (Thm 40), which is not stated in raw parameter count.
- Notation / class-name slips. The stacked analysis switches from (Thms 13-14) to layers (footnote 2) without flagging the identification. The class "PNC[1]" is defined loosely and used to argue tightness against a stated NC[1] = PNC[1] separation; standard characterizations of iterated integer matrix product are typically stated in terms of #NC[1] / GapNC[1] / BWBP, and the tightness argument only lands with a precise anchor.
- Random-features / RKHS bound (Thm 32) has an undefined term. The chaining argument in Thm 32 introduces a term "proportional to ", but is never defined. The subsequent Interpretation says the error is "governed solely by" and , which is inconsistent with an undefined length-dependent extra. Also the Lipschitz justification (step 2) argues with respect to initialization parameters, but step 3's covering argument needs Lipschitzness in the input under weighted Hamming; the two notions are conflated.
- Local slips. "Spectral radius" is used where the analysis requires "spectral norm" (Sec. 5.2 remark). The Thm 35 formula shows a constant additive term, but the surrounding prose promises "grows at most linearly" — the two are only consistent in the case. The Base-Case / RoPE lemma proof recovers the additive sinusoidal encoding, not the multiplicative rotary encoding as claimed. The Parity construction uses , outside the default SO Cayley parameterization; the paper acknowledges this in a remark but never demonstrates a Parity solve within the default regime.
Reproducibility & code
- Modulo-5 data generation and Transformer baseline underspecified. The Modulo-5 task's alphabet, target definition (running count vs. final residue), sampling scheme, and dataset size are not documented. The Transformer baseline is only ", 2 layers"; heads (Table 2's "4 heads" appears to describe the Rational Transductor, not the baseline), positional-encoding type, dropout, and weight decay are unstated. A reimplementer can plausibly recover the qualitative gap but not the exact intermediate-length numbers.
- Long-Integer Addition training window inconsistent. Table 2 states for the addition curriculum, while Section 6.4 (and Figure 6a's caption) refer to a training window of . These are not reconcilable as printed.
- Base-2 LSTM baseline underspecified. Only "standard PyTorch initialization" is given for the LSTM baseline; hidden dim, layer count, and gating configuration are not stated. Since the reported MSE floor of for LSTM/Transformer is used to argue an -order-of-magnitude gap, this omission is material. Training set construction over (uniform sampling? number of examples? train/eval split?) is also not documented.
- Latency-benchmark baselines not pinned. The Sequential RNN's state dimension, precision, and framework are unstated; FlashAttention version is unpinned; warmup and CUDA-sync protocols are not described; trial count varies over a range ("20–100"). Small variations in these can shift the reported crossover between Rational Transductor and Sequential RNN meaningfully. The underlying latency-vs-length numeric arrays for Figure 5 do not appear to be released.
- Length-generalization hero result depends on a footnoted Float64 workaround. Table 2 lists Float32 precision for the Length Generalization experiment, but the accuracy at is only achievable under Float64 (per footnote 1). Naive re-execution following Table 2 will not reach the reported number.
- Seed values and per-seed results not disclosed. Appendix E.4 reports 5/5 successes with but does not disclose the seeds used or per-seed step-to-convergence numbers. Given the "converged within 500-1000 steps in every trial" claim is a central stability argument, publishing per-seed numbers would substantially strengthen it.
Recommended Changes
Essential
- Fix Theorem 22's statement. Replace "strictly subsumes" with "is exactly characterized by" (or "equals") under the hard-attention hypothesis; move the strict-subsumption claim to a separate corollary with an explicit soft-attention hypothesis.
- Repair Theorem 14. Either (i) prove a Hankel-rank lower bound on the minimal single-linear-WFA dimension realizing the non-linear cascade's transduction, or (ii) restate the theorem as an upper bound on the specific tensor-lift construction. Also, reconcile the / / conditions in the proof.
- Reconcile Virtual Tensorization with Theorem 14. Either weaken the Sec. C.4 payoff to "compute bilinear features over pairs of states" (a static readout, not a -state automaton), or explain how the sidecar's attention layer accesses tensor-product dynamics that Thm 14 says are unreachable to any linear single-head system.
- Fix TC[0]/Parity framing. State explicitly that Parity is in TC[0], make the negative claim about Transformers about uniform learnability / length generalization (with a proper appeal to Merrill et al. and C-RASP), and remove "approximate" wording that suggests TC[0] does not exactly represent Parity.
- Fill the Krohn-Rhodes proof gap. Show explicitly how an FFN's output becomes the input-dependent transition matrix of the next rational head. Flag that this construction is the non-linear stacked variant and does not preserve the sidecar's parallel scan.
- Add an SSM/Mamba baseline on Modulo-5 counting and Long-Integer Addition. If the paper's positioning as "minimal theoretical proxy" is retained, the empirical section still needs at least one linear-recurrence baseline (Mamba with unitary init or a small S4) to show the Regular-Gap results are attributable to the paper's specific architectural choices rather than to any modern linear recurrence.
- Add a experiment that tests Thm 7. All current experiments run in ; the fading-memory error bound (Thm 7) is never validated. A "forget-and-check" task with a decay-regime rational head would tie theorem to evidence.
- Publish data-generation, baselines, and per-seed numbers in the appendix: (i) Modulo-5 alphabet / target / sampling scheme / dataset size; (ii) full Transformer baseline (heads, PE, dropout, init); (iii) LSTM baseline (hidden dim, layers, gating); (iv) latency-benchmark baselines (RNN config, FlashAttention version, warmup, sync); (v) seed values and per-seed step-to-convergence; (vi) resolve the vs. inconsistency for Long-Integer Addition.
- Correct the precision claim. If the length-generalization number requires Float64, update Table 2 to reflect this and move the caveat into the main text — not a footnote.
- Ablate Deep Rational Injection vs. concatenation. Report a head-to-head comparison on at least one task; the signal-attenuation defense in Sec. 2 currently rests on narrative.
Suggested
- Add a small language-modeling experiment (or explicitly scope down the "no constraints on the Transformer backbone" claim). Even a modest character-LM ablation would go a long way toward supporting the modularity claim.
- Test the random-features regime. Freeze the transitions at initialization and train only the linear readout on Parity/Modulo-5 to empirically demonstrate the Compactness Gap of Prop 33.
- Clarify PNC[1]. Anchor the class definition to a standard reference (#NC[1] / GapNC[1] / BWBP) and cite the specific completeness result for iterated integer matrix product; the tightness argument depends on this being a standard, not bespoke, class.
- Define in Theorem 32 (metric diameter, log-covering number, or otherwise) and either carry it into the "Interpretation" paragraph or explain why it drops out.
- Fix local slips. Replace "spectral radius" with "spectral norm" in the Thm 35 remark; split the Thm 35 bound into two regimes (: constant; : linear ); relabel the RoPE lemma as reproducing the additive sinusoidal encoding, or extend it to the true rotary construction; reconcile the / notation switch in the stacked analysis; qualify the context-free "bridging" claim by explicitly noting it requires unbounded real precision and does not survive the paper's own bounded-precision regime.
- Show a Parity solve within the default SO parameterization using or the augmented embedding suggested in the remark, so that the headline "" is achievable in the paper's own default parameterization.
- Publish underlying numeric arrays for Figures 4, 5, and 6 in a table appendix so plots can be re-drawn without re-running the experiments.