Guaranteed Optimal Compositional Explanations for Neurons

This repository provides the official implementation of the algorithms introduced in the paper "Guaranteed Optimal Compositional Explanations for Neurons." Biagio La Rosa and Leilani H. Gilpin. ICML 2026 (Oral).

The repository includes:

• the algorithm for computing guaranteed-optimal compositional explanations
• the proposed decomposition of the IoU score (dIoU),
• heuristics for estimating explanation quality
• heuristic-guided beam-search variants
• utilities for comparing optimal and non-optimal explanations,
• scripts for reproducing the analyses presented in the paper.

We also release a Python package for easy integration into other projects including both the optimal compositional explanations and the beam-search explanations. The package can be installed via pip:

pip install compositional_explanations

Contributions and their Implementation

Decomposed IoU

Compositional explanations aim to identify the logical combination of concepts that maximizes the spatial alignment between neuron activations and semantic annotations in a probing dataset.

The alignment objective is defined through the Intersection over Union (IoU):

$$IoU(L, N, M) = \frac{ |\mathbf{1}_N \cap \mathbf{1}_{\theta(M,L)}| }{ |\mathbf{1}_N \cup \mathbf{1}_{\theta(M,L)}| }$$

where $N$ is the neuron activation matrix, $L$ is a logical explanation, and $\theta$ is the binary mask generated by applying the logical formula $L$ over the concept masks $M$.

To make optimal search feasible, the paper introduces a decomposition of the IoU score (dIoU) into fundamental quantities governing alignment quality. Specifically, the decomposition separates unique intersections ($I^U$), common intersections ($I^C$), unique extras ($E^U$), and common extras ($E^C$). The resulting formula is:

$$dIoU(N,L)= \frac{ \sum_{x \in D} |I^U(L)_x| + |I^C(L)_x| }{ |\mathbf{1}_N| + \sum_{x \in D} |E^U(L)_x| + |E^C(L)_x| }$$

The decomposition is mathematically equivalent to the original IoU score.

Implementation: the main utilities for computing the dIoU and for managing quantities are included in the optimal_utils.py file. For example the compute_quantities_vector function computes the concept quantities and the compute_max_iou_from_label_info function estimates the maximum achievable IoU using the previous formula.

---

Heuristics

The heuristics use the decomposed IoU quantities to estimate:

• the alignment of candidate expansion (i.e. IoU of $L$)
• the best possible alignment reachable from a node $L$ in the search tree by attaching additional concepts to the current formula ($L^*$).
• upper and lower bounds for unexplored paths ($IoU_{min}$ and $IoU_{max}$).

Two estimation strategies are possible:

• sample-based estimation, which computes quantities per sample and is more precise;
• aggregated estimation, which operates on global statistics and is significantly faster.

These heuristics are admissible and therefore preserve the optimality guarantees of the search algorithm.

Implementation: the sample-based estimation is included in the compositional/optimal_sample_heuristic.py file, while the aggregated estimation is included in the compositional/optimal_sum_heuristic.py file.

---

Optimal Algorithm

The algorithm performs a bounded informed search guided by the proposed heuristic.

At a high level, the algorithm:

1. Computes exact quantities for all concepts
2. Estimates upper and lower bounds for explanation paths
3. Expands only explanations whose maximum achievable alignment can improve over the current best solution
4. Refines heuristic estimates progressively
5. Prunes branches whose upper bound falls below the current optimum.

Unlike beam search methods, the proposed algorithm is guaranteed to recover the explanation with the highest possible IoU within the specified search space.

Implementation: the optimal algorithm is implemented in the compositional/optimal.py file, while the computation of path estimates and bounds is implemented in the compositional/path_heuristic.py file.

---

Beam Variant

This repository also includes the beam-search variant guided by our proposed heuristics. This variant is implemented in the compositional/beam_optimal.py file.

Other Beam Variants

The repository also includes the implementation of M-MESH in the compositional/mmesh.py file and a vanilla beam search without heuristic guidance in the compositional/vanilla_search.py file.

Quick Start

Option 1: Use the Python package

The easiest way to use the algorithms is to install the package via pip:

pip install compositional_explanations

In the home directory of the repository (https://pypi.org/project/compositional-explanations/), you can find the documentation for the package, including an usage example.

Option 2: Use this repository

Setup the environment

Follow step 1, 2, and 3 described in PAPER.md to set up the environment and download the required assets.

Main CLI Flags

All the scripts share a set of user arguments that define the configuration for which explanations are generated.

Common flags are defined in utils/common_flags.py. The main ones are:

• --dataset: default cityscapes_fine_sem_seg_val (also supports ade20k_full_sem_seg_freq_val_all and broden)
• --model: default resnet18 (also supports alexnet, densenet161)
• --layer: default layer4 (for alexnet and densenet161 use features)
• --heuristic: optimal, beam_optimal, mmesh, none
• --units : list of units to explain (default none, which means all the units)
• --random_units: number of random units to select (default: 0, which means all units)
• --device: compute device (default cuda)

In addition, the user can set the following flags

• --length: explanation length (default 3)
• --beam_limit: beam size (default 5)
• --num_clusters: number of clusters to consider (default and the settings used in the paper 1)
• --seed: random seed (default 0)
• --preload_masks: whether to preload masks in memory (default True, can speed up computation but can be memory intensive)
• --step_size: Number of concepts (and their masks across the whole dataset) to keep in memory during the generation of the masks from the annotation dataset (default 60). Larger values can speed up the computation but can be memory intensive.
• --fast_impl: whether to use a faster implementation for the generation of the masks from the annotation dataset (default False, can speed up computation but can be memory intensive)

Finally, the user can set the following flags to specify the output directories (if not using the default ones):

• --root_datasets: root directory where the datasets are stored (default data/dataset)
• --root_models: root directory where the models are stored (default data/model)
• --root_results: root directory where the results are stored (default data/results)
• --root_segmentations: root directory where the segmentation masks are stored (default data/cache/segmentations)
• --root_activations: root directory where the activations are stored (default data/cache/activations)
• --root_optimal_info: root directory where the optimal information is stored (default data/cache/optimal_info)

Compute Explanations

To compute explanations, use the run_clustering.py script with the common flags set appropriately.

python3 run_clustering.py \
	--dataset=<DATASET> \
	--model=<MODEL> \
	--layer=<LAYER> \
	--heuristic=<HEURISTIC> \
	--num_clusters=1

Explanation are saved in data/results/. Specifically, the explanations for different heuristics are saved in separate subdirectories: - data/results/optimal/ - data/results/beam_optimal/ - data/results/mmesh/ - data/results/no_heuristic/

Compare Methods

To compare behavior between different heuristics and optimal:

python3 compare_methods.py \
	--dataset=<DATASET> \
	--model=<MODEL> \
	--layer=<LAYER> \
	--num_clusters=1

This script uses all common flags except heuristic, since it searches for all available results across methods. Note that the script computes statistics only for configurations (units, model, layer, and so on) for which run_clustering.py has already been executed with matching flags. The output is displayed in the terminal.

Generate Qualitative Difference Visualizations

Generate qualitative difference figures between M-MESH and Optimal:

 python3 generate_images.py \
	--dataset=broden \
	--num_clusters=1

• Visualizations are saved under output/explanations/.

Reproducibility Notes

• Set --seed (default 0) for deterministic behavior where possible.
• Add CUBLAS_WORKSPACE_CONFIG=:4096:8 for CUDA reproducibility.
• Larger concept spaces can be memory-intensive. In these cases, consider tuning either --preload_masks, --step_size or --fast_impl.

Citation

If you use this repository, please cite:

@inproceedings{
LaRosa2026,
title={Guaranteed Optimal Compositional Explanations for Neurons},
author={Biagio {La Rosa} and Leilani H. Gilpin},
booktitle={Forty-third International Conference on Machine Learning},
year={2026},
url={https://openreview.net/forum?id=MHiiwC3oFR}
}