The repo also trains three other architectures: CEM and ProbCBM from the official mateoespinosa/cem package, and ECBM, which is included and follows the authors' code. Install the CEM package with

./scripts/install_cem_repo.sh

or by hand:

git clone https://github.com/mateoespinosa/cem.git third_party/cem
python -m pip install "pytorch-lightning>=1.6,<2.0" "torchmetrics<1.0"
python -m pip install -r third_party/cem/requirements.txt
python -m pip install -e third_party/cem

Then pass --cbm-family cem, probcbm or ecbm to the robot or sudoku pipeline (the text pipeline trains the CBM only). Their settings (cem_emb_size, training_mode, probcbm_n_samples_inference, …) are fields of the config; see docs/installation.md.

Common parameters apply to both image and text modalities; scope-specific parameters are ignored when the other modality is selected.

from concept_benchmark.robots import DatasetGenerator, LabelFormula, F

dataset = DatasetGenerator(
    # ── Common (image + text) ──
    seed=1014,                       # random seed (default: 1014; the text pipeline uses 1337)
    data_type="image",               # "image" (default) or "text"
    concepts={                           # 9 features (default: ROBOT_CONCEPTS)
        "head_shape": ["square", "round"],
        "body_shape": ["square", "round"],
        "has_knees": ["false", "true"],
        "has_elbows": ["false", "true"],
        "has_antennae": ["false", "true"],
        "ears_shape": ["square", "triangle"],
        "mouth_type": ["closed", "open"],
        # the two shapes list their subtypes; each becomes one binary concept (round/edgy, flat/pointy)
        # unless it is named in expand_concepts
        "hand_shape": ["round_circle", "round_oval", "round_oval2",
                       "edgy_triangle", "edgy_square", "edgy_trapezoid"],
        "foot_shape": ["flat_trapezoid", "flat_rounded", "flat_square", "flat_5sided", "flat_lshaped",
                       "pointy_trapezoid", "pointy_rounded", "pointy_square", "pointy_3sided", "pointy_4sided"],
    },
    label_rule="balanced",           # "balanced" (default) or "sparse"; see Labeling rules below
    label_formula=None,              # or your own LabelFormula (see below), which replaces the rule's formula
    concept_preset="foot_subtypes",  # "ground_truth" or "foot_subtypes" (expands foot_shape into subtypes)
    renders_per_robot=4,             # samples per unique robot config (image: 4, text: 1)
    expand_concepts=["foot_shape"],                 # which features expand into subconcepts
    # ── Image-only (data_type="image") ──
    image_size="medium",             # "small" (8px), "medium" (32px), or "large" (600px)
    color_mode="color",              # "color" or "grayscale"
    render_images=True,              # set False to skip rendering PNGs (faster)
    # ── Text-only (data_type="text") ──
    template_complexity="high",      # template complexity level
).generate()

This guide shows how to evaluate your own concept bottleneck model on the benchmarks provided by this package. All examples below use the robot benchmark, but the same approach works for sudoku.

Prerequisite: You need the full repository (not just pip install concept-benchmark) to run the pipeline scripts and examples.

Getting data for your model

Generate a dataset as shown in the Quick Start, then access the splits:

train, val, test = dataset.train, dataset.val, dataset.test

Each split is a ConceptDatasetSample with these attributes:

Access formats:

# NumPy arrays; X (also .inputs) holds the image file names, the texts or the feature rows
X_train, C_train, y_train = train.X, train.C, train.y

# PyTorch DataLoader
loader = train.loader(batch_size=64, shuffle=True)
for x_batch, c_batch, y_batch in loader:
    ...

# Pandas DataFrame
df = train.to_dataframe()

Splitting

sample() splits the dataset into train/val/test. Sizes can be absolute counts or fractions:

# Absolute counts
dataset.sample(test_size=10000, val_size=1000, train_size=3800, seed=42)

# Fractions (of total dataset size)
dataset.sample(test_size=0.2, val_size=0.2, seed=42)

# Stratified — preserves class proportions in each split
dataset.sample(test_size=0.2, val_size=0.2, stratify=dataset.y, seed=42)

# Group-based — no group appears in multiple splits (group_ids: one id per row, e.g. the robot's identity)
dataset.sample(test_size=0.2, val_size=0.2, groups=group_ids, seed=42)

# Skewed training set — ensure min-fraction of specific concept patterns
dataset.sample(
    test_size=10000, val_size=0.2, train_size=3800,
    sampling_constraints=[{"concepts": {"my_concept": 1}, "min_fraction": 0.3}],
    seed=42,
)

The paper's robot split (generate_splits()) requires a minimum share of each training foot subtype this way; the constraints name subtypes that the preset drops, so split first:

from concept_benchmark.config import PRESET_EXCLUDED_CONCEPTS

dataset = DatasetGenerator(seed=1014, concept_preset="foot_subtypes").generate()
dataset.sample(
    test_size=10000, val_size=0.2, train_size=3800, seed=1014,
    sampling_constraints=[
        {"concepts": {"foot_shape_pointy_4sided": 1}, "min_fraction": 0.30},
    ],
)
# split first: the constraints name subtypes that the preset drops
dataset.drop_concepts(PRESET_EXCLUDED_CONCEPTS["foot_subtypes"])

You can re-split at any time by calling sample() again.

Concept missingness

Simulate missing concept annotations (e.g., incomplete labels from crowdsourcing). Use the generator transforms to produce a new dataset with missingness applied:

from concept_benchmark.transforms import ConceptMissingnessGenerator

dataset.sample(test_size=0.2, val_size=0.2, seed=42)

# MCAR: each concept label independently missing with probability p
dataset = ConceptMissingnessGenerator(dataset, p=0.2, mechanism="mcar", seed=99).generate()

# MNAR: missingness depends on concept value (present concepts more likely observed)
dataset = ConceptMissingnessGenerator(
    dataset, p=0.2, mechanism="mnar", seed=99,
    mnar_config={"present_prob": 0.8, "absent_prob": 0.1},
).generate()

Similarly, ConceptNoiseGenerator adds symmetric or asymmetric label flips, and LabelNoiseGenerator corrupts target labels:

from concept_benchmark.transforms import ConceptNoiseGenerator, LabelNoiseGenerator

dataset = ConceptNoiseGenerator(dataset, p=0.1, seed=99).generate()
dataset = LabelNoiseGenerator(dataset, p=0.05, seed=99).generate()

Wrapping your concept detector

The intervention runner requires a ConceptDetector subclass. Wrap your model:

from experiments.models import ConceptDetector

class MyConceptDetector(ConceptDetector):
    def __init__(self, my_model):
        super().__init__()
        self._my_model = my_model

    def predict_proba(self, dataset, **kwargs):
        """Must return (N, n_concepts) float array in [0, 1]."""
        return self._my_model.predict_concept_probs(dataset.X)

Key point: Override predict_proba() — it receives a ConceptDatasetSample, not raw arrays. Access inputs via dataset.X. The base class predict() calls predict_proba() and thresholds at 0.5 automatically.

Using a PyTorch module directly

If your model is already a PyTorch nn.Module, you can pass it directly instead of subclassing. Your module's forward() must:

• Accept a batched input tensor (e.g. (B, 3, 32, 32) for images)
• Return (B, n_concepts) raw logits (pre-sigmoid) — sigmoid is applied internally by predict()

The return value can also be a tuple (first element is used) or a dict with a "logits" key.

cd = ConceptDetector(model=my_pytorch_module)
cd.fit(train, val, fit_params={"epochs": 50, "lr": 1e-3})

You can also split your model into a backbone + concept head using the embedding_model parameter. The detector will probe the backbone's output shape and attach an MLP head automatically:

cd = ConceptDetector(embedding_model=my_backbone)
cd.fit(train, val, fit_params={"epochs": 50, "lr": 1e-3})

Wrapping your label predictor

Subclass FrontEndModel and override predict() and predict_proba():

from experiments.models import FrontEndModel

class MyFrontEnd(FrontEndModel):
    def __init__(self, my_classifier):
        super().__init__()
        self._clf = my_classifier

    def predict(self, C):
        """C is (N, n_concepts) binary 0/1. Returns (N,) int labels."""
        return self._clf.predict(C)

    def predict_proba(self, C):
        """C is (N, n_concepts) binary 0/1. Returns (N, n_classes) probs."""
        return self._clf.predict_proba(C)

Key point: C is already binarized at 0.5 — binary 0/1 values, not probabilities.

For a simple logistic regression baseline, the built-in FrontEndModel() works out of the box:

fe = FrontEndModel()
fe.fit(train.C, train.y)

Assembling and evaluating

Combine your concept detector and label predictor into a ConceptBasedModel:

import numpy as np
from experiments.models import ConceptBasedModel

cbm = ConceptBasedModel(
    concept_detector=MyConceptDetector(my_concept_model),
    label_predictor=MyFrontEnd(my_classifier),
)

predictions = cbm.predict(test)
accuracy = np.mean(predictions == test.y)
print(f"CBM accuracy: {accuracy:.4f}")

For running interventions and alignment on your model, see the Interventions and Alignment section and examples/robot_pipeline_example.py.

Metrics and Plots

concept_benchmark.evaluation has the metrics (accuracy, gain, selective_accuracy, coverage, abstention_threshold, net_work_automated, …) and the plots (plot_intervention_curve, plot_automation, plot_concept_report, …); experiments.evaluate has the blocks that need a model (train_cbm, intervention_table, automation_table, …). The plots take the results tables the pipelines write, average over seeds and draw a standard-error band:

from pathlib import Path
import pandas as pd
from concept_benchmark.evaluation import plot_intervention_curve

# the paper's results: paper-results.zip from the v0.4.0 release; the CBM corrected with the true values
files = Path("paper-results/robot/balanced_rule").glob("*__arch-cbm__isrc-perfect__*__results.csv")
results = pd.concat(pd.read_csv(f) for f in files)
results = results[results["concept_source"].isin(["ground_truth", "human_concepts"])]  # true and human concepts
fig, ax = plot_intervention_curve(results, group="concept_source", baseline_accuracy=0.88)

Interventions

Interventions are the core benefit of concept bottleneck models: at test time, a user (or automated system) can inspect and correct the model's concept predictions before the final label is determined.

Oracle interventions (manual approach)

The simplest way to run interventions is to directly manipulate concept predictions. After training a ConceptDetector and FrontEndModel, the workflow is:

1. Get concept probabilities from the detector
2. Identify the most uncertain concepts per sample
3. Replace them with ground-truth values
4. Re-predict with the label model

import numpy as np
from experiments.models import ConceptDetector, FrontEndModel

# cd and fe are a trained concept detector and label predictor, e.g. cbm.concept_detector and
# cbm.label_predictor of a CBM from experiments.evaluate.train_cbm

# Step 1: Get concept probabilities
concept_probs = cd.predict_proba(test)

# Step 2-3: For each sample, replace the k most uncertain concepts
# with ground-truth values
for k in [1, 3]:
    C_intervened = concept_probs.copy()
    uncertainty = np.abs(concept_probs - 0.5)  # distance from decision boundary
    for i in range(len(test)):
        most_uncertain = np.argsort(uncertainty[i])[:k]
        C_intervened[i, most_uncertain] = test.C[i, most_uncertain]

    # Step 4: Threshold to binary and predict
    preds = fe.predict((C_intervened > 0.5).astype(np.float32))
    acc = np.mean(preds == test.y)
    print(f"k={k}: accuracy={acc:.4f}")

Using the intervention API

For more complex intervention scenarios (budgets, strategies, batched evaluation), use the ConceptInterventionRunner:

import numpy as np
from experiments.models import ConceptBasedModel
from experiments.intervention import ConceptInterventionRunner, InterventionConfig
from experiments.kflip import KFlipInterventionStrategy

# Combine detector and label predictor into a CBM
cbm = ConceptBasedModel(concept_detector=cd, label_predictor=fe)

# Configure the intervention
config = InterventionConfig(
    per_instance_budget=3,  # correct up to 3 concepts per sample
    score_threshold=0.2,          # intervene when the label would change
                                  # with probability 0.2 or more
)

# Run interventions
runner = ConceptInterventionRunner(model=cbm)
strategy = KFlipInterventionStrategy()
result = runner.run(strategy, config, test)

# InterventionResult contains y_prob_before/after and y_pred_after
acc_before = np.mean(np.argmax(result.y_prob_before, axis=1) == test.y)
acc_after = np.mean(result.y_pred_after == test.y)
print(f"Accuracy before: {acc_before:.4f}")
print(f"Accuracy after:  {acc_after:.4f}")
print(f"Concepts corrected: {result.mask.sum()}")

Key classes

• InterventionConfig — controls intervention budgets, thresholds, and per-instance caps
• KFlipInterventionStrategy — the default strategy: evaluates all subsets of up to k concepts per sample and selects the intervention that maximizes predicted confidence
• ConceptInterventionRunner — coordinates intervention execution and before/after evaluation

Writing a custom strategy

Subclass InterventionStrategy and implement propose(model, batch, config), which returns a StrategyProposal with a boolean mask of the concepts to replace by their true values; prepare() is an optional step run once on the validation set.

import numpy as np
from experiments.intervention import InterventionStrategy, StrategyProposal

class UncertaintyStrategy(InterventionStrategy):
    def __init__(self):
        super().__init__(name="uncertainty")

    def propose(self, model, batch, config):
        k = int(min(config.per_instance_limit(batch.n_concepts), batch.n_concepts))
        mask = np.zeros((batch.n_samples, batch.n_concepts), dtype=bool)

        # Rank concepts by uncertainty (closeness to 0.5)
        uncertainty = 0.5 - np.abs(batch.C_pred - 0.5)
        for i in range(batch.n_samples):
            top_k = np.argsort(uncertainty[i])[-k:]
            mask[i, top_k] = True

        return StrategyProposal(mask=mask)

result = runner.run(
    strategy=UncertaintyStrategy(),
    config=InterventionConfig(per_instance_budget=3, score_threshold=0.2),
    dataset=test,
)

class MyStrategy(InterventionStrategy):
    def __init__(self):
        super().__init__(name="my_strategy")

    def prepare(self, model, batch, config):
        """Called once on the validation set before run()."""
        # Compute concept importance from validation data
        self._state["concept_order"] = compute_importance(model, batch)

    def propose(self, model, batch, config):
        order = self._state["concept_order"]
        # ... use precomputed order

runner.prepare(strategy, config, validation_dataset=val)
result = runner.run(strategy, config, dataset=test)

Concept sources and intervention sources

The robot pipeline varies two things independently: who annotates the concepts a model is trained on, and who answers when the model asks about a concept at test time.

python scripts/robot_pipeline.py --seed 1014 --concept-preset foot_subtypes \
    --concept-sources human_concepts machine_annotation \
    --intervention-sources perfect expert self --budgets 1 3 max

LLM interventions need an API key (--llm-api-key or GEMINI_API_KEY); --llm-cache-only reuses saved answers.

Regimes

--regimes names the pairings of the first version of the benchmark: baseline (human concepts, corrected with the true values), expert (corrected by an 80%-accurate expert), subjective (concepts with 20% label noise, expert), machine (CLIP-scored concepts, expert), llm and clip (LLM-written or CLIP-chosen concepts, corrected by an LLM).

python scripts/robot_pipeline.py --seed 1014 --concept-preset foot_subtypes --regimes baseline expert subjective
python scripts/robot_text_pipeline.py --seed 1337 --regimes baseline expert subjective

For a complete end-to-end example with training, interventions (through the intervention_table block, which runs ConceptInterventionRunner with the paper's policy) and alignment, see examples/robot_pipeline_example.py.

Alignment

Alignment constraints force the label predictor's concept weights to match a user's prior expectations about concept-label relationships. For example, if a domain expert knows that "has knees" should positively predict the Glorp class, alignment constrains that weight to be positive during retraining.

Why alignment matters

In standard CBM training, the label predictor (logistic regression on concept activations) learns weights freely from data. This can produce counterintuitive weights — e.g., has_knees getting a negative weight even when knees truly indicate Glorp — because the model exploits correlations among imperfect concept predictions.

The paper shows that an alignment constraint can:

• Raise accuracy before interventions — the constrained model no longer relies on a concept its detector reads at chance.
• Remove the benefit of interventions — after every concept is corrected, the unconstrained model is the more accurate one.

The constraint closes the pathway through which a corrected concept reaches the label.

Usage

After training a ConceptBasedModel, use run_alignment() to retrain the frontend with sign constraints and compare:

from experiments.utils import run_alignment

results = run_alignment(
    concept_based_model=cbm,
    train_dataset=train,
    test_dataset=test,
    monotonicity_constraints={"has_knees": 1},  # require a non-negative weight
)

print(f"Original accuracy: {results['original_accuracy']:.4f}")
print(f"Aligned accuracy:  {results['aligned_accuracy']:.4f}")
print(f"Accuracy change:   {results['accuracy_change']:+.4f}")

The monotonicity_constraints dict maps concept names to their required sign: +1 for positive weight, -1 for negative.

Expected results

Mean accuracy over the 10 seeds of the paper (balanced rule):

For a complete end-to-end example with training, interventions, and alignment, see examples/robot_pipeline_example.py.