code/Taotern_SSM/README.md

Gamma Space Model

Gamma Space Model is a PyTorch codebase for experimenting with a Gamma-structured state space model (SSM) and an S4-inspired enhanced Gamma SSM while preserving a fixed lower-bidiagonal ternary-friendly transition matrix.

This repository is now organized as an installable Python package:

1. gamma_space_model/ contains the models we are actively developing.
2. csrc/ contains optional pure-Python acceleration helpers used when the runtime supports them.
3. output/jupyter-notebook/ contains the benchmark notebooks and saved Colab runs.

Earlier versions of this repository included a copied s4-main/ reference tree. That vendored copy has been removed so the project can behave like a normal package. S4 remains an external theory/reference dependency, cited below, rather than source code we ship inside this repo.

The current project direction is:

• keep the Gamma / ternary transition structure for deployment compatibility
• borrow the most useful ideas from S4 for stability and full-sequence efficiency
• benchmark both training-time full-sequence behavior and deployment-time recurrent behavior

Repository Layout

gamma_ssm_s4_v2/
|-- gamma_space_model/
|   |-- modules/
|   |   |-- ssm_gamma.py        # original Gamma SSM core
|   |   |-- block.py            # original residual Gamma block
|   |   |-- ssm_gamma_s4.py     # S4-inspired enhanced Gamma SSM core
|   |   |-- block_s4.py         # enhanced Gamma blocks
|   |   `-- normalization.py    # LayerNorm, RMSNorm
|   |-- ops/
|   |   `-- selective_scan_interface.py
|   `-- __init__.py
|-- csrc/
|   `-- tilelang/                 # optional acceleration helpers
|-- output/jupyter-notebook/
|   |-- gamma-s4-sinewave-benchmark.ipynb
|   |-- gamma-s4-research-benchmark.ipynb
|   |-- gamma-s4-challenge-benchmark.ipynb
|   `-- *_rN.ipynb              # saved Colab benchmark runs
|-- scripts/
|   |-- generate_gamma_benchmark_notebook.py
|   `-- generate_gamma_challenge_benchmark_notebook.py
|-- tests/
|-- EXPERIMENT_RECORD.md        # run-by-run experiment history
|-- pyproject.toml              # package metadata for pip install
|-- setup.py                    # compatibility shim for editable installs
`-- Gamma Distributed Ternary HiPPO.pdf

What We Are Modeling

At the highest level, all variants in this repo are state space models. The standard continuous-time form is

\dot{h}(t) = A h(t) + B u(t), \qquad y(t) = C h(t) + D u(t)

where:

• u(t) is the input sequence
• h(t) is the latent state
• y(t) is the output sequence
• A controls the state dynamics
• B injects the input into the state
• C reads the state back out
• D is a direct skip term

The main design choice in this project is to keep A structured and deployment-friendly.

The Gamma Transition Matrix

The original Gamma SSM in this repo uses a fixed lower-bidiagonal matrix:

A_{n,n} = -1, \qquad A_{n,n-1} = 1

and zero everywhere else.

In matrix form:

[-1  0  0  ...]
[ 1 -1  0  ...]
[ 0  1 -1  ...]
[ .  .  .   . ]

This is important for the project because:

• it is sparse
• it only uses the values -1, 0, and 1
• it maps naturally to the ternary computing direction we care about

That structural constraint is the main reason we do not simply replace Gamma with a standard S4 parameterization.

The Original Gamma SSM

The original implementation is in gamma_space_model/modules/ssm_gamma.py.

It uses Euler discretization:

h_{t+1} = h_t + \Delta t \, (A h_t + B u_t)

and output readout:

y_t = C h_{t+1}

Key properties of the baseline:

• fixed Gamma A
• learned B and C
• scalar fixed delta_t
• recurrent forward pass
• optional TileLang/Triton accelerated scan path when available

This version is the simplest and most deployment-aligned baseline in the repo.

The S4-Inspired Enhanced Gamma SSM

The enhanced implementation is in gamma_space_model/modules/ssm_gamma_s4.py.

The goal is not to become full S4. The goal is to keep the Gamma transition structure while borrowing ideas that made S4 powerful and stable:

• learned positive timestep dt
• stable discretization
• optional direct skip term D
• full-sequence kernel view for parallel training/inference
• recurrent stepping for deployment
• Mamba-inspired input selection before the Gamma SSM

Discretization

The enhanced model supports three discretizations:

Euler

\bar{A} = I + \Delta t A, \qquad \bar{B} = \Delta t B

Bilinear (Tustin)

\bar{A} = (I - \tfrac{1}{2}\Delta t A)^{-1}(I + \tfrac{1}{2}\Delta t A)

\bar{B} = (I - \tfrac{1}{2}\Delta t A)^{-1}\Delta t B

ZOH

\bar{A} = e^{\Delta t A}

\bar{B} = A^{-1}(\bar{A} - I)B

The default practical choice is bilinear discretization, because it has been the most stable and effective in our current experiments.

Kernel View

For a linear time-invariant discrete SSM, the output can also be written as a causal convolution:

y_t = \sum_{\ell=0}^{t} K_{\ell} u_{t-\ell} + D u_t

with kernel

K_{\ell} = C \bar{A}^{\ell} \bar{B}

This gives us two useful execution styles for the same model:

• recurrent stepping: useful for deployment and streaming
• full-sequence kernel/convolution view: useful for parallel whole-sequence computation

In this repo, the enhanced model exposes:

• kernel_mode="recurrent"
• kernel_mode="conv"
• kernel_mode="auto"

The auto mode switches to the kernel path only when the sequence is long enough to justify it.

Why This Is Not "Full S4"

This is the central design tradeoff of the project.

S4 gets its strongest speedups from a structured parameterization of A that can be diagonalized or reduced to efficient kernel computations such as Cauchy/Vandermonde operations.

Our Gamma model intentionally keeps:

• a fixed lower-bidiagonal A
• ternary-friendly values
• deployment-oriented structure

So we borrow selected S4 ideas, but we do not inherit the entire original S4 kernel machinery.

That means:

• we can improve stability and full-sequence performance substantially
• but the exact original S4 fast-kernel theory does not transfer directly

Block Design

Original Block

gamma_space_model/modules/block.py defines GammaSingleBlock.

Its structure is:

x
-> LayerNorm (if prenorm)
-> Gamma SSM
-> Dropout
-> Residual add
-> optional postnorm

Enhanced Block

gamma_space_model/modules/block_s4.py defines GammaS4Block.

Its structure is:

x
-> LayerNorm
-> optional input-selection gate
-> SSMGammaS4
-> activation
-> optional gate
-> optional output linear
-> layer scale
-> dropout
-> residual add

In equations, the enhanced block is roughly:

\tilde{x} = \mathrm{Norm}(x)

u = \tilde{x} \odot \sigma(W_{in\_gate}\tilde{x} + b_{in\_gate})

s = \mathrm{SSMGammaS4}(u)

z = \mathrm{OutputLinear}(\sigma(s) \odot \mathrm{Gate}(\tilde{x}))

\mathrm{Block}(x) = x + \alpha z

where:

• sigma is the chosen activation
• the input-selection gate is a per-token channel gate inspired by Mamba's selective input flow
• Gate may be absent
• alpha is a learnable layer-scale parameter

The input-selection gate does not change the fixed Gamma transition matrix A. It modulates the input before it enters the state, so the model can learn to emphasize information-bearing tokens and suppress blanks/noise while keeping the ternary-friendly state transition intact.

There is also a lighter variant:

• GammaS4MinimalBlock

which removes the richer input-gating/output-gating/output-linear pathway and keeps only the core S4-inspired stability changes.

Stacked Model Pattern

The benchmark notebooks use a simple stacked forecaster pattern:

x_0 = W_{in} u

x_{\ell+1} = \mathrm{Block}_{\ell}(x_{\ell})

\hat{y} = W_{out} x_L

This pattern is implemented inside the notebooks rather than as a dedicated package module because we are still iterating rapidly on benchmark design.

Full-Sequence vs Recurrent vs Deployment-Lite

The benchmarks report several distinct execution modes.

1. Full-sequence

The whole sequence is passed to forward(...) in one call.

This is the relevant mode for:

• standard training
• offline batch inference
• full-sequence throughput comparisons

Important detail:

• full-sequence does not always mean convolution is being used
• it only means the whole sequence is evaluated in one forward pass

2. Recurrent

The model is stepped token by token using step(...) while carrying state.

This is the relevant mode for:

• streaming inference
• autoregressive deployment
• hardware-style stateful execution

3. Deployment-lite

This is an experimental recurrent inference mode for GammaS4Block.

It uses the same trained weights but simplifies part of the block-time recurrent computation to reduce runtime. The current lite path is intentionally aggressive: it skips the input-selection gate and post-SSM gate/output branch during recurrent stepping. The benchmark reports both:

• speed improvement
• output mismatch versus the standard full-sequence prediction

For the baseline model, deployment and recurrent are effectively the same path, so those numbers match exactly.

4. Balanced deployment

This is a middle-ground recurrent inference mode for GammaS4Block.

It keeps the trained output projection, but replaces the input-selection and post-SSM gates with static gates derived from the learned gate biases. This is meant to test whether we can recover more fidelity than deployment-lite while still avoiding the full gate cost at every token.

In the notebooks, these metrics appear as balanced_deploy_*.

Current Empirical Status

The benchmark notebooks live in:

• output/jupyter-notebook/gamma-s4-sinewave-benchmark.ipynb
• output/jupyter-notebook/gamma-s4-research-benchmark.ipynb
• output/jupyter-notebook/gamma-s4-challenge-benchmark.ipynb

The challenge notebook is separate from the quick and research notebooks. It tracks harder capability-style tests:

• permuted MNIST for long-range image-as-sequence memory
• selective copying for sparse content recall
• induction-style recall for key-value association across a sequence
• token-memory curriculum tiers for easy, moderate, and hard selective/induction diagnostics

The research and challenge notebooks also include inference-oriented sections. These separate full-sequence prefill-style inference from recurrent decode-style inference, report deployment-lite and balanced deployment variants where available, and include quality/fidelity columns so speed is not interpreted independently from accuracy or output mismatch.

Saved Colab runs are also committed in the same folder using _rN suffixes.

Experiment record guide

The run-by-run history is summarized in EXPERIMENT_RECORD.md.

Use that file when you want to answer:

• what changed between _r1, _r2, ..., _r10
• which model variants were tested in each run
• which notebook and task configuration produced each result
• which results are complete, partial, or not recorded
• what we learned from each run

The saved notebooks under output/jupyter-notebook/ are the raw experiment artifacts. They include the actual Colab outputs, plots, printed metrics, and configuration cells. The experiment record is the human-readable index over those artifacts.

The naming convention is:

• gamma_s4_sinewave_benchmark_r1.ipynb and gamma_s4_sinewave_benchmark_r2.ipynb
  • early single-notebook experiments before the quick/research split
• gamma-s4-sinewave-benchmark_rN.ipynb
  • saved quick benchmark runs
  • use these for fast regression history
• gamma-s4-research-benchmark_rN.ipynb
  • saved practical/research benchmark runs
  • use these for presentation and deeper analysis
• gamma-s4-challenge-benchmark_rN.ipynb
  • saved challenge benchmark runs
  • use these for permuted MNIST, selective copying, and induction-style recall history

For the current project state, the most useful records are:

• _r9 research benchmark: best presentation artifact
• _r10 quick/research benchmarks: latest deployment fidelity comparison with balanced deployment metrics
• _r8 research benchmark: first mature run with token-lite enabled and strong long-context conv results
• _r2: early evidence that gamma_s4_enhanced could outperform baseline on harder sequence tasks

When presenting to others, start with the README for theory and architecture, then use EXPERIMENT_RECORD.md for the experiment timeline, then open the latest _rN research notebook for plots and raw outputs.

What the quick benchmark is for

The sinewave/quick notebook is the fast regression loop. It is used to answer:

• does the model still train correctly?
• does full-sequence performance remain strong?
• does recurrent/deployment behavior regress?

What the research benchmark is for

The research notebook is the more presentable and more practical benchmark. It currently includes:

• current_reference: medium practical forecasting task
• long_context: harder long-range forecasting task
• token-lite: lightweight character-level next-token benchmark

Latest recorded highlights

From the latest committed _r10 GPU runs:

• Quick benchmark:

  • enhanced val loss is much lower than baseline on both simple and moderate
  • enhanced full-sequence throughput is higher than baseline
  • enhanced recurrent inference is still slower than baseline

• Research benchmark:

  • current_reference
    • enhanced validation loss: 0.019951
    • baseline validation loss: 0.709749
  • long_context
    • enhanced validation loss: 0.011708
    • baseline validation loss: 27.229956
    • enhanced mean epoch time: 15.86s
    • baseline mean epoch time: 40.35s
    • enhanced full-sequence throughput: 16957 tokens/s
    • baseline full-sequence throughput: 2395 tokens/s
    • enhanced balanced deployment match MSE: 0.001692
    • enhanced deployment-lite match MSE: 0.200325
  • token-lite
    • enhanced validation CE: 2.4868
    • baseline validation CE: 3.1322
    • enhanced perplexity: 12.02
    • baseline perplexity: 22.92

The challenge notebook has been added after _r10, so its first saved run should be treated as the first challenge-task record.

Practical interpretation

The current picture is:

• the enhanced Gamma SSM is now clearly the stronger model for full-sequence training and harder tasks
• the long-context conv/full-sequence path is finally showing meaningful advantages
• recurrent deployment remains the main remaining weakness
• challenge recall tasks remain near random, so the next question is whether the model can solve easier curriculum tiers before investing in harder selective-memory variants

Installation

The distribution package is named gamma-ssm-s4-enhanced; the Python import package remains gamma_space_model.

Install From This Checkout

pip install -e .

Install From GitHub

pip install "git+https://github.com/StarMists/gamma_SSM_S4_enhanced.git"

Install From A Private GitHub Repo

Use a GitHub personal access token with read access to the private repository. In Colab or a shell, prefer keeping the token in an environment variable instead of hard-coding it into notebooks:

export GITHUB_TOKEN="ghp_your_token_here"
pip install "git+https://${GITHUB_TOKEN}@github.com/StarMists/gamma_SSM_S4_enhanced.git"

In Google Colab:

import os

os.environ["GITHUB_TOKEN"] = "ghp_your_token_here"
!pip install "git+https://${GITHUB_TOKEN}@github.com/StarMists/gamma_SSM_S4_enhanced.git"

Install Optional Extras

For development:

pip install -e ".[dev]"

For benchmark notebooks:

pip install -e ".[notebook]"

For optional performance dependencies:

pip install -e ".[performance]"

For a private GitHub install with extras, use the PEP 508 form:

pip install "gamma-ssm-s4-enhanced[notebook,performance] @ git+https://${GITHUB_TOKEN}@github.com/StarMists/gamma_SSM_S4_enhanced.git"

Notes:

• CUDA / Triton / TileLang acceleration is optional
• the code automatically falls back to pure PyTorch paths when those kernels are unavailable
• installing from GitHub gives downstream LLM code access to the latest committed package without cloning the repo manually

Quick Start

Original Gamma block

import torch
from gamma_space_model import GammaSingleBlock

block = GammaSingleBlock(
    d_model=64,
    hidden_dim=128,
    delta_t=0.1,
    prenorm=True,
)

x = torch.randn(2, 128, 64)
y, h_T = block(x)
print(y.shape, h_T.shape)

Enhanced Gamma block

import torch
from gamma_space_model import GammaS4Block

block = GammaS4Block(
    d_model=64,
    hidden_dim=128,
    discretization="bilinear",
    kernel_mode="auto",
    kernel_threshold=384,
    input_gate=True,
    gate=True,
    use_D=True,
)

x = torch.randn(2, 512, 64)
y, h_T = block(x)
print(y.shape, h_T.shape)

Export discretized matrices for deployment

from gamma_space_model import SSMGammaS4

ssm = SSMGammaS4(state_dim=64, hidden_dim=128)
deployment_mats = ssm.export_inference_matrices()
print(deployment_mats.keys())

Tests

The main test files are:

• tests/test_ssm_gamma.py
• tests/test_block.py
• tests/test_ssm_gamma_s4.py

These cover:

• forward and step behavior
• state initialization
• recurrent/full-sequence agreement
• enhanced conv vs recurrent consistency
• deployment cache paths

Roadmap / Next Improvement Points

These are the main next-step items identified from the latest benchmark runs. They are intentionally listed as a to-do record, not yet implemented in this README update.

1. Improve enhanced recurrent inference further.

   • Full-sequence behavior is now strong.
   • Recurrent deployment remains the largest performance gap.

2. Improve deployment-lite fidelity.

   • It speeds up recurrent inference, but its output mismatch on long_context is still too large.

3. Improve structured kernel generation further.

   • The current conv path is much better than before, but it still uses a direct kernel-building loop rather than the strongest possible structured kernel derivation.

4. Strengthen the token benchmark.

   • The current token-lite task is useful for relative comparison, but still too small to serve as a final language-model benchmark.

5. Add more presentation-oriented result summaries.

   • The research notebook now includes task previews and prediction/error plots.
   • A future pass could add automated summary tables or figure exports for slide decks.

References

Internal project note

• Gamma Distributed Ternary HiPPO.pdf in the repository root

External references

• Albert Gu, Karan Goel, Christopher Re. Efficiently Modeling Long Sequences with Structured State Spaces. ICLR 2022.

  • arXiv
  • official S4 repository

• Albert Gu, Tri Dao, Stefano Ermon, Atri Rudra, Christopher Re. HiPPO: Recurrent Memory with Optimal Polynomial Projections. NeurIPS 2020.

  • arXiv

• Albert Gu, Tri Dao. Mamba: Linear-Time Sequence Modeling with Selective State Spaces. 2023.

  • arXiv

Citation

Yes, adding citations is good practice here. This repo builds directly on ideas from HiPPO and S4, and the README should make that explicit.

If you use this repository in academic or technical writing, please cite the upstream S4, HiPPO, and Mamba papers above, and mention that this codebase studies a Gamma-structured SSM with S4-inspired enhancements.

Example BibTeX entries:

@inproceedings{gu2022s4,
  title={Efficiently Modeling Long Sequences with Structured State Spaces},
  author={Gu, Albert and Goel, Karan and Re, Christopher},
  booktitle={International Conference on Learning Representations},
  year={2022}
}

@inproceedings{gu2020hippo,
  title={HiPPO: Recurrent Memory with Optimal Polynomial Projections},
  author={Gu, Albert and Dao, Tri and Ermon, Stefano and Rudra, Atri and Re, Christopher},
  booktitle={Advances in Neural Information Processing Systems},
  year={2020}
}

@article{gu2023mamba,
  title={Mamba: Linear-Time Sequence Modeling with Selective State Spaces},
  author={Gu, Albert and Dao, Tri},
  journal={arXiv preprint arXiv:2312.00752},
  year={2023}
}