Source-linked AI summary

StateSMix: Online Lossless Compression via Mamba State Space Models and Sparse N-gram Context Mixing

Roberto Tacconelli

arXiv:2605.02904v1cs.LGcs.IT

TL;DR

StateSMix targets practical lossless compression without the large external weights and GPU inference associated with some neural compressors. It combines an online-trained Mamba SSM with sparse n-gram logit biasing and arithmetic coding, and beats xz on moderate-size enwik8 inputs while crossing over to xz superiority at approximately 30 MB.

  • Problem

    Neural compression faces a trade-off between prediction quality and model cost, while some strong systems require large external weights and GPU inference.

  • Method

    StateSMix trains a Mamba SSM online from scratch and combines its predictions with sparse n-gram logit biasing and arithmetic coding in a self-contained compressor.

  • Results

    StateSMix beats xz by 8.7% on 1 MB, 5.4% on 3 MB, and 0.7% on 10 MB enwik8 inputs.

  • Takeaways & Limitations

    The SSM is the primary compression engine, while n-gram tables add exact local and long-range memorisation as a complementary gain.

  • Takeaways & Limitations

    The crossover to xz superiority at approximately 30 MB reflects LZMA’s ability to exploit long-distance repetitions beyond fixed-horizon n-gram tables.

Abstract

from arXiv · show

We present StateSMix, a fully self-contained lossless compressor that couples an online-trained Mamba-style State Space Model (SSM) with sparse n-gram context mixing and arithmetic coding. The model is initialised from scratch and trained token-by-token on the file being compressed, requiring no pre-trained weights, no GPU, and no external dependencies. The SSM (DM=32, NL=2, approximately 120K active parameters per file) provides a continuously-updated probability estimate over BPE tokens, while nine sparse n-gram hash tables (bigram through 32-gram, 16M slots each) add exact local and long-range pattern memorisation via a softmax-invariant logit-bias mechanism that updates only non-zero-count tokens. An entropy-adaptive scaling mechanism modulates the n-gram contribution based on the SSM's predictive confidence, preventing over-correction when the neural model is already well-calibrated. On the standard enwik8 benchmark, StateSMix achieves 2.123 bpb on 1 MB, 2.149 bpb on 3 MB, and 2.162 bpb on 10 MB, beating xz -9e (LZMA2) by 8.7%, 5.4%, and 0.7% respectively. Ablation experiments establish the SSM as the dominant compression engine: it alone accounts for a 46.6% size reduction over a frequency-count baseline and beats xz without any n-gram component, while n-gram tables provide a complementary 4.1% gain through exact context memorisation. OpenMP parallelisation of the training loop yields 1.9x speedup on 4 cores. The system is implemented in pure C with AVX2 SIMD and processes approximately 2,000 tokens per second on commodity x86-64 hardware.

1. Introduction

StateSMix addresses the quality–cost tension in neural compression with fully online Mamba-based prediction and sparse n-gram mixing. Its compact, self-contained design targets practical compression without pre-trained weights or GPU inference, while performance eventually gives way to xz on larger files.

  • Motivation: Neural compressors trade prediction quality against transmitted model size and inference cost.LLM-based systems can require hundreds of megabytes of external weights and GPU inference.
  • Approach: StateSMix trains a Mamba-style SSM online from random initialisation, avoiding pre-trained weights, external weights, and GPU inference.The model uses DM=32 and NL=2, with linear-time inference and a compact recurrent state.
  • Approach: Entropy-adaptive mixing scales n-gram bias with SSM uncertainty, increasing n-gram contribution when the SSM is uncertain and reducing it when confidence is high.The logit-bias mechanism exploits softmax translation invariance and updates only tokens with non-zero counts.
  • Efficiency: 10–30% lower head-projection cost follows from modelling only tokens present in the current file.The effective vocabulary is reduced from 49,152 to 18K–44K.
  • Results and boundary: At approximately 30 MB, xz becomes superior because LZMA exploits long-distance repetitions beyond fixed-horizon n-gram tables.Before that crossover, the SSM alone beats xz on enwik8 3M and the full system achieves a further 3.7% reduction.

2. Related Work

StateSMix is situated among dictionary, context-mixing, and neural compressors that differ in compression quality, compute, memory, and model-weight requirements. Its distinguishing regime is online training from scratch with a self-contained output.

  • Classical compression: Classical dictionary compressors exploit byte repetition, while arithmetic coding approaches source entropy using adaptive symbol probabilities.xz −9e (LZMA2) achieves ∼1.99 bpb on enwik8 and is the primary baseline.
  • Context mixing: PAQ8px and CMIX achieve lower reported enwik8 bpb than classical tools but require extreme compute, memory, or slow processing.PAQ8px reaches ∼1.27 bpb, while CMIX reaches ∼1.17 bpb at 0.5–5 KB/s with 16–64 GB RAM.
  • Neural compression: Recent neural compressors report strong ratios using large pre-trained or quantized recurrent models and arithmetic coding.The cited systems include Chinchilla 70B, LLaMA-3-8B with LoRA, and RWKV-169M with 8-bit quantisation.
  • Online compression: LLM-based compressors require weights to be transmitted or pre-shared, whereas StateSMix trains online and transmits the model implicitly through the compressed output.This contrasts with StateSMix’s self-contained design.
  • State-space models: Mamba introduces input-dependent SSM selection while retaining O(N) inference complexity, and StateSMix applies this style of SSM to lossless compression.The paper states this is the first such application to lossless compression.

3. Background

The paper frames lossless compression as online probability prediction followed by arithmetic coding. StateSMix maintains encoder–decoder synchrony by updating the predictor with each recovered token, while Mamba supplies linear-time recurrent processing.

  • Arithmetic coding: Arithmetic coding assigns code length −log2 q(t_i | t_<i) bits, making expected code length equal to model cross-entropy.The decoder must reproduce the same sequence of distributions for lossless reconstruction.
  • Lossless synchronisation: StateSMix guarantees encoder–decoder agreement by updating both models with the true token after every step.This maintains identical recurrent state and probability distributions.
  • SSM formulation: A continuous-time SSM evolves a hidden state through input and state-transition terms, then produces an output from the state and input.The discrete recurrence uses step ∆ to update h_i and y_i.
  • Mamba selectivity: Mamba makes B, C, and ∆ input-dependent, enabling per-token control over state retention while retaining O(N) inference complexity.This supports forgetting irrelevant context and focusing on salient tokens during online processing.
  • Online learning: Online learning predicts before observing each token and updates afterward, with stochastic gradient descent on chunked cross-entropy providing a practical algorithm.The regret formulation measures cumulative −log q(t_i) loss.

4. Method

StateSMix combines BPE tokenisation, compact vocabulary remapping, an online Mamba predictor, sparse n-gram biasing, and arithmetic coding in a mirrored compression–decompression pipeline. The implementation reduces vocabulary-dependent work and maintains recurrent and convolutional state across tokens.

  • System Overview: The compression pipeline tokenises input with BPE, remaps the active vocabulary, predicts and encodes online, updates the model, and serialises the output.Decompression mirrors the same predict-update loop.
  • Tokenisation and vocabulary: BPE converts raw bytes into discrete token IDs, and compact remapping allocates model computations only over tokens present in the file.For enwik8, ve = 44,298 gives a ∼10% reduction; for 1 MB excerpts, ve = 18,058 gives a 63% reduction.
  • SSM predictor: The predictor uses two Mamba layers, layer normalisation, and a linear language-model head over the effective vocabulary.At ve = 44,298 and DM = 32, the head projection is approximately 1.4M MADs per token.
  • Mamba layer: Each Mamba layer combines normalisation, input projection, depthwise convolution, selective SSM recurrence, gating, and a residual output projection.The recurrent state is carried across tokens, alongside convolution buffers.
  • Parameterisation: The model’s parameter count scales with the effective vocabulary as 19,776 + 2 ve DM.The fixed architecture uses DM=32, DS=16, DI=64, and NL=2.

4.4 Online Training

StateSMix trains its SSM online in chunks while combining sparse n-gram evidence through softmax-invariant logit updates. Hash tables capture contexts from bigrams through 32-grams, with entropy-adaptive scaling controlling their influence.

  • Online optimization: 32-token chunks trigger Adam updates on cross-entropy loss while encoding proceeds simultaneously.The SSM state is detached at chunk boundaries, implementing truncated BPTT.
  • Online optimization: 8, 4, and 2 Adam iterations are applied to chunks 1–10, 11–30, and 31+, respectively.The warm-up schedule uses more iterations early, before n-gram tables have accumulated observations.
  • Sparse context mixing: Only tokens with non-zero n-gram counts receive logit updates, while unseen tokens retain zero bias.Softmax translation invariance makes this sparse update equivalent to a renormalised probability adjustment.
  • Sparse context mixing: O(fan-out) operations per token make sparse n-gram mixing memory- and compute-efficient without a dense probability vector.Natural-language fan-out is typically 5–30.
  • Context storage: 16M-slot hash tables cover orders 2–8, 16, and 32, while the bigram uses a collision-free direct array.Higher-order contexts use open addressing with linear probing depth 8 and sparse per-context count arrays.
  • Context storage: 3.45 and 6.91 logit boosts result from single observations in the 16-gram and 32-gram tables, respectively.These correspond to approximately 32× and 1,000× probability increases under the stated settings.
  • Adaptive mixing: Entropy-adaptive scaling reduces n-gram influence to about 0.2 for confident SSM predictions and raises it to about 2.5 under high uncertainty.The mechanism is described as preventing n-gram over-correction when the SSM is well calibrated.

4.6 Additional Context Models

Additional context models supplement the combined SSM and n-gram logits with specialised predictors and recency information. The implementation uses integer arithmetic coding and AVX2-accelerated C kernels.

  • Specialised predictors: The LZ hash predictor stores the most recent next-token prediction for each two-token context and boosts that continuation according to confidence count.It targets associations too specific for the probabilistic n-gram model.
  • Specialised predictors: The last 64 tokens receive an exponentially decaying recency bonus to capture within-sentence repetition patterns.The normalised age ranges from 0 to 1, with λr = 0.05.
  • Specialised predictors: A count-based baseline adds λc log(ctotal + 1) to provide a smooth prediction before the SSM has been trained.The cited implementation sets λc = 0.1.
  • Combined prediction: The final probability distribution is softmax over the SSM logits plus entropy-scaled sparse n-gram biases and other context contributions.Before the first SSM forward pass, the SSM logits are zero and the scale is 1.
  • Arithmetic coding: 32-bit range arithmetic coding uses a scale of 65,536, gives every token frequency 1, and maintains a minimum interval width of 32,768.The stated width remains above the 32-bit representability threshold.
  • Arithmetic coding: 0.6 bits/token is the approximate quantisation redundancy for ve = 44,298, relative to approximately 6.8 bits/token prediction error.The quantisation overhead is characterised as small relative to prediction error.
  • Implementation: AVX2 SIMD accelerates projections and Adam updates in pure C without Python, CUDA, BLAS, or external dependencies.The head projection dominates forward-pass cost because it loops over ve tokens.

5. Theoretical Analysis

The analysis interprets StateSMix as an adaptive-memory model whose components reduce different parts of prediction excess. It connects the design to PPM and PAQ while identifying scaling limits and the SSM’s dominant role.

  • Adaptive memory: Input-dependent time constants give the diagonal Mamba recurrence short memory for rapidly changing contexts and long memory for slowly varying contexts.Online training further differentiates time constants toward those useful for the compressed file.
  • Bayesian interpretation: Bayes’ rule in log space recovers the sparse n-gram logit update, with λ controlling likelihood strength and α controlling count smoothing.Larger λ sharpens the posterior, while smaller α makes count evidence more influential.
  • Information-theoretic view: SSM, n-gram, and specialised predictors target global patterns, frequent exact contexts, and highly specific repetitions, respectively.The analysis assigns the SSM a dominant role early and n-gram contributions especially after approximately 100K tokens.
  • Scaling analysis: The collision-free fraction is approximately min(1, M/Nk), with sublinear distinct-context growth predicted by Zipf’s law.The expression relates table capacity M to the number of unique k-gram contexts Nk.
  • Scaling analysis: Below approximately 30 MB the tables are lightly loaded and competitive with LZMA, whereas beyond approximately 100 MB saturation limits further gains.The passage reports lower load for higher-order tables because longer contexts are exponentially rarer.
  • Classical connections: StateSMix combines a neural background model with bounded PPM-like context refinement and a simplified PAQ-like confidence-weighted mixer.Unlike PPM escape probabilities or a learned PAQ mixer, it uses fixed λk values and analytically determined entropy scaling.
  • Empirical implication: 2.158 bpb versus 3.568 bpb is reported for SSM-only versus n-grams-only compression on enwik83M.The ablation supports the SSM as the stronger compression component in this comparison.

6. Experiments

StateSMix is evaluated on enwik8 excerpts, component ablations, online compression progress, runtime, memory, and per-order n-gram behavior. It outperforms xz through 10 MB, with the SSM providing the dominant gain and n-grams adding complementary improvements, while longer runs face saturation and collision limits.

  • Main Results: StateSMix consistently outperforms xz on enwik8 file sizes up to 10 MB, with its advantage largest at 1 MB and crossing over at approximately 30 MB.The reported advantage is −8.6% at 1 MB and decreases with file size.
  • Main Results: The 100 MB comparison positions StateSMix below the best neural compressors but above classical methods without pre-training or GPU hardware.Its competitive per-file performance extends up to 10 MB, while NNCP incurs stored-weight overhead on short files.
  • Ablation Study: The SSM is the dominant component: removing it increases output size by +95%, whereas the SSM alone reduces size by 46.6% over count-only and beats xz by 1.3%.The SSM-only variant produces 840 KB on enwik8 3M.
  • Ablation Study: N-grams add a complementary 4.1% reduction on top of the SSM, yielding 806 KB and placing the full system 5.4% below xz on enwik8 3M.The 16-gram and 32-gram tables contribute approximately 2 KB by capturing repeated multi-token patterns beyond the 8-gram window.
  • Compression Progress: Online performance improves from approximately 8.1 bpt at initialization to 7.1 bpt after 50K tokens and approximately 6.83 bpt by 3M tokens, then improves marginally beyond 10M tokens.The reported causes of the later plateau are saturated SSM learning and n-gram table collisions.
  • Speed and Memory: OpenMP parallelisation yields a 1.9× speedup on 4 cores, while the nine n-gram hash tables dominate the reported 6.1 GB memory use.Online training accounts for approximately 75% of runtime, and actual memory use is proportional to hash-table load because untouched pages are not physically allocated.
  • Per-order N-gram Contribution: Per-order statistics show a trade-off: bigrams have over 99% hit rate with approximately 8 continuations, whereas eightgrams have sub-5% hit rate with approximately 2 continuations.Removing any order k ≥3 increases bpb by approximately 0.1–0.3%, while removing bigrams increases it by approximately 0.5%.

7. Discussion

StateSMix combines an online-trained SSM with sparse n-gram memorisation, outperforming xz on smaller natural-language files but losing its advantage beyond roughly 30 MB.

  • Why xz wins at 100 MB: At roughly 100 MB, xz wins because it copies repeated multi-KB blocks, whereas StateSMix encodes tokens individually through probability boosting.Table saturation accounts for only ∼13 KB of the 1.76 MB gap, so the main boundary is architectural.
  • The role of the SSM as compression core: The SSM alone beats xz on enwik83M, while n-gram tables add exact memorisation of frequent local transitions.The SSM uses online backpropagation; n-grams complement its generalisation.
  • Comparison with NNCP: StateSMix’s advantages over NNCP are self-contained output, portability through pure C and no GPU, and better performance than xz up to 10 MB.NNCP achieves ∼1.19 bpb on enwik8 but uses a much larger Transformer whose weights are stored in the compressed output.
  • Practical niche: StateSMix occupies a niche for individual natural-language files up to ∼10 MB when GPUs and large pre-trained models are unavailable.Suggested applications include embedded systems, encrypted backups, and bandwidth-constrained environments.
  • Limitations: The main limitations are ∼700 KB/s speed on 4 cores, ∼6.1 GB RAM for 100 MB input, and xz superiority beyond ∼30 MB.OpenMP provides 1.9× speedup on 4 cores, but further gains require fewer training iterations or GPU offloading.

8. Future Work

Future work targets long-range repetition, faster training, adaptive context weighting, stronger initialization, and a copy channel to address StateSMix’s large-file limitations.

  • BWT preprocessing: BWT preprocessing could cluster identical contexts and amplify n-gram effectiveness for long-range patterns.The motivation is to improve competitiveness with LZMA on large corpora.
  • GPU acceleration: GPU acceleration could reduce per-token cost by 50–100× and enable larger SSMs with DM = 128 and NL = 4.The head projection and Adam update are described as embarrassingly parallel.
  • Adaptive n-gram weighting: An exponential-weights meta-learner could replace fixed λk values with online per-order adaptation to each file’s statistical signature.This proposal follows PAQ-style adaptive context mixing.
  • Variable-order back-off: Variable-order back-off could use the longest matching context with confidence-weighted fallback to lower orders.The proposed design is intended to improve accuracy when high-order contexts are reliable.
  • Larger SSM with pre-trained initialisation: Pre-training the SSM on a reference corpus could provide better initialization for per-file fine-tuning, while distributing the weights as a codec file.This would combine pre-trained and online-compression approaches.
  • LZ match channel: A predict-or-copy architecture could encode long repeated token sequences as offset–length pairs, bypassing arithmetic coding.The proposal addresses the gap between per-token encoding and xz’s block copying; table saturation was not the bottleneck.

9. Conclusion

StateSMix is a self-contained compressor built from an online Mamba SSM and sparse n-gram biasing. It substantially improves over xz at moderate file sizes, with the SSM providing most of the gain and performance crossing over beyond roughly 30 MB.

  • Contribution: StateSMix combines an online-trained Mamba SSM with sparse n-gram logit biasing and arithmetic coding without pre-trained weights, a GPU, or external dependencies.The model is trained online during compression.
  • Ablation: 46.6% size reduction over a frequency-count baseline comes from the SSM alone, which also beats xz by −1.3% without n-grams.The ablation identifies the SSM as the primary compression engine.
  • Ablation: The n-gram tables provide a consistent additional 4.1% gain by memorising exact local and long-range transitions.The cited contexts include 16-gram and 32-gram matching.
  • Benchmark outcome: StateSMix beats xz by 8.7% on 1 MB, 5.4% on 3 MB, and 0.7% on 10 MB natural-language text.The comparison is reported for moderate file sizes.
  • Scientific implication: The SSM demonstrates rapid online adaptation to file-specific patterns despite a tiny parameter budget, eventually outperforming a mature classical compressor.This conclusion concerns sequential learning as well as compression.
Loading 2605.02904v1…