PaperlayerView exact revision on arXiv

Source-linked AI summary

Non-Commutative State Tracking with Input-Dependent Low-Rank Updates in Mamba-3

Hiroki Fujii, Masaki Yamakita

arXiv:2609.28273v2eess.SYcs.LGcs.SY

TL;DR

State tracking requires retaining information while composing observed operations in order, a challenge for diagonal-transition models. The paper adds an input-dependent reflection term to Mamba-3 and evaluates it on discrete group problems and a continuous-observation shell game. It reports improved tracking over standard Mamba-3 under the evaluated conditions, with stronger performance on longer timing-jitter shell-game sequences than selected alternatives.

  • Problem

    State tracking must distinguish retaining information from composing operations, while diagonal-transition models face limits on exact non-commutative tracking at arbitrary lengths under TC0 ≠ NC1.

    Source: §I · p. 1

  • Method

    The paper adds an input-dependent rank-one reflection term to Mamba-3 and evaluates it on group word problems and a behavioral-cloning shell game.

    Source: §I · p. 1, B. Shell Game · p. 3

  • Results

    The proposed model improves tracking over standard Mamba-3 on S5 generators and the continuous-observation shell game, and shows an advantage over selected alternatives on longer timing-jitter sequences.

    Source: C. Generalization under Timing Jitter · p. 5, §VII · p. 7

  • Takeaways & Limitations

    Mamba-3 + NPLR provides an additional Mamba-3-based option for non-commutative state tracking beyond existing delta-based models.

    Source: §VII · p. 7

    Evaluations cover discrete inputs and low-dimensional observations, not image observations or physical systems; shell-game actions do not affect subsequent observations.

    Source: C. Applications and Limitations · p. 6

Abstract

from arXiv · show

State tracking from sequential observations can require both retaining information and updating it by composing observed operations. We extend Mamba-3's diagonal transition with an input-dependent low-rank reflection term to support noncommutative state tracking, in which the order of operations matters. The rank-one update couples state coordinates along an input-dependent direction, enabling non-diagonal state transitions within a single Mamba-3 block. The extension preserves Mamba-3's exponential-trapezoidal discretization, rotary embeddings (RoPE), and readout. For training, we adapt chunkwise computation to parallelize the proposed recurrence within each chunk. Experiments cover group word problems with discrete inputs and a shell game with continuous observations, in which a policy is trained by behavioral cloning. Among the models selected for their strong performance under fixed timing, the proposed model maintains higher tracking success on longer swap sequences in the shell game with continuous observations and timing jitter. These experiments show that the proposed method achieves high accuracy on the evaluated non-commutative tracking tasks, improving on standard Mamba-3. The extension thus offers a Mamba-3-based approach to non-commutative state tracking.

I. INTRODUCTION

State tracking must preserve history while composing observed operations in order, a challenge for diagonal-transition models on non-commutative tasks. The paper introduces an input-dependent reflection update for Mamba-3 and evaluates it on group word problems and a continuous-observation shell game.

I. INTRODUCTION

  • Evaluating state tracking requires distinguishing information retention from operation composition, since an object's initial location no longer identifies its current position after an exchange.SSMs encode histories in time linear in sequence length, and Mamba-3 improves state tracking through complex-valued state updates.Source: §I · p. 1
  • Ordered permutation composition motivates non-commutative transitions; fixed-depth Transformers and diagonal-transition SSMs are analyzed as uniform TC0, while fixed finite non-solvable group word problems are NC1-complete.Under the conjecture TC0 ≠ NC1, the analyzed architectures cannot solve all such words exactly at arbitrary length; Mamba required increasing depth for high accuracy on longer sequences [3].Source: §I · p. 1
  • The proposed input-dependent reflection update enables non-commutative state tracking within one Mamba-3 block and is evaluated on S5 and a continuous-observation shell game.The paper also positions this as an additional Mamba-3-based option beyond delta-based models.Source: §I · p. 1

II. RELATED WORK

  • Prior work studies shallow circuits and Transformers for solvable groups, input-dependent transitions, and reflection-based delta models, including DeltaProduct and GDN.Mamba-3 had been evaluated on parity and modular arithmetic but not word problems over non-solvable groups; imitation-learning applications motivate testing whether a Mamba backbone tracks locations through successive swaps.Source: §II · pp. 1–2

A. Input-Dependent Low-Rank State Updates

Mamba-3 + NPLR adds an input-dependent rank-one reflection to the diagonal scan transition while retaining the block's existing discretization and readout components.

B. Input-Dependent Low-Rank State Updates

  • The extension adds a low-rank reflection term to Mamba-3's diagonal scan transition while retaining rotary embeddings, trapezoidal correction, normalization, and gates.Removing the low-rank term recovers the Mamba-3 state update.Source: A. Input-Dependent Low-Rank State… · pp. 2–3
  • Each head has P channels with separate n-dimensional states, while the input x_t is obtained by embedding the observation.The scalar decay, time step, and trapezoidal-correction coefficient follow Mamba-3; the preceding-input contribution is zero only at sequence start.Source: A. Input-Dependent Low-Rank State… · p. 2
  • Cumulative RoPE rotates the projected B/C vectors, and the reflection direction uses the same cumulative rotation.The no-RoPE ablation sets the rotation to identity for B/C and k_t while retaining the reflection term.Source: A. Input-Dependent Low-Rank State… · pp. 2–3
  • Additional input projections generate a normalized direction k_t and coefficient β_t for a rank-one generalized Householder update.The coefficient and direction are input-dependent; the term modifies the previous state along that direction and can attenuate or reverse its sign.Source: A. Input-Dependent Low-Rank State… · p. 3

B. Non-Commutative Transitions from a Reflection Term

The reflection term makes Mamba-3 transitions input-dependent and potentially non-commutative, allowing order-dependent state updates under specific conditions.

B. Non-Commutative Transitions from a Reflection Term

  • Writing the scan transition as A_t = d_tI + N_t isolates scalar decay d_t from the input-dependent reflection term N_t.Without the low-rank term, N_t is zero and the transitions commute.Source: B. Non-Commutative Transitions from a… · p. 3
  • When both reflection coefficients are nonzero and their direction vectors are neither parallel nor orthogonal, transitions do not commute and can represent order-dependent updates.The reflection term is N_t = −β_t k_t k_t^T.Source: B. Non-Commutative Transitions from a… · p. 3
  • Non-commutativity applies to the transition alone, excluding input injection and readout, and does not guarantee that a group word problem will be learned successfully.Source: B. Non-Commutative Transitions from a… · p. 3

C. Chunkwise Computation for Training

Training adapts chunkwise computation to the reflection recurrence, while evaluation includes group-prefix composition tasks over cyclic and symmetric groups.

C. Chunkwise Computation for Training

  • Chunkwise training applies a lower-triangular solve to scalar decay, rank-one reflection, and trapezoidal injection, exactly rearranging the recurrence.Channels within a head share a coefficient matrix but use channel-specific right-hand sides; final states and preceding-injection factors carry across chunk boundaries.Source: C. Chunkwise Computation for Training · p. 3
  • For fixed chunk length c, forward state updates and readouts cost O(LnP + Lc(n + P)) per sequence and head, linear in sequence length L.This cost is stated after coefficient generation.Source: C. Chunkwise Computation for Training · p. 3

A. Group Word Problems

  • A group word problem predicts the group element represented by each input prefix, rather than only whether the prefix product equals the identity.The prefix product follows s_t = g_t ◦ s_(t−1), starting from the identity.Source: A. Group Word Problems · p. 3
  • The evaluation uses cyclic addition modulo five for Z5 and classifies S5 prefixes into 120 classes from adjacent transpositions.Training sequences contain at most 64 group elements, and accuracy is evaluated at each position and at the final position on length-64 and longer sequences.Source: A. Group Word Problems · p. 3

B. Shell Game

The shell game tests whether a policy can track a ball through visible cup swaps when the ball is revealed only initially. Models are evaluated with fixed timing and timing jitter, alongside comparisons of their architectures and task-specific mappings.

B. Shell Game

  • The ball is revealed only initially, and the policy must compose visible swaps to reach the cup containing it from continuous observations.Five visually identical cups undergo up to 16 adjacent swaps; cup identities are not observed.Source: B. Shell Game · p. 3
  • Models are first compared with fixed waiting and motion durations, then selected strong performers are evaluated with randomized swap intervals and speeds.Timing jitter varies durations within each episode.Source: B. Shell Game · p. 3
  • Fixed-timing models are trained by behavioral cloning for 20,000 updates with K ≤16 and sequence length T=213.Direct training uses the full swap-count range from the outset, with greater loss weight on steps after the cue.Source: B. Shell Game · p. 3
  • Success means the model’s output is nearest to the correct slot during the response interval; the diffusion-based prior provides a no-tracking comparison.The constant-prediction chance level is 20%; success at K = 0 indicates retention, while composition evidence requires exceeding the prior at K > 0.Source: B. Shell Game · pp. 3–4

C. Models Compared

  • Timing-jitter comparisons include Mamba-3 + NPLR, DeltaProduct (nh = 4), GDN, and a no-RoPE ablation, selected for strong fixed-timing performance.GDP is also evaluated as a supplemental control.Source: C. Models Compared · p. 4
  • Mamba-3 + NPLR uses one block, and its parameter counts are not matched across architectures; it has fewer parameters than DeltaProduct, GDP, GDN, and IDS4 on both tasks.Core block configuration and recurrent-state size are held fixed across tasks, while task-specific input and output mappings differ.Source: C. Models Compared · p. 4

V. EXPERIMENTAL RESULTS

Mamba-3 + NPLR achieves strong results on the evaluated group word problems and shell-game conditions, while its length extrapolation remains limited. In the shell game, it achieves perfect success under direct fixed-timing training, and on S5 it reaches perfect training-length accuracy where one-block Mamba-3 remains at chance.

A. Learning Group Word Problems with Low-Rank Updates

  • Mamba-3 + NPLR achieves perfect final-position accuracy at training length on S5 (generators) for all three runs, while one-block Mamba-3 without reflection remains at chance level.Increasing Mamba-3 depth improves accuracy, but four blocks still do not achieve perfect tracking at training length.Source: A. Learning Group Word Problems with… · p. 4
  • On Z5, Mamba-3 + NPLR and IDS4 achieve high accuracy at training length, while DeltaProduct configurations exceed chance but still make errors, as does GDN.These results apply to the evaluated configurations and training conditions.Source: A. Learning Group Word Problems with… · p. 4
  • The proposed model maintains high accuracy at twice the training length but degrades on longer S5 words; disabling RoPE yields perfect accuracy through evaluation length 512 for all five runs.At greater lengths, the no-RoPE variant’s accuracy declines and differences across runs remain.Source: A. Learning Group Word Problems with… · pp. 4–5, Figure 2 · p. 5

B. State Tracking in the Shell Game

  • Mamba-3 + NPLR achieves perfect success in both standard and extended fixed-timing shell-game conditions for all three runs.DeltaProduct also achieves high success in both conditions, with a small number of errors.Source: B. State Tracking in the Shell Game · p. 5
  • Under the swap-count curriculum, GDN also achieves near-perfect success in the standard shell-game condition.The corresponding comparison across models is reported in Table IV in Appendix D.Source: B. State Tracking in the Shell Game · p. 5

C. Generalization under Timing Jitter

In the evaluated timing-jitter shell game, Mamba-3 + NPLR retains high success on longer swap sequences, while disabling RoPE performs better at the longest evaluated lengths. Comparisons also show that component presence alone does not explain task performance.

C. Generalization under Timing Jitter

  • 87.83% at 128 swaps for Mamba-3 + NPLR, versus 21.81% for DeltaProduct and 21.16% for GDP, with mean success above 95% through 96 swaps.It achieves perfect success through 48 swaps in all three runs; these results indicate stronger but finite extrapolation in the evaluated timing-jitter setting.Source: C. Generalization under Timing Jitter · p. 5, Figure 4 · p. 6
  • 95.74% mean success at 128 swaps without RoPE, compared with 87.83% for the proposed model, and perfect success through 64 swaps in all three runs.The paper concludes that RoPE does not improve performance in this evaluated setting.Source: C. Generalization under Timing Jitter · p. 5

A. State Tracking Performance

  • Adding the input-dependent low-rank term improves shell-game success over standard Mamba-3 under the evaluated conditions, while the proposed model uses fewer parameters than several alternatives and no larger recurrent state.The proposed model uses one rank-1 reflection update per input, compared with two or four in the evaluated DeltaProduct configurations.Source: A. State Tracking Performance · p. 5

B. Architectural Components and Task Performance

  • GDP and the GDN extension achieve high accuracy on the generator task but show shell-game success rates from 20.87% to 100% across runs, so component presence alone does not explain performance.The no-RoPE variant matches shorter-length performance and has higher mean success at longer lengths in the timing-jitter setting; the evaluations provide no evidence of a RoPE benefit there.Source: B. Architectural Components and Task… · p. 6

C. Applications and Limitations

The evaluations cover discrete inputs and low-dimensional observations, but not image observations, physical systems, or settings where actions influence subsequent observations. The paper presents Mamba-3 + NPLR as an additional architectural option and identifies more realistic imitation-learning tasks as future work.

C. Applications and Limitations

  • The evaluations use discrete inputs and low-dimensional observations; effectiveness with image observations or physical systems has not been tested.In the shell game, model outputs do not affect subsequent observations, leaving action-observation feedback unexamined.Source: C. Applications and Limitations · p. 6
  • Future work will examine imitation-learning tasks that combine composing observed operations with using their timing to select actions.The evaluations leave open when RoPE’s additional temporal flexibility benefits non-commutative state tracking.Source: C. Applications and Limitations · p. 7

VII. CONCLUSION

  • Mamba-3 + NPLR combines Mamba-3’s input and readout mechanisms with reflection-based updates as an architectural option for non-commutative state tracking.The conclusion reports improved tracking over standard Mamba-3 on S5 generators and a continuous-observation shell game under the evaluated conditions.Source: §VII · p. 7

APPENDIX A CHUNKWISE STATE-UPDATE DERIVATION

The proposed rank-one recurrence is reorganized into a chunkwise linear system, allowing exact within-chunk computation and transfer of each chunk’s final state. The derivation also shares work across channels and avoids reconstructing every intermediate state for readout.

APPENDIX A CHUNKWISE STATE-UPDATE DERIVATION

  • Each chunk’s final state initializes the next chunk, while the trapezoidal injection retains preceding factors even at chunk boundaries.Those factors contribute zero only at the beginning of the full sequence.Source: Appendix A · p. 7
  • The rank-one recurrence combines the injection terms into b_t and updates h_t from the prior state through a reflection term.The chunk’s incoming state is h_0, and the diagonal factors d_l define products used in the chunkwise expansion.Source: Appendix A · p. 7
  • Within a chunk, the scalars q satisfy a unit lower-triangular system, whose solution permits exact state reconstruction rather than an approximation.The system uses M, q, and v; the resulting states are reconstructed after solving it.Source: Appendix A · p. 7
  • All P channels in a head share the coefficient matrix, so their channel-specific right-hand sides can be solved jointly.The channels share d_t, β_t, and k_t but retain their own injections and incoming states.Source: Appendix A · p. 7
  • The readout substitutes the chunkwise expression instead of reconstructing every time-step state, sharing n-dimensional inner products across P channels.Only the final state of each chunk is reconstructed and passed forward.Source: Appendix A · p. 7

APPENDIX B TASK AND TRAINING SETTINGS

The appendix specifies group-word and shell-game tasks, their training and evaluation settings, and configurations for Mamba-3 and DeltaProduct comparisons.

APPENDIX B TASK AND TRAINING SETTINGS

  • Repeated runs use different random seeds and freshly generated batches, and shell-game training is compared under direct training and a swap-count curriculum.The appendix reports the curriculum comparison in Table IV.Source: Appendix B · p. 8

A. Group word problems

  • Group-word tasks use sequences up to 64 elements, with direct training for Z5 and a length curriculum for S5 generators.Inputs are sampled uniformly from five Z5 increments or four adjacent S5 transpositions; S5 accuracy is compared with a parity-constrained 1/60 baseline.Source: A. Group word problems · p. 8

B. Shell-game task

  • The fixed-timing shell game uses five-slot observations, sampled adjacent swaps, and a behavioral-cloning loss weighted more heavily after the cue.Evaluation uses 4,096 episodes in standard (T = 213, K ≤16) and extended (T = 405, K ≤32) conditions, with success reported by swap count and overall.Source: B. Shell-game task · p. 8

C. Timing-jitter shell game

  • Timing-jitter training varies motion and waiting durations, while evaluation tests sequences with 16 to 128 swaps.Motion durations are sampled from U(4, 16); waiting durations use U(0, 8) with probability 0.8 and U(16, 64) otherwise.Source: C. Timing-jitter shell game · p. 8

A. Mamba-3 configurations

  • Mamba-3 configurations retain RoPE, trapezoidal correction, and normalization, with the proposed model adding reflection updates to one block.The configurations use d_model=64, d_state=16, head dimension 16, and expansion factor 4; task-specific input and output mappings are shared across the comparison.Source: A. Mamba-3 configurations · p. 8

B. DeltaProduct configuration

  • DeltaProduct uses the authors’ one-layer block with four heads and state dimension 32 per head, tested at n_h = 2 and 4.The decay gate is disabled and negative eigenvalues are allowed; optimization uses Adam at learning rate 10^-3 without early stopping.Source: B. DeltaProduct configuration · p. 8

C. Gated DeltaNet configuration

The shell-game comparison reports perfect success for Mamba-3 + NPLR in both evaluation conditions under either training procedure, while the appendix also describes comparator configurations and model costs.

C. Gated DeltaNet configuration

  • The GDN configuration extends the coefficient range to β ∈(0, 2), which improves S5-generator accuracy over the original (0, 1) range under the same training conditions.It retains the short convolution, output gate, normalization, residual connections, and MLP.Source: C. Gated DeltaNet configuration · p. 9

D. IDS4 configuration

  • The IDS4 comparison uses one layer with its public default dimensions and evaluates each condition over three runs.Its configuration uses d_model = 255 and d_state = 110, with AdamW at learning rate 10^-4 and weight decay 0.01.Source: D. IDS4 configuration · p. 9

E. Model sizes and state-update costs

  • Mamba-3 + NPLR has fewer parameters than DeltaProduct, GDP, GDN, and IDS4 on both tasks, while sharing recurrent-state size with one-block Mamba-3, DeltaProduct, and GDP.The reported transition-update costs exclude coefficient generation, input injection, and readout.Source: E. Model sizes and state-update costs · p. 9, Table III · p. 9

A. Shell-game training comparison

  • Mamba-3 + NPLR achieves perfect shell-game success in standard and extended conditions under both direct training and the swap-count curriculum.Even four-block Mamba-3 makes errors in the extended condition; one IDS4 run predicts only the central slot, and the other two runs err even without swaps.Source: A. Shell-game training comparison · p. 9, Table IV · p. 9

B. Depth comparison

The appendix reports depth-dependent S5 accuracy, defines the length-extrapolation and swap-count evaluations, and tests whether the timing-jitter advantage persists with a decay gate.

B. Depth comparison

  • At S5 generator length L = 64, final-position accuracy ranges from 3.52% to 6.64% for two-block Mamba-3 and from 47.46% to 61.91% for four-block Mamba-3 across three runs.Source: B. Depth comparison · p. 9

D. Success rates by swap count

C. Length extrapolation

  • Length extrapolation evaluates models trained at length 64 using final-position accuracy, with mean, minimum, and maximum reported across runs.The evaluation uses five runs for S5 generators and three for Z5, with shared evaluation sequences at each length.Source: C. Length extrapolation · p. 9, Table V · p. 10

E. Supplemental comparison with Gated DeltaProduct

  • With the decay gate enabled, Mamba-3 + NPLR retains its advantage on longer timing-jitter sequences, while direct-training GDP runs show greater fixed-timing variability.One fixed-timing run converged to a constant prediction, and one timing-jitter run predicted only the central slot.Source: E. Supplemental comparison with Gated… · p. 10, Table VIII · p. 11

Related papers

Loading 2609.28273v2…