Mamba 3 model updated to track order sensitive states better
Non-Commutative State Tracking with Input-Dependent Low-Rank Updates in Mamba-3
Machine Learning
Summary
Tracking the order of events is important in many tasks because the sequence affects the outcome. The authors improved a model called Mamba-3 by adding a new way to update its internal memory that depends on the input and links different parts of the memory together. This allows the model to better remember sequences where order matters. They tested it on challenges involving tracking swaps and games, showing it works better than before when things happen in complex orders. This makes it more reliable for tasks where the order of observations can't be mixed up.
What this means in practice
- •For robotic system developers: Improve robot controllers that must track changing object arrangements where the order of moves is important, using the enhanced Mamba-3 model.
- •For interactive game developers: Create game AI that better understands sequences of player actions where the order changes the state, enabled by noncommutative state tracking.
Authors
Hiroki Fujii, Masaki Yamakita
Abstract
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.