maxpagels/axiom-base

Axiom (base)

Causal discovery for tabular data. In one forward pass over a whole table (rows x variables), Axiom labels every pair of variables causes (a -> b), caused_by (b -> a), undirected (linked, direction not identifiable), confounded (no edge, hidden common cause) or none, with probabilities, and assembles the result into an acyclic graph. Missing values and rows from known experiments are supported.

Treat the graph as hypotheses to check against domain knowledge: uncertain pairs are candidates for experiments; confident ones can feed causal effect identification and estimation.

11.9M parameters, fp32 safetensors (48 MB). Project page: maxpagels.com/axiom.

Recommended use (ballpark)

	recommendation
rows	256 or more works best; usable from ~100; below that, results get unreliable. Over 2048 rows, discover averages 8 random 2048-row subsets (raise n_subsets to use more of a big table).
variables	trained on 2-25. On simulated data it held up to 100 with no loss in accuracy; beyond 25 on real data is untested. Memory grows quickly with the variable count (~1 GB on CPU at 100, ~3 GB at 200).
interventions	if any rows come from experiments, label all of them: unlabeled experimental rows act like hidden causes and can produce wrong or reversed edges.
output	treat edges and especially directions as ranked hypotheses to check, not a verdict (see Limitations).

Usage

pip install "git+https://huggingface.co/maxpagels/axiom-base"

Upgrading from an earlier version: add -U (pip install -U "git+https://huggingface.co/maxpagels/axiom-base"). From 0.3.0 on, the package loads the weights released with it, so code and weights always match. Before 0.3.0 the weights always came from main: a 0.2.0 install loads the current weights without an error but gives wrong results. Upgrade, or pin it to its own release (see Versions).

from axiom import Axiom

model = Axiom.from_pretrained("maxpagels/axiom-base")
g = model.discover(df)           # pandas DataFrame or (n, d) array; NaN = missing
print(g)                         # the relations found, strongest first

# rows from known experiments: one entry per row (None = observational),
# or an (n, d) bool mask
targets = [None, "PKC", ["PKA", "PIP3"], None, ...]   # column(s) each row's experiment targeted
g = model.discover(df, interventions=targets)

discover standardizes each column as in training (do not pass raw model inputs yourself unless you use axiom.preprocess). Tables with more than 2048 rows are split into random 2048-row subsets whose predictions are averaged.

Reading the result

For every pair of variables a and b, Axiom gives a Relation(a, b, kind, p, probs):

• kind: Axiom's conclusion about how a and b are causally linked, one of the five below.
• p: the model's probability for that kind.
• probs: the probabilities of all five kinds (they sum to 1); kind is the decision, probs the evidence behind it.

kind	meaning	example
causes	a directly causes b (a -> b)	smoking -> tar in lungs
caused_by	b directly causes a (b -> a)	the same pair, read from the other side
undirected	a and b are directly linked, but the data cannot tell which way	linear-Gaussian data, where both directions fit equally well
confounded	no direct link, but an unmeasured variable causes both	ice cream sales and drownings (hidden cause: hot weather)
none	no direct link

A kind is always read from the first variable: g.relation("PKA", "PKC").kind == "causes" and g.relation("PKC", "PKA").kind == "caused_by" describe the same edge.

Only direct links count. In a -> b -> c, the pair a, c is none, even though a affects c, because that effect passes through b. Drop b from the table and a -> c becomes a direct link: what counts as direct depends on which variables you include.

g.relations                      # [Relation(a, b, kind, p, probs), ...]: the relations found,
                                 #   strongest first; directed ones read cause first
g.edges                          # [(cause, effect, p), ...]: directed edges only
[r for r in g.relations if r.kind == "confounded"]   # one kind only: here, the pairs
                                                     #   likely to share a hidden cause

g.relation("PKA", "PKC")         # any one pair, also "none" ones: Relation(a='PKA', b='PKC',
                                 #   kind='causes', p=0.85, probs={'none': 0.0, 'causes': 0.85, ...})

g.to_frame()                     # the relations found as a pandas DataFrame: a, b, kind, p, p_<kind>...
g.to_frame(all_pairs=True)       # every pair, plus `top`, the most probable kind

A pair only gets a kind other than none if its probability clears thresh (default 0.5; lower it with e.g. model.discover(df, thresh=0.3)). The raw probabilities and 0/1 matrices behind all this are on g too; see help(g).

Versions

Each release is a tagged commit in this repo; main is the latest. To use a specific one, install its code and load its weights by tag:

pip install "git+https://huggingface.co/maxpagels/[email protected]"

model = Axiom.from_pretrained("maxpagels/axiom-base")                     # this package's release: v0.3.1
model = Axiom.from_pretrained("maxpagels/axiom-base", revision="v0.2.0")  # weights of another release

• v0.3.1: weights load strictly: code and weights that do not match raise an error instead of silently dropping weights; same weights as 0.3.0.
• v0.3.0: classical pair statistics (correlations, partial correlations, pairwise direction cues) as model input; per-row condition ids (discover(conditions=...)); trained on 2-25 variables with a more realistic prior (experiment mixtures, skewed noise, feedback loops, measurement artefacts); weights pinned to the package version.
• v0.2.0: the release of October 2026.

Training

Trained only on synthetic data: every step simulates fresh random structural causal models (no dataset is seen twice).

• Graphs: Erdos-Renyi and scale-free DAGs, 2-25 variables, 0.5-3.5 expected parents per node; 64-2048 rows per dataset.
• SCM families: linear-Gaussian (target: the CPDAG), linear non-Gaussian, nonlinear (MLP, additive, interaction, random-Fourier mechanisms) and mixed; several noise families, heteroscedastic noise, post-nonlinear transforms.
• Hidden confounders, perfect and soft interventions (reported or not), measurement error, discretization, outliers, and MCAR/MAR/MNAR missingness.

Evaluation

Held-out synthetic data: fixed datasets simulated from the training distribution but never trained on, 256 rows each.

graphs	edge AUROC	skeleton F1	orientation accuracy	confounded-pair F1
all (4-25 variables)	0.92	0.70	0.91	0.29
small (4-7 variables)	0.85	0.82	0.77	0.00
medium (10-16 variables)	0.90	0.70	0.92	0.29
large (19-25 variables)	0.92	0.68	0.93	0.31

Real data

Datasets with a known causal graph; every method sees the same 500 standardized rows per draw, default settings, mean over 5 draws.

• Sachs: real protein-signaling data, 11 proteins, 18-edge consensus graph; 500 random cells per draw out of 7466 (pooled over 9 experiments)
• Sachs+interventions: the same cells, plus which experiment each came from and which protein it targeted (only for methods that can use it)
• Chamber: real physical system (light tunnel: light source, polarizers, sensors), 20 variables, 39-edge graph known from physics; 500 random rows per draw out of 26000 (10000 observational, 16000 from experiments that set one input)
• Chamber+interventions: the same rows, plus which experiment each came from and which variable it set (only for methods that can use it)

Edge detection (AUROC, 0.5 = random):

method	Sachs	Sachs+interventions	Chamber	Chamber+interventions
Axiom	0.75	0.79	0.80	0.79
AVICI	0.78	0.81	0.72	0.80
TabCausal	0.72	0.77	0.78	0.83
CDFM	0.77	—	0.76	—
PC	0.57	—	0.55	—
GES	0.61	—	0.69	—
DirectLiNGAM	0.67	—	0.77	—
partial corr.	0.64	—	0.73	—
correlation	0.60	—	0.72	—

Direction (% of identifiable true edges pointed the right way, 50% = coin flip):

method	Sachs	Sachs+interventions	Chamber	Chamber+interventions
Axiom	33%	37%	94%	98%
AVICI	47%	48%	41%	100%
TabCausal	56%	37%	60%	98%
CDFM	40%	—	97%	—
PC	50%	—	51%	—
GES	52%	—	73%	—
DirectLiNGAM	64%	—	68%	—
partial corr.	50%	—	50%	—
correlation	50%	—	50%	—

Sachs direction by evidence. Of the consensus edges, 8 link proteins that are visibly correlated within experiments; the other 10 show no such link, so the only evidence about them is how whole experiments shift the proteins. Those shifts do not fit the consensus graph at the level of the measured proteins, because the reagents act on enzyme activity rather than on measured levels. The consensus graph also contains a cycle (PIP3 -> plcg -> PIP2 -> PIP3), so no acyclic graph gets every direction. DirectLiNGAM's Sachs score rests on one experiment: without the b2cAMP cells it is 37% instead of 64%.

method	Sachs, 8 linked edges	Sachs, 10 unlinked edges	Sachs+interventions, 8 linked edges	Sachs+interventions, 10 unlinked edges
Axiom	45%	24%	42%	32%
AVICI	48%	46%	62%	36%
TabCausal	60%	52%	57%	20%
CDFM	52%	30%	—	—
PC	60%	42%	—	—
GES	46%	56%	—	—
DirectLiNGAM	64%	64%	—	—
partial corr.	50%	50%	—	—
correlation	50%	50%	—	—

Tübingen cause-effect pairs

102 real two-variable datasets (altitude and temperature, age and height, ...) with a known direction; up to 500 random rows each. Only direction is scored, as the weighted accuracy of the dataset's own weights (Mooij et al., 2016). Pairs taken from the same source share one unit of weight; ± is a bootstrap standard error over pairs. Axiom is trained on 2-25 variables.

Direction accuracy (50% = coin flip):

method	weighted	unweighted
Axiom	74% ± 6	70%
AVICI	58% ± 6	59%
TabCausal	61% ± 7	65%
CDFM	60% ± 6	67%
PC	50% ± 0	50%
GES	50% ± 0	50%
DirectLiNGAM	54% ± 6	50%
partial corr.	50% ± 0	50%
correlation	50% ± 0	50%

Synthetic benchmarks

Standard simulated benchmarks from the literature, 5 random graphs each.

• LinGauss: linear-Gaussian SCM on ER-2 graphs (NOTEARS, Zheng et al. 2018); direction only identifiable for some edges; 10 variables, 500 rows
• LiNGAM: linear SCM with non-Gaussian (exponential) noise on ER-2 graphs (Shimizu et al. 2006); 10 variables, 500 rows
• GP-add: nonlinear additive SCM, functions drawn from a Gaussian process, ER-2 graphs (CAM, Buhlmann et al. 2014); 10 variables, 500 rows
• MLP-SF: nonlinear SCM, random 1-hidden-layer MLP per variable, scale-free SF-2 graphs (NOTEARS-MLP, Zheng et al. 2020); 20 variables, 500 rows

Direction (% of identifiable true edges pointed the right way, 50% = coin flip):

method	LinGauss	LiNGAM	GP-add	MLP-SF
Axiom	88%	82%	96%	83%
AVICI	81%	58%	80%	84%
TabCausal	73%	100%	91%	89%
CDFM	69%	90%	90%	77%
PC	54%	50%	61%	48%
GES	49%	46%	56%	69%
DirectLiNGAM	43%	97%	42%	37%
partial corr.	50%	50%	50%	50%
correlation	50%	50%	50%	50%

Robustness

The synthetic benchmarks pooled, under common data problems. Missing values are mean-imputed for every method except Axiom and CDFM, which take them as given.

• clean: no changes
• missing: 20% of cells missing at random
• outliers: 2% of cells replaced by gross outliers (5-10 sd)
• meas.noise: measurement noise added (20% of each variable's variance)
• n=100: only 100 rows

Direction (% of identifiable true edges pointed the right way, 50% = coin flip):

method	clean	missing	outliers	meas.noise	n=100
Axiom	87%	84%	84%	81%	83%
AVICI	76%	56%	62%	64%	61%
TabCausal	88%	56%	65%	79%	83%
CDFM	82%	75%	66%	67%	76%
PC	54%	54%	52%	56%	53%
GES	55%	62%	53%	62%	54%
DirectLiNGAM	55%	49%	49%	53%	61%
partial corr.	50%	50%	50%	50%	50%
correlation	50%	50%	50%	50%	50%

Limitations

• Synthetic training only; real data that differs from the prior (cycles, time series, selection bias, truly discrete mechanisms) is out of distribution, and more than 25 variables is only tested on simulated data.
• Hidden-confounder detection is weak; treat the confounded output as a hint.
• Direction scores on pairs Axiom gives little link probability carry little information: on Sachs they follow shifts between experiments and mostly come out reversed, which is why its Sachs direction is below chance overall (on edges the data visibly links it is near chance, like the other methods). Read directions together with p, not on their own.
• The decoded graph can have too many edges: on Sachs its SHD is worse than that of an empty graph even when the edge ranking is good; raise thresh for a sparser graph.
• Causal discovery from observational data rests on assumptions; validate important edges with domain knowledge or experiments.