tmu-rs

A Rust port of the cair/tmu Tsetlin Machine library.

Implements the core Tsetlin Machine variants — multiclass classifier, coalesced classifier, regressor, convolutional (1-D and 2-D), autoencoder, composite classifier, and a sparse classifier with absorbing actions — with bit-packed clause banks, bit-parallel training, optional Rayon multi-threading, and a fast type-safe booleanizer.

For a full breakdown of what has been ported and what is missing, see PORTING_STATUS.md. For a Python vs Rust throughput and accuracy comparison, see BENCHMARKS.md.

---

Use as a library dependency

Add to your project's Cargo.toml:

[dependencies]
tmu-rs = { git = "https://github.com/ozzykampha/tmu-rust" }

# Optional: pin to a specific tag for reproducible builds
# tmu-rs = { git = "https://github.com/ozzykampha/tmu-rust", tag = "v1.0.0" }

# Optional: enable multi-threaded training
# tmu-rs = { git = "https://github.com/ozzykampha/tmu-rust", features = ["parallel"] }

Then use it:

use tmu_rs::{TsetlinMachine, Encoder};

// Build encoder from training data (binary features in this example)
let encoder = Encoder::binary(n_features);
let train_x = encoder.encode_batch(&raw_train_x);

// Create and train the classifier
let mut tm = TsetlinMachine::with_config(
    n_classes, clauses_per_class, n_features,
    threshold, specificity, max_states, boost_true_positive, seed,
);
for _ in 0..epochs {
    tm.fit_epoch(&train_x, &train_y);
}
let accuracy = tm.accuracy(&test_x, &test_y);

---

Features

• Bit-packed clause bank for cache-efficient inference and training
• Five model types:
  • TMClassifier — weighted multiclass classification
  • TMCoalescedClassifier — one shared clause bank with signed per-class weights
  • TMRegressor — continuous output from binary features
  • ConvolutionalTsetlinMachine — 1-D and 2-D sliding-window clause banks (weight-tied patches)
  • TMCompositeClassifier — ensemble of per-class Tsetlin Machines with independent clause banks
  • TMAutoEncoder — binary reconstruction via positive-only clause banks
  • TMSparseClassifier — sparse clause bank with absorbing actions: literals are permanently dropped from each clause as training converges, so memory and per-clause evaluation scale with the number of active literals (a big win in high-dimensional, sparsely-relevant feature spaces)
• Optional multi-threaded training via Rayon (--features parallel)
• AVX2 fast paths for clause update loops with runtime dispatch — u8 TA counters processed 32-wide (4× smaller working set vs u32; scalar fallback on non-AVX2 targets)
• Type-safe Encoder for binary, numeric (quantile booleanization), and categorical inputs
• Fast booleanizer for continuous-valued inputs
• Save/load trained models and encoders to disk via the SaveLoad trait (on by default through the serde feature) — serde + bincode preserve all learned state and RNG streams, so a reloaded model predicts identically and can resume training deterministically. Build with --no-default-features for a dependency-free build without save/load
• Ports of the core TMU demos including classification, regression, convolution, autoencoder, and composite

---

Getting started

git clone --recurse-submodules https://github.com/OzzyKampha/tmu-rust.git
cd tmu-rust
cargo build --release

Run a self-contained example:

cargo run --release --example noisy_xor

For multi-threaded training:

cargo run --release --features parallel --example mnist

For maximum performance, compile with native CPU optimizations (enables AVX2 and other extensions at compile time):

RUSTFLAGS="-C target-cpu=native" cargo build --release

AVX2 fast paths are also activated at runtime automatically when the CPU supports it, even without target-cpu=native.

---

Examples

The examples reproduce the cair/tmu demos with matching hyperparameters (e.g. MNIST: 2000 clauses, T=50, s=10.0; IMDb: 2000 clauses, T=80, s=10.0). See PORTING_STATUS.md for the full status.

Classification demos

TMU demo	Example	Data required	Command
XORDemo	xor	—	cargo run --release --example xor
NoisyXORDemo	noisy_xor	—	cargo run --release --example noisy_xor
InterpretabilityDemo	interpretability	—	cargo run --release --example interpretability
TMCoalescedClassifier	coalesced	—	cargo run --release --example coalesced
TMSparseClassifier	sparse	—	cargo run --release --example sparse
BreastCancerDemo	breast_cancer	scikit-learn	see Data preparation
MNISTDemo / MNISTDemoWeightedClauses	mnist	MNIST	see Data preparation
IMDbTextCategorizationDemo	imdb	Keras IMDb	see Data preparation

New model types (run python scripts/gen_shared_data.py once first to generate shared data)

Model	Example	Command
TMRegressor	regression	cargo run --release --example regression
ConvolutionalTM 1-D	convolutional	cargo run --release --example convolutional
ConvolutionalTM 2-D	convolutional_2d	cargo run --release --example convolutional_2d
TMCompositeClassifier	composite	cargo run --release --example composite
TMAutoEncoder	autoencoder	cargo run --release --example autoencoder
TMCoalescedAutoEncoder	coalesced_autoencoder	cargo run --release --example coalesced_autoencoder

Extras

Example	Description
sparse_vs_dense	Dense vs sparse head-to-head: accuracy parity, memory footprint, train/inference time
save_load	Train → save → load → predict/resume round-trip
ndr_flows	Synthetic network-flow detection (booleanizer + rule extraction)
sysmon / sysmon_windows / sysmon_mordor	Sysmon event classification
bench_training	Training throughput benchmark (sequential vs parallel, IMDB-scale)
bench_autoencoder	AutoEncoder throughput + accuracy vs Python TMU
absorb_timing	Per-epoch accuracy and absorbing-state fraction at various state_bits

bench_training uses a synthetic dataset — no download required. Compare with and without --features parallel.

---

Data preparation

Three examples require datasets generated by the Python scripts in scripts/. Generated files are written to data/ and are not tracked by git.

Breast Cancer (requires scikit-learn):

pip install scikit-learn
python scripts/prepare_breast_cancer.py
cargo run --release --example breast_cancer

MNIST (requires tensorflow or scikit-learn):

python scripts/prepare_mnist.py
cargo run --release --features parallel --example mnist

IMDb (requires tensorflow):

python scripts/prepare_imdb.py
cargo run --release --features parallel --example imdb

Regression / Convolutional / Composite shared data (requires numpy):

pip install numpy
python scripts/gen_shared_data.py        # writes 14 binary files to data/
cargo run --release --example regression
cargo run --release --example convolutional
cargo run --release --example convolutional_2d
cargo run --release --example composite

---

Project structure

src/
  encoder.rs              # Type-safe Encoder (binary / numeric / categorical)
  booleanizer.rs          # Quantile booleanization (used by Encoder)
  clause_bank/            # Bit-packed clause storage and update logic
  models/
    classification/
      vanilla_classifier.rs      # TMClassifier
      coalesced_classifier.rs    # TMCoalescedClassifier
      convolutional_classifier.rs # ConvolutionalTsetlinMachine (1-D and 2-D)
      composite_classifier.rs    # TMCompositeClassifier
    regression/
      vanilla_regressor.rs       # TMRegressor
    autoencoder/
      vanilla_autoencoder.rs     # TMAutoEncoder
      coalesced_autoencoder.rs   # TMCoalescedAutoEncoder
  rng.rs                  # Fast SplitMix64 RNG
examples/               # Demo programs (ports of TMU + extras)
benches/                # Criterion throughput benchmarks
scripts/                # Python data preparation and comparison scripts
data/tmu/               # cair/tmu submodule (reference implementation)

---

License

MIT

Original TMU library: cair/tmu (MIT).