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.

3.8M parameters, fp32 safetensors (15 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 512 rows, `discover` averages 8 random 512-row subsets (raise `n_subsets` to use more of a big table).
variables	trained on 3-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"

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 512 rows are split into random 512-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).

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, 3-25 variables, 0.5-3.5 expected parents per node; 64-512 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.93	0.75	0.92	0.28
small (4-7 variables)	0.93	0.76	0.82	0.10
medium (10-16 variables)	0.91	0.76	0.92	0.30
large (19-25 variables)	0.93	0.75	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 protein each experiment 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 variable each experiment set (only for methods that can use it)

Edge detection (AUROC, 0.5 = random):

method	Sachs	Sachs+interventions	Chamber	Chamber+interventions
Axiom	0.73	0.74	0.78	0.77
AVICI	0.78	0.81	0.72	0.80
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	30%	37%	70%	97%
AVICI	47%	48%	41%	100%
PC	50%	—	51%	—
GES	52%	—	73%	—
DirectLiNGAM	64%	—	68%	—
partial corr.	50%	—	50%	—
correlation	50%	—	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	77%	88%	97%	85%
AVICI	81%	58%	80%	84%
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, which takes 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%	83%	86%	81%
AVICI	76%	56%	62%	64%	61%
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.
Directions on real data can come out reversed. The synthetic training data lets the model infer direction from which variable is more predictable from the others, and real systems often have the opposite pattern; this is why direction accuracy on Sachs is below chance. A fix to the training data is planned.
Causal discovery from observational data rests on assumptions; validate important edges with domain knowledge or experiments.

加载中…