torchtsetlin: Fast, Interpretable Tsetlin Machines for the PyTorch Ecosystem

Models, data preparation, training, evaluation, interpretation and visualisation in one package, GPU-enabled and PyTorch-native.

A Tsetlin machine learns propositional rules (IF x0 AND NOT x1 THEN class 1) with teams of learning automata instead of gradients. torchtsetlin implements the family as ordinary torch.nn.Modules: the automata live in tensors, so a model moves between CPU and GPU with .to(device), is saved with state_dict(), and learns from mini-batches with a single model.update(x, y) call.

import torchtsetlin as tt

x_train, y_train = tt.data.make_noisy_xor(5000, noise=0.4)          # Boolean features, int labels
x_test, y_test = tt.data.make_noisy_xor(2000, noise=0.0, seed=1)

model = tt.TsetlinMachine(n_features=12, n_classes=2, n_clauses=20, T=15, s=3.9).to("cuda")

# PyTorch-style loop: update() is the analogue of loss.backward(); optimizer.step()
for epoch in range(50):
    model.train()
    for i in range(0, 5000, 10):
        model.update(x_train[i:i+10].cuda(), y_train[i:i+10].cuda())
    model.eval()
    acc = (model(x_test.cuda()).argmax(1) == y_test.cuda()).float().mean()

print(model.rules()[:2])   # ['IF x0 AND NOT x1 THEN 1', 'IF x1 AND NOT x0 THEN 1']

Or with the built-in trainer:

trainer = tt.Trainer(model, batch_size=10, callbacks=[tt.EarlyStopping(patience=20)])
trainer.fit((x_train, y_train), epochs=100, val_data=(x_test, y_test))
trainer.evaluate((x_test, y_test))          # {'accuracy': 0.99}

Features

Models	TsetlinMachine (multi-class, optional integer clause weights), CoalescedTsetlinMachine (shared clauses; multi-class or multi-label), RegressionTsetlinMachine, ConvTsetlinMachine / Conv1dTsetlinMachine / Conv3dTsetlinMachine (volumes) and convolutional coalesced / regression variants, SegmentationTsetlinMachine / CoalescedSegmentationTsetlinMachine for dense per-pixel prediction
Learning	vote margin T, specificity s, memory depth, boosted true-positive feedback, clause-size constraint, drop-clause / drop-literal, focused negative sampling, batched (GPU-friendly) or exact sequential feedback
Data	thermometer, one-hot, bit-plane, adaptive-threshold, colour-thermometer, hypervector and multi-scale pyramid encoders; synthetic datasets (including segmentation scenes and 3D volumes); torchvision helpers for classification and segmentation
Training	Trainer for tensors / Dataset / DataLoader, callbacks (early stopping, checkpoints, CSV logs, hyper-parameter schedules), History
Evaluation	accuracy, confusion matrix, precision/recall/F1, regression and multi-label metrics, IoU / Dice / pixel accuracy / boundary-F1 for segmentation, calibration, trustworthiness curves
Interpretation	rule extraction, clause activity/precision, closed-form global & local feature importance, per-example and per-pixel explanations
Visualisation	memory plots, automata heat-maps, convolutional clause patches, segmentation overlays and vote maps, confusion matrices, vote distributions

Installation

torchtsetlin is published on PyPI: https://pypi.org/project/torchtsetlin/. It needs Python 3.9+ and PyTorch 2.0+.

1. Install PyTorch first, choosing the CUDA or CPU build for your machine from pytorch.org, e.g.

   pip install torch                                                     # default build
   pip install torch --index-url https://download.pytorch.org/whl/cpu    # CPU-only
   

2. Install torchtsetlin from PyPI:

   pip install torchtsetlin
   

3. Optional extras:

   pip install "torchtsetlin[viz]"      # matplotlib for torchtsetlin.viz
   pip install "torchtsetlin[vision]"   # torchvision for the MNIST / torchvision / segmentation dataset helpers
   pip install "torchtsetlin[sklearn]"  # scikit-learn interop in the examples
   pip install "torchtsetlin[all]"      # everything, including docs and dev tools
   

Check the installation:

import torch, torchtsetlin as tt
print(tt.__version__, torch.cuda.is_available())

From source

git clone https://github.com/vlbthambawita/torchtsetlin
cd torchtsetlin
pip install -e ".[dev]"
pytest            # run the test-suite (uses the GPU when available)
mkdocs serve      # browse the documentation locally

How learning works (in one paragraph)

Each clause is an AND of literals (x_k or NOT x_k); every (clause, literal) pair has a Tsetlin automaton whose integer state decides whether the literal is included. For a training example, clauses of the true class receive Type I feedback (memorise the example's True literals, forget the rest) and clauses of a random other class receive Type II feedback (add a False literal so the clause stops matching), each with a probability controlled by the vote margin T. torchtsetlin evaluates a mini-batch with one matrix product, counts the feedback events per (clause, literal) with another, and turns the counts into binomial state transitions — the whole update is a handful of tensor ops. See the concept pages for details and references.

Documentation

• Getting started, concepts, guides, examples and the full API reference: https://vlbthambawita.github.io/torchtsetlin/ (or mkdocs serve locally).
• Measured CPU vs GPU performance (throughput, scaling, memory, where the time goes): Benchmarks -> CPU vs GPU; the harness lives in benchmarks/.
• Runnable scripts in examples/ (see its README): Noisy XOR, tabular data with thermometer encoding, regression, MNIST flat and convolutional, four convolutional walkthroughs covering 2-D clauses (shapes_conv.py), 1-D signals (conv1d_ramps.py), convolutional regression (conv_regression_blobs.py) and shared clause pools (mnist_conv.py --coalesced), and semantic segmentation on generated street scenes (segmentation_scenes.py).
• Worked notebooks in examples/notebooks/: Iris (rules you can read), MNIST (convolutional clauses and GPU throughput) and CIFAR-10 (booleanizing colour images), plus two that work through the Tsetlin-machine segmentation literature — the Convolutional Regression TM (ICMLT 2021) and CTM-UNet (ISTM 2025), the latter building a dense per-pixel Tsetlin segmentation head.

Citation

If you use torchtsetlin in research, please cite the five-page technical report, together with the Tsetlin machine papers whose algorithms you rely on (see the documentation). The PDF and other formats are at https://vajira.info/torchtsetlin/cite/, and GitHub's "Cite this repository" button reads CITATION.cff.

@techreport{thambawita2026torchtsetlin,
  author      = {Thambawita, Vajira},
  title       = {torchtsetlin: Fast, Interpretable Tsetlin Machines for the PyTorch Ecosystem},
  institution = {Simula Metropolitan Center for Digital Engineering},
  type        = {Technical report},
  year        = {2026},
  note        = {Software version 0.3.2},
  url         = {https://github.com/vlbthambawita/torchtsetlin}
}

License

MIT — see LICENSE.