Expressivity-Efficiency Tradeoffs for Hybrid Sequence Models

Summary

This paper takes a serious first step toward a formal understanding of why hybrid Transformer–SSM sequence models can outperform either constituent architecture alone. The authors define a family of "structured two-layer function composition" (S2FC) tasks whose targets factor as $F(u(\vec{\mathbf{x}}), v(\vec{\mathbf{x}}))$, where $u$ is a moderately long local context and $v$ is a control variable extracted from the wider context. On this family they prove two complementary lower bounds: under an injectivity condition, any $k$-layer SSM must satisfy $\sum_i \log|S_i| \ge \Omega(m\log|V| - q\log|Y|)$ (Theorem 3.3), and under a global-sensitivity condition, any $k$-layer sliding-window Transformer must have total window $\sum_i W_i \ge R$ (Theorem 3.7). They then construct explicit shallow hybrids for selective copying (Theorem 4.3) and associative recall with decoding (Theorem 4.6) whose parameter count and working memory scale polylogarithmically in $|V|$ and $L$, provably escaping both lower bounds simultaneously. Empirically, they show that learned SSM→TF hybrids match or beat pure baselines with roughly $6\times$ fewer parameters on selective copying and MKAR, that hybrids uniquely reach $>50\%$ accuracy on associative recall with decoding at the tested scales, and that hybrids generalise better under length and bit-proportion shifts.

The conceptual contribution — a hardness–construction dichotomy that formally frames "SSM as encoder feeding a small-window Transformer" as the source of the hybrid advantage — is genuinely useful and, to my knowledge, novel in this precise form. That said, the current draft has multiple presentational and technical inconsistencies (probability thresholds that do not match between theorem statement and proof, a $\sum$-vs.-$\prod$ error in Lemma 3.5, a wrong-task copy-paste in the statement of Theorem 4.6, and an overloaded selective-copying definition). The empirical side is confined to synthetic tasks at $L\le 100$, uses only sliding-window attention baselines, and includes an unresolved observation (Transformer accuracy dropping as its window grows) that partially undercuts the expressivity narrative. When the release was executed under a constrained environment only the theory-validating construction notebooks and a heavily downscaled slice of the selective-copy sweep completed; the qualitative signature (SSM→TF standing out) is directionally reproduced, but almost all headline empirical claims — the decode-recall, MKAR, needle, length-generalisation, and OOD tables — could not be exercised and rest on the released code alone.

Strengths

• Conceptual contribution. The S2FC family cleanly isolates the two orthogonal resources (input-independent state size vs input-dependent window) that pure architectures have to trade off; it gives a first principled framework in which the hybrid win can be stated and proved as a strict separation.
• Matched lower and upper bounds. Theorem 3.3 and Theorem 3.7 use clean, standard tools (Yao's principle + Fano; a sensitivity-based indistinguishability argument), and the hybrid constructions (Theorems 4.3, 4.6) are proved to escape both bounds with poly-logarithmic parameters and sub-linear window — this is exactly the "best-of-both-worlds" claim, formally.
• Constructive Mamba realisation, validated end-to-end. The constructions are explicit enough to be implemented (Appendix D.5) and, in the executed replication, both constructions/construction_var_copy.ipynb and constructions/construction_decode_recall.ipynb reproduce the correct last-position output on their sample inputs — a meaningful strengthening over pure existence proofs.
• Empirical breadth beyond the theory. The paper reports learnability experiments on four synthetic tasks (selective copy, AR-with-decoding, MKAR, needle-in-a-haystack), plus length-generalisation and OOD bit-proportion experiments at roughly 100M-parameter scale; this is a more thorough empirical picture than most theory-first papers provide.
• Honest reporting. The authors flag the MKAR case where $v$ is empty and S2FC does not directly predict a hybrid–Transformer gap, and they explicitly acknowledge the empirical anomaly of Transformer accuracy decreasing with larger context windows.
• Code release. Both the constructions (as notebooks) and the learnability sweeps (micro_hf/, mini/) are shipped, which makes the coverage of the empirical claims in principle re-runnable.

Weaknesses

• Theorem 3.7 proof does not deliver the stated $2/3$ threshold. The two-input indistinguishability construction only forces error probability $1/2$ under a uniform mixture, so the theorem's claim about models achieving success probability $\ge 2/3$ is not proved. Either weaken the statement to "$> 1/2$" or amplify (e.g. many-input construction) to force error above $1/3$.
• Theorem 3.3 proof requires a stronger success threshold than the statement asserts. The Fano step explicitly uses $\mathrm{err} < 1/8$, but the theorem promises the bound at success probability $\ge 1/2$. Tighten the statement to $\ge 7/8$ or extend the proof to the near-random regime.
• Theorem 3.3 bound can be vacuous under Assumption 3.2 alone. The injection $G: V^m \to Y^q$ implies $m\log|V| \le q\log|Y|$, so the raw bound $\Omega(m\log|V| - q\log|Y|)$ is non-positive without further structure on $|Y|$ and $q$. The end-of-proof step ("$q\log|V| = O(m\log|V|)$") uses $|V|$ where $|Y|$ was needed and does not follow from the assumption; the statement needs an explicit control on $q$ or $|Y|/|V|$.
• Lemma 3.5 statement uses $\sum$ where $\prod$ is meant. The lemma bounds $|S'| \le \sum_j |S_j|$, but the proof constructs a Cartesian product and the downstream information-theoretic argument in Theorem 3.3 uses $\sum_i \log|S_i|$ — this is a real typo that a careful reader will catch immediately.
• Theorem 4.6 statement names the wrong task. As written, the formal theorem states the hybrid "can solve the selective copying task with probability $99\%$", contradicting its own preamble (associative recall with decoding).
• Selective copying definition is internally inconsistent. Definition 4.1 writes $i = \arg\max_{1\le i\le L}\mathbf{x}_i\in\mathcal{N}$ (which reads as an argmax over values), while the figure and the Mamba construction use the position of the most recent number token. Additionally, $\mathbf{x}_{L+1-\mathbf{x}_i}$ tacitly identifies $\mathcal{N}=\{\#1,\dots,\#N\}$ with $\{1,\dots,N\}$; this identification should be stated.
• Definition 4.4 applies $\Phi$ to the wrong argument. $\Phi(\vec{\mathbf{x}})$ is defined but $\Phi(v(\vec{\mathbf{x}}))$ is used; also $v(\vec{\mathbf{x}})\in\{0,1\}^{\log|V|}$ should presumably be $\{0,1\}^{\log|\mathcal{M}|}$.
• MKAR key overlaps the query region. As stated, $K = \mathbf{x}_{L-k:L}$ trivially matches itself in the search for "the last position $i$ where $\mathbf{x}_{i:i+k}=K$", so the target position is undefined unless the search is restricted to $i < L-k$.
• Proof of Theorem 4.2 assigns two distributions to the same coordinate range. The construction of $D_T$ describes $\vec{\mathbf{x}}_{L/2:L}$ as uniform on $\mathcal{M}$ and uniform on $V$ in the same sentence — the second range is presumably $\vec{\mathbf{x}}_{1:L/2}$.
• Receptive-field additivity is used without proof. The claim that a $k$-layer sliding-window Transformer's last-position output depends only on the last $\sum_i W_i$ tokens is asserted rather than shown; a short induction over layers should be added.
• MKAR and NH sit outside the S2FC hardness result. The paper flags that with $v$ empty the S2FC Transformer-hardness argument does not fire, but the Results paragraph nonetheless asserts the "same separation" with 6× fewer parameters. Either state that the MKAR/NH separation is observational, or extend the framework to cover $v = \emptyset$.
• MKAR mechanism gap. On MKAR the pure SSM sits at ~0.52 while the SSM→TF hybrid saturates near 1.0. Under the paper's own mechanism ("SSM as encoder for long-range control") it is not clear why the SSM alone fails so hard on a task the paper explicitly says has no long-range control; a short mechanistic hypothesis (e.g. SSM read-out is the retrieval bottleneck) would connect the empirical result back to the theory.
• Empirical claim about Transformer window is unresolved. Increasing the Transformer's attention window degrades performance on both selective copy and MKAR, which contradicts Theorem 3.7's implication that a larger window can only help. Attributing this to "learnability" is plausible but not controlled for (no LR sweep, no optimiser ablation), and the same paragraph reports that hybrids "also frequently suffer" at large windows — this partially blunts the "hybrids are strictly better than pure Transformers" reading.
• Ablation gap: no full-attention Transformer baseline at hybrid parameter counts. Comparing sliding-window Transformers to hybrids conflates architectural expressivity with the window schedule imposed on the Transformer baseline; a full-attention Transformer at matched parameters would isolate the SSM-as-summariser mechanism.
• Empirical scope: $L\le 100$ across the small-scale sweeps. The theoretical story is asymptotic in $L$ (window $\Omega(L)$), but the parameter-efficiency claim is drawn from $L=100$ tasks; an $L$-sweep is needed to link the two.
• Length-generalisation attribution not isolated. The 10% long-sequence gap is drawn from a single-configuration comparison using RoPE only; without per-seed variance and at least one alternative positional encoding (e.g. ALiBi or learned), the gap cannot be cleanly attributed to the hybrid architecture rather than to RoPE's frequency schedule.
• Table 1 layout is ambiguous. The train-proportion × architecture entries are flattened into one cell, so the SSM/TF/Hybrid mapping cannot be recovered from the rendered table, and the bolding convention is not stated.
• Adversarial-distribution framing conflates "hard for each" with "hard for both". The mixture $D_S/2 + D_T/2$ is hard for SSMs on $D_S$ and hard for Transformers on $D_T$, not "hard for both models on half of the instances".
• Scope: hardness for Transformers is stated for sliding-window attention only. The abstract and intro speak of "any pure Transformer", but Theorem 3.7 covers only sliding-window (with full attention as $W = L$); this restriction should be flagged up front.
• Figure 4 caption overstates. "Hybrid models" (plural) attain perfect accuracy at ~2000 parameters, but only the SSM→TF ordering does; TF→SSM is at 0.433, comparable to the pure models.
• Conclusion overstates the SSM comparison. "Substantially smaller working memory than either pure counterpart" conflates working memory (the pure-Transformer limitation) with parameter count (the pure-SSM limitation).
• Probability thresholds vary between theorems without a unified explanation. Theorems 3.3 ($1/2$), 3.7 ($2/3$), and 4.2/4.5/4.6 ($90\%$/$99\%$) all use different regimes with no accompanying remark.
• GSSM read-out alphabet vs state budget. Section 2 defines $r: S\to V$, yet the S2FC argument treats $\log|S|$ as capacity for outputs in $Y$; a compatibility statement (or a reduction to $Y \subseteq V^n$) is needed.
• Appendix restatements silently narrow the main-text. Lemma C.1 restricts to "auto-regressive SSM" (Lemma 3.5 does not); the appendix Theorem 3.3/3.7 restatements are the ones that carry the quantitative content.
• Local presentational issues. Claim numbering (D.3 vs D.7) and index labelling ($n_1$ vs $n_L$) are inconsistent between statement and proof in the selective-copying construction; $\psi$ vs $\psi'$ are used interchangeably; $M$ is overloaded (set vs marker token) in the NH definition.

Reproducibility & code

• Constructions reproduce cleanly but only on one input. Re-executing constructions/construction_var_copy.ipynb and constructions/construction_decode_recall.ipynb (both re-run from scratch, execution_count reset to sequential 1, 2, 3…) yields correct last-position outputs for the hand-set-weight hybrid on both selective copy (target/pred at last = 2) and associative recall with decoding (target/pred at last = 20), and the residual-stream decode diagnostics agree. However, both notebooks only display a single input batch and never assert correctness across many random draws; Theorem 4.3's "for every input sequence" and Theorem 4.6's "with probability $\ge 99\%$" therefore rest on visual inspection of one example.
• Headline aggregates: only the ~12000-param column of Figure 4 was reached, and pure SSM there collapsed to chance. Under a single-core SLURM cap (with the mamba_ssm CUDA fast path unavailable and the pure-PyTorch fallback scan in use), the micro_hf/ var-copy sweep completed only single-seed runs at ~12000 params (with a few smaller-scale diagnostic points). At ~12000 params the qualitative signature holds — SSM→TF = 1.000 vs the paper's 1.000 — but Pure TF (0.734 vs paper 0.923), TF→SSM (0.605 vs 0.908) and especially Pure SSM (0.044 vs 0.931; a clear downscaling artefact from num_examples=100/epochs=4) are all well outside tolerance, and the ~1000/~2000/~6000 columns and the process.ipynb aggregation into Figure 4 were never reached. The 6× parameter-efficiency claim is therefore neither confirmed nor refuted from this run.
• Not attempted: decode-recall (Figure 5), MKAR (Figure 6), NH (Figure 7), length generalisation (Figure 8), and OOD Table 1. No result files exist for any of these under micro_hf/results/ or mini/results/. The environment (single-core CPU, no mamba_ssm CUDA kernels) is only part of the reason: the release itself needs manual toggling to run several of these — e.g. experiments = [experiment5_params] in micro_hf/run.ipynb sets num_layers=1 by default so the 2-layer Figure 6 configuration is not what runs out of the box, and the head/state-dim, window, adversarial-distribution and vocabulary-size sweeps live on inactive branches. There is also no aggregation code for Figure 8 or Table 1 in the shipped notebooks.
• Length-generalisation command in the release is broken as shipped. mini/cmd.txt uses --min_train_len/--max_train_len but mini/main.py's argparse expects --min_train_length/--max_train_length; the sample values (30–50) also do not match the paper's (20–50). Combined with a hard-coded num_train_epochs=1 in mini/train_utils.py that silently overrides any --epochs flag, the shipped invocation will not reproduce Figure 8 even given the compute.
• Seed count mismatch and undocumented LR-selection rule. Appendix E.1 reports 11 seeds and a $10^{-4}$–$10^{0}$ LR grid, but micro_hf/run.ipynb iterates range(5) and does not set --seed, and the paper does not describe which LR the reported curves collapse the sweep to (best-of-sweep? median? per-config best on val?). The release's auto_lr logic in micro_hf/main.py is not documented. This makes the 10/90 percentile bands in Figures 4/6/7 hard to reconstruct even when the sweeps do run.
• "Trained to convergence" is not implemented in the training scripts. micro_hf/main.py defaults to --epochs 4 --num_examples 1000, and mini/train_utils.py hard-codes num_train_epochs=1. The paper's "4× compute after plateau" stopping rule is not encoded in the release.
• Environment fragility. mamba_ssm (and causal-conv1d) could not be built during the replication because the upstream setup.py hardcodes nine GPU-arch gencodes and ignores TORCH_CUDA_ARCH_LIST; the shipped model correctly falls back to a pure-PyTorch sequential Mamba scan, but that path is far slower and was a major contributor to sweeps not finishing. The compute assumptions and expected wall time per figure are not stated anywhere in the release.

Recommended Changes

Essential.

• Fix Theorem 3.7 to match its proof. Either weaken the statement to "success probability $> 1/2$" or extend the proof (e.g. via a many-instance amplification) to force error above $1/3$ for models achieving success $\ge 2/3$ — as raised in Theorem 3.7 proof does not deliver the stated $2/3$ threshold.
• Reconcile Theorem 3.3's threshold with the Fano step. Either restate the theorem with success probability $\ge 7/8$ (matching $\mathrm{err} < 1/8$ used in the proof), or provide a Fano-style argument valid at $\mathrm{err} \le 1/2$ — cf. Theorem 3.3 proof requires a stronger success threshold.
• Add an explicit assumption that makes Theorem 3.3 non-trivial. State a bound on $q$ (or on the ratio $\log|Y|/\log|V|$) that ensures $m\log|V| - q\log|Y| = \Omega(m\log|V|)$, so the bound is not driven by an unconstrained free parameter — cf. Theorem 3.3 bound can be vacuous.
• Replace $\sum$ with $\prod$ in Lemma 3.5 (and C.1). The lemma statement currently reads $|S'|\le \sum_j |S_j|$; the proof and downstream use both require $\prod$.
• Correct Theorem 4.6's task name. Replace "the selective copying task" with "the associative recall with decoding task" in the formal statement.
• Rewrite the selective-copying definition unambiguously. State the argmax is over positions of number tokens, and explicitly identify $\mathcal{N} = \{\#1,\dots,\#N\}$ with $\{1,\dots,N\}$ before using $\mathbf{x}_i$ as an integer offset — cf. Selective copying definition.
• Fix the domain of $\Phi$ and the associated bit length in Definition 4.4. Write $\Phi(v(\vec{\mathbf{x}}))$ (or define $\Phi$ on $\{0,1\}^{\log|\mathcal{M}|}$) and replace $\log|V|$ with $\log|\mathcal{M}|$.
• Fix the MKAR definition. Restrict the search for the matching key to $i < L-k$ so the query region does not match itself.
• Correct the coordinate range in the $D_T$ construction (proof of Theorem 4.2). Change one of the two "$\vec{\mathbf{x}}_{L/2:L}$" to "$\vec{\mathbf{x}}_{1:L/2}$" so the two sub-ranges are disjoint.
• Add the sliding-window receptive-field lemma. Prove by induction that a $k$-layer sliding-window Transformer's last-position output depends only on the last $\sum_i W_i$ tokens.
• Align mini/cmd.txt with mini/main.py's argparse and remove the hard-coded epoch cap. Rename flags to --min_train_length/--max_train_length, update the sample train range to 20–50, and lift the num_train_epochs=1 line in mini/train_utils.py — cf. Length-generalisation command in the release is broken.
• Ship a README mapping figure → notebook cell / flag, plus aggregation scripts for Figure 8 and Table 1. Also make the 2-layer Figure 6 config the default (rather than an inactive branch), so the paper's headline empirical claims are one-command reproducible — cf. Not attempted: decode-recall, MKAR, NH, length generalisation, and OOD Table 1.

Suggested.

• Reformulate Table 1 with one column per architecture and state the bolding convention so the row-to-value mapping is unambiguous — cf. Table 1 layout.
• Qualify the Figure 4 caption to say "the SSM→TF hybrid" (not "hybrid models") attains perfect accuracy at ~2000 parameters, matching the main-text observation that the reverse hybrid tracks the pure models.
• Split the conclusion sentence about "either pure counterpart" into two clauses, matching the two separate resource axes (parameter count vs window size) — cf. Conclusion overstates the SSM comparison.
• Add a full-attention Transformer baseline at hybrid parameter counts on selective copy and MKAR, to disentangle the SSM-as-summariser benefit from the effect of the imposed sliding window on the Transformer baseline — cf. Ablation gap.
• Add a length ($L$) sweep for the small-scale claims. Show whether the 6× parameter gap persists as $L$ grows past 100, connecting the $L=100$ experiments to the asymptotic hardness — cf. Empirical scope.
• Add a controlled experiment on the "Transformer worsens with larger window" phenomenon. A joint LR/warmup/optimiser sweep (or a citation to prior work reporting the same trainability issue) would either strengthen the "learnability" attribution or reveal a scope limitation of the hybrid advantage.
• Add a mechanistic hypothesis for the MKAR SSM collapse. Discuss (or ablate) why the pure SSM fails on a task the paper describes as having no long-range control — cf. MKAR mechanism gap.
• Add an alternative positional encoding and per-seed variance for length generalisation. Report at least one non-RoPE variant (e.g. ALiBi or learned) and per-seed bands so the 10% long-sequence gap is attributable to architecture — cf. Length-generalisation attribution not isolated.
• Address MKAR and NH within (or explicitly outside) the S2FC framework. Either add a modified hardness assumption that covers the $v=\emptyset$ case, or explicitly mark the MKAR/NH results as observational — cf. MKAR and NH sit outside the S2FC hardness.
• Signpost the sliding-window scope of Theorem 3.7 in the abstract/intro so "any pure Transformer" is not misread as excluding sparse or context-dependent attention.
• Standardise probability thresholds across Theorems 3.3, 3.7, and 4.2/4.5/4.6 (or add a remark explaining why hardness statements use $1/2$–$2/3$ while concrete tasks use $90\%$–$99\%$).
• Rephrase the adversarial-distribution paragraph so it says "SSMs perform poorly on the $D_S$ half and Transformers on the $D_T$ half", rather than suggesting each half is hard for both — cf. Adversarial-distribution framing.
• Fix the seed configuration in micro_hf/run.ipynb (iterate 11 seeds and set a --seed flag) to match the reported 10/90 percentile bands, or update the paper text to say 5 seeds and re-check the bands — cf. Seed count mismatch.
• Document the LR-selection rule used to draw the reported curves, and cache the chosen LR per configuration — cf. Undocumented LR-selection rule.
• Implement the "trained to convergence" rule. Add an explicit stopping criterion (patience-based early stop on validation) in micro_hf/main.py and mini/train_utils.py, or state the compute per configuration used for the reported curves — cf. "Trained to convergence" is not implemented.
• Add a randomised correctness assertion to the construction notebooks (e.g. 1000 sampled inputs, exact match at the last position) so the "$\ge 99\%$ on every input" claim of Theorem 4.6 is checked programmatically — cf. Constructions reproduce cleanly but only on one input.
• Signpost compute assumptions and the mamba_ssm build path in a README. Document required GPU archs / expected wall time per figure, and either loosen the mamba_ssm gencodes or ship a fallback-only path with an explicit performance note — cf. Environment fragility.
• Unify local notation. Choose $\psi$ vs $\psi'$ consistently; rename the NH marker token so it does not collide with $\mathcal{M}$; fix the "Claim D.3 / Proof of Claim D.7" numbering and the $n_1$ vs $n_L$ index in the selective-copying construction.