Graph Tsetlin machine — clauses that read a node and its neighbours¶
Paper. O.-C. Granmo et al., The Tsetlin Machine Goes Deep: Logical Learning and Reasoning With Graphs, 2025 (arXiv:2507.14874).
A flat Tsetlin machine reads one Boolean vector; a convolutional one reads every patch of an image and ORs the result. Both look at one thing at a time. The graph Tsetlin machine (GraphTM) reads a graph: each node carries a Boolean feature vector, and a clause is a deep clause, an AND over several layers. Layer 0 tests the node's own features. Every further layer tests messages: which clauses matched, one layer down, at the nodes that have an edge into this one. A depth-2 clause can therefore say "this node is an A, and an A sits on each side of it", something no single node knows.
This notebook covers:
- the paper's first example: does a sequence contain three consecutive A's?
- how sequences become graphs (
SymbolTable,GraphBatch); - why depth 1 cannot learn it and depth 2 can;
- reading the four learned clauses, and tracing one clause through the graph (
explain_node); - a harder sequence task: counting the length of a run;
- the GPU fast path for graphs (padding a dataset to one shape);
- images as graphs: at depth 1 a GraphTM is a convolutional machine, bit for bit.
Runs in a few minutes on a GPU (about a minute on an RTX 3090) and also on the CPU. In Colab, Runtime → Change runtime type → T4 GPU uses the CUDA fast paths.
# Setup: on Google Colab this installs torchtsetlin from PyPI (torch, torchvision, scikit-learn,
# matplotlib and pandas are already there). Elsewhere: `pip install torchtsetlin`, or
# `pip install -e ".[dev]"` in a checkout of the repository.
import sys
IN_COLAB = "google.colab" in sys.modules
if IN_COLAB:
%pip install -q "torchtsetlin>=0.7"
import time
import numpy as np
import pandas as pd
import torch
import torchtsetlin as tt
from torchtsetlin.data import GraphBatch, SymbolTable
SEED = 0
tt.seed_everything(SEED)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
DATA_ROOT = "data" if IN_COLAB else "../../data" # MNIST (section 7) downloads here
print("torchtsetlin", tt.__version__, "| torch", torch.__version__, "| device", device)
if device.type == "cuda":
print("gpu:", torch.cuda.get_device_name(0))
1. The task: three A's in a row¶
The paper's Section 3.1. Each example is a sequence of five letters from A–E; about a third of
them have AAA planted somewhere. The label says whether AAA occurs. Training labels get 1% noise,
as in the paper, and the test labels are clean.
Any single position is uninformative: an A occurs in most sequences, with or without the triple.
The answer lives in the relation between neighbouring positions.
def aaa_task(n, rng, noise=0.0):
seqs, ys = [], []
for _ in range(n):
s = [str(c) for c in rng.choice(list("ABCDE"), 5)]
if rng.random() < 1 / 3:
k = int(rng.integers(0, 3))
s[k:k + 3] = ["A"] * 3
y = int("AAA" in "".join(s))
if rng.random() < noise:
y = 1 - y
seqs.append(s)
ys.append(y)
return seqs, torch.tensor(ys)
rng = np.random.default_rng(SEED)
seq_train, y_train = aaa_task(32000, rng, noise=0.01)
seq_test, y_test = aaa_task(8000, rng)
for s, y in zip(seq_test[:6], y_test[:6]):
print("".join(s), "->", "AAA" if y else "no AAA")
print(f"\n{len(seq_train)} training / {len(seq_test)} test sequences, "
f"{float(y_test.float().mean()):.0%} of the test set contains AAA")
2. Sequences as graphs¶
Two objects do the encoding:
SymbolTablemaps symbols to node features. Withencoding="one_hot"every symbol owns one feature, so the clauses stay exactly readable. ("random"hypervectors keep the width fixed for large vocabularies.)GraphBatchholds a batch of graphs: node features(B, P, F), each graph's node count, and typed, directed edges.GraphBatch.from_sequencesmakes one node per position. Nodeksends tok − 1over edge type 0 (left) and tok + 1over type 1 (right).
symbols = SymbolTable(list("ABCDE"), encoding="one_hot")
g_train = GraphBatch.from_sequences(seq_train, symbols)
g_test = GraphBatch.from_sequences(seq_test, symbols)
print(g_train)
one = g_test[0]
print("\nsequence:", "".join(seq_test[0]))
print("node features (one row per position, columns A..E):")
print(one.node_features[0].int())
names = ["left", "right"]
print("edges (source -> destination, type):")
print(", ".join(f"{int(s)}->{int(d)} {names[int(t)]}" for (s, d), t in zip(one.edges, one.edge_type)))
A GraphBatch behaves like a tensor along the graph axis (len(g), g[3], g[2:10], .to(...)),
so the Trainer takes it wherever it takes a tensor.
Pad the dataset to one shape. Graphs may differ in size, so a batch is stored padded to its
largest graph. pad(P, E) fixes the node and edge counts for the whole dataset. On CUDA, every
slice then has the same shape without looking at the data, which lets updates run on the fused
kernels and be replayed from CUDA graphs (section 6). Padding nodes never match and never send, so
the results do not change.
g_train, g_test = g_train.pad(5, 8), g_test.pad(5, 8) # 5 positions, 4 left + 4 right edges
print(g_train)
3. Depth 1 cannot see the triple; depth 2 can¶
The paper's settings: 4 clauses, T = 40, s = 2, 4 epochs. Messages are one_hot: each
(clause, edge type) pair has its own inbox bit, which keeps every message readable.
At depth 1 a clause sees only its own node's letter. The machine then is a convolutional machine with a one-letter window, and the best it can do is a guess from letter counts. At depth 2, a clause can also require that given clauses matched at the left and right neighbours.
def train_graph_tm(depth, g_tr, y_tr, g_te, y_te, n_classes=2, n_clauses=4, T=40, s=2.0,
epochs=4, batch_size=25, **kwargs):
tt.seed_everything(SEED)
model = tt.GraphTsetlinMachine(
n_features=symbols.dim, n_outputs=n_classes, n_clauses=n_clauses, T=T, s=s,
depth=depth, message_encoding="one_hot",
n_edge_types=2, edge_type_names=["left", "right"], **kwargs,
)
trainer = tt.Trainer(model, device=device, batch_size=batch_size, verbose=False)
t0 = time.time()
history = trainer.fit((g_tr, y_tr), epochs=epochs, val_data=(g_te, y_te))
return model, history, time.time() - t0
rows, models = [], {}
for depth in (1, 2):
models[depth], hist, secs = train_graph_tm(depth, g_train, y_train, g_test, y_test)
rows.append({"depth": depth, **{f"epoch {i + 1}": round(a, 4)
for i, a in enumerate(hist["val_accuracy"])},
"seconds": round(secs, 1)})
pd.DataFrame(rows).set_index("depth")
Depth 1 stays well below 100%: one letter on its own cannot decide whether AAA occurs. Depth 2
solves the task, as the paper reports (100% test accuracy).
4. Reading the clauses¶
clause_expression prints a deep clause in the paper's notation. Layer-0 literals are symbols. A
message literal right1:3 reads "clause 3's layer-0 component matched at the node a right-going
edge came from", that is, at the left neighbour. (left1:3 is the right neighbour.) The weights
are the clause's votes for no AAA and AAA.
model = models[2]
pd.set_option("display.max_colwidth", None)
pd.DataFrame(
[{"clause": j, "votes (no AAA, AAA)": model.weights[j].tolist(),
"expression": model.clause_expression(j, symbols=symbols)} for j in range(model.n_clauses)]
).set_index("clause")
The two clauses voting for AAA (0 and 3) say exactly what the task is about: this node is an A, and at both neighbours clauses 0 and 3 (an A) and clause 2 (neither C nor E) matched, i.e. the node is the middle of a triple. The clauses voting for no AAA are less tidy; clause 1, for instance, needs a node that is neither A, B nor D and has no A-message from either side.
explain_node traces one clause through one graph: which node it holds at, and, for each message
literal, which neighbour sent it and why that neighbour's clause matched.
j = int(model.weights[:, 1].argmax()) # the clause voting most for AAA
i = int(torch.nonzero(y_test == 1)[0]) # a test sequence with AAA
print("sequence:", "".join(seq_test[i]), "\n")
print(model.explain_node(g_test, j, graph=i, symbols=symbols))
And on a sequence without the triple, the same clause holds at no node. Asking about a specific node shows which component fails:
k = int(torch.nonzero((y_test == 0) & torch.tensor(["A" in s for s in seq_test]))[0])
node = seq_test[k].index("A")
print("sequence:", "".join(seq_test[k]), f"| asking about node {node}\n")
print(model.explain_node(g_test, j, graph=k, symbols=symbols))
print()
print(model.explain_node(g_test, j, graph=k, node=node, symbols=symbols))
5. A harder task: how long is the run?¶
The reference implementation's sequence demo (also examples/graph_sequence.py). Sequences have
3–10 positions and a single run of 1, 2 or 3 consecutive A's; the class is the run length. Graphs
now have different sizes, so padding matters.
Two settings change from section 3:
- More clauses (40) for three classes and longer patterns.
- A literal budget (
max_included_literals=8). Early in training every clause matches everywhere and sends messages from every node, so one boosted Type Ia event can include all message literals at once. Without a budget, clauses can collapse into "everything matched around me", which holds in every graph.
def run_task(n, max_len=10, noise=0.0, seed=0):
g = torch.Generator().manual_seed(seed)
seqs, ys = [], []
for _ in range(n):
L = int(torch.randint(3, max_len + 1, (1,), generator=g))
y = int(torch.randint(0, 3, (1,), generator=g)) # run of y + 1 A's
end = int(torch.randint(y, L, (1,), generator=g))
s = [[] for _ in range(L)] # empty positions
for p in range(end - y, end + 1):
s[p] = ["A"]
if float(torch.rand(1, generator=g)) < noise:
y = int((y + torch.randint(1, 3, (1,), generator=g)) % 3)
seqs.append(s)
ys.append(y)
return seqs, torch.tensor(ys)
run_symbols = SymbolTable(["A"], encoding="one_hot")
r_train, ry_train = run_task(10000, noise=0.01, seed=0)
r_test, ry_test = run_task(2000, seed=1)
print("examples:", ["".join("A" if p else "." for p in s) + f" -> {int(y) + 1}"
for s, y in zip(r_test[:4], ry_test[:4])])
rg_train = GraphBatch.from_sequences(r_train, run_symbols)
rg_test = GraphBatch.from_sequences(r_test, run_symbols)
P = max(rg_train.shape[1], rg_test.shape[1])
E = int(max(rg_train.edges_per_graph.max(), rg_test.edges_per_graph.max()))
rg_train, rg_test = rg_train.pad(P, E), rg_test.pad(P, E)
print(rg_train)
symbols_backup, symbols = symbols, run_symbols # train_graph_tm reads `symbols`
rows, run_models = [], {}
for depth in (1, 2):
run_models[depth], hist, secs = train_graph_tm(
depth, rg_train, ry_train, rg_test, ry_test, n_classes=3, n_clauses=40, T=200, s=3.0,
epochs=10, max_included_literals=8)
rows.append({"depth": depth, "final test accuracy": round(hist["val_accuracy"][-1], 4),
"best": round(max(hist["val_accuracy"]), 4), "seconds": round(secs, 1)})
symbols = symbols_backup
pd.DataFrame(rows).set_index("depth")
Depth 1 stays near chance: a lone A does not know how long its run is. Depth 2 learns the task.
With one-hot messages a clause names the clauses that matched next door, so the strongest clauses
are read in two steps: the clause itself, then what the cited clauses test at layer 0.
import re
rm = run_models[2]
order = rm.weights.abs().max(dim=1).values.argsort(descending=True)[:6].tolist()
cited = set()
print("strongest clauses (votes for runs of 1, 2, 3):")
for j in order:
expr = rm.clause_expression(j, symbols=run_symbols)
cited |= {int(k) for k in re.findall(r"(?:left|right)1:(\d+)", expr)}
print(f" clause {j:2d} {str(rm.weights[j].tolist()):18s} {expr}")
print("\nwhat the cited messages vouch for (layer 0 of each cited clause):")
for k in sorted(cited):
print(f" clause {k:2d}: {rm.clause_expression(k, symbols=run_symbols, layers=[0])}")
With 40 clauses the rules are correct but no longer as tidy as in section 4: several clauses share
the work, and a message such as right1:26 stands for "clause 26 matched at the left neighbour",
where clause 26 may only say A or nothing at all (TRUE). The weighted vote combines them.
6. The GPU fast path¶
Graph updates join the CUDA fast paths of the convolutional models, with nodes for patches:
- Batched updates run in events mode: one Triton kernel computes, samples and applies the feedback.
- Padded batches also run sequential updates (one graph at a time, the classical algorithm) through a compiled per-example commit, and replay every update from a CUDA graph.
The cell below times one epoch of sequential updates on 2,000 graphs, flat versus padded. On a CPU nothing changes with padding, so it only runs on CUDA. These are rough timings, not benchmarks.
if device.type == "cuda":
flat_tr = GraphBatch.from_sequences(r_train[:2000], run_symbols)
rows = []
for name, data in (("flat", flat_tr), ("padded", rg_train[:2000])):
tt.seed_everything(SEED)
m = tt.GraphTsetlinMachine(run_symbols.dim, 3, 40, T=200, s=3.0, depth=2,
message_encoding="one_hot", n_edge_types=2,
max_included_literals=8, feedback_mode="sequential").to(device)
data = data.to(device)
ys = ry_train[:2000].to(device)
tr = tt.Trainer(m, batch_size=25, verbose=False)
tr.fit((data, ys), epochs=1) # warm-up: compile and capture
torch.cuda.synchronize()
t0 = time.time()
tr.fit((data, ys), epochs=1)
torch.cuda.synchronize()
rows.append({"layout": name, "ms per graph": round(1000 * (time.time() - t0) / len(data), 3)})
display(pd.DataFrame(rows).set_index("layout"))
else:
print("CPU runtime: padding changes nothing here, skipping the timing.")
7. Images as graphs: depth 1 is a convolutional machine¶
GraphBatch.from_images makes one node per image patch and no edges. Each node's features are the
patch pixels plus thermometer-coded coordinates, exactly the features of ConvTsetlinMachine. A
depth-1 clause is then an AND over one patch, and the class score ORs it over the nodes: that is
the convolutional coalesced Tsetlin machine.
In torchtsetlin the two are not just equivalent on paper. With the same random stream they learn the same automata, bit for bit (depth-1 graph batches with every node present run on the convolutional kernel). The cell trains both on 10,000 MNIST digits and compares their states.
x_tr, y_tr = tt.data.load_mnist_boolean(DATA_ROOT, train=True, threshold=0.3, device=device)
x_te, y_te = tt.data.load_mnist_boolean(DATA_ROOT, train=False, threshold=0.3, device=device)
x_tr, y_tr, x_te, y_te = x_tr[:10000], y_tr[:10000], x_te[:2000], y_te[:2000]
img_train = GraphBatch.from_images(x_tr, patch_size=10, stride=1)
img_test = GraphBatch.from_images(x_te, patch_size=10, stride=1)
print(img_train, "| full_nodes:", img_train.full_nodes)
settings = dict(n_clauses=1000, T=1000, s=10.0)
tt.seed_everything(SEED)
graph_tm = tt.GraphTsetlinMachine(img_train.shape[2], 10, depth=1, **settings)
tt.seed_everything(SEED)
conv_tm = tt.ConvCoalescedTsetlinMachine(10, patch_size=10, stride=1, input_shape=(1, 28, 28), **settings)
rows = []
for name, m, (a, b) in (("GraphTsetlinMachine, depth 1", graph_tm, (img_train, img_test)),
("ConvCoalescedTsetlinMachine", conv_tm, (x_tr, x_te))):
tt.seed_everything(SEED)
tr = tt.Trainer(m, device=device, batch_size=100, verbose=False)
t0 = time.time()
h = tr.fit((a, y_tr), epochs=3, val_data=(b, y_te))
rows.append({"model": name, **{f"epoch {i + 1}": round(v, 4) for i, v in enumerate(h["val_accuracy"])},
"seconds": round(time.time() - t0, 1)})
display(pd.DataFrame(rows).set_index("model"))
print("same automata:", torch.equal(graph_tm.ta_state, conv_tm.ta_state),
"| same clause weights:", torch.equal(graph_tm.weights, conv_tm.weights))
The paper's Table 3 uses this construction for MNIST and Fashion-MNIST. Our reproduction with 5
seeds (experiments/paper_experiments/results/graphtm/REPORT.md) lands within a point of every
published number. On MNIST, for example, the GraphTM reaches 98.07% against the paper's 98.42%.
Summary¶
- A graph Tsetlin machine clause is an AND over layers. Layer 0 reads the node, and each further layer reads which clauses matched at the in-neighbours one layer down.
- Depth is reach. A pattern spanning
hhops needs depthh + 1. At depth 1 the model is a convolutional machine over the nodes, bit for bit. - Clauses stay readable.
clause_expression,rulesandexplain_nodeprint them in the paper's notation and trace messages back to the nodes that sent them. - On a GPU, pad the dataset to one shape (
GraphBatch.pad) to get the fused kernels and CUDA graphs.
Further reading: the Graph input guide and Graph Tsetlin machines in the documentation.