InterMol is an open-source mechanistic interpretability toolbox that leverages sparse autoencoders (SAEs) to interpret chemical language models (cLMs). By default, InterMol uses MoLFormer-XL as its underlying cLM. Our interactive visualizer of the discovered features is available at intermol.co.
Learn more about the implementation details and findings by reading our preprint.
.
├── intermol/
│ ├── main/ # Core SAE pipeline: training, normalization, and inference
│ └── interp/ # Interpretability tools: latent profiling, concept generation, labelling, and SMILES variant generation
├── notebooks/ # Jupyter notebooks for visualization and analysis
├── scripts/ # CLI interpretability utils
├── LICENSE
├── README.md
├── pyproject.toml
└── requirements.txt
Prerequisites: Python >= 3.9
# Package only
pip install git+https://github.com/ckennetha/intermol.git
# Clone the repository (recommended for local development)
git clone https://github.com/ckennetha/intermol.git
cd intermol
pip install -e .We provide SAE weights trained on top of the open-source version of MoLFormer-XL at layers 1, 3, 6, 9, and 12. MoLFormer-XL weights are fetched on-the-fly from HuggingFace, while SAE weights must be downloaded separately. To extract SAE activations using a pretrained model:
from intermol.main.inference import SAEInferenceConfig, SAEWithBaseModel
MODEL_NAME = "ibm/MoLFormer-XL-both-10pct"
USE_MOLFORMER = True # MoLFormer requires different HF setting
SMILES = "c1ccccc1"
config = SAEInferenceConfig(
layer=1 # MoLFormer-XL layer
hidden_dim=3072,
k=128,
weights_path="norm-MOL-1-3072-128.pt", # normalized SAE weights
)
sae = SAEWithBaseModel(config, MODEL_NAME, USE_MOLFORMER)
mf_acts, sae_acts = sae.encode(SMILES)
# mf_acts: MoLFormer-XL hidden states, sae_acts: SAE activationsTo extract SAE activations in bulk, we provide a command line interface (run-precomp-acts) that efficiently stores the results as decomposed CSC sparse matrix format using h5py. Multiple SAE configs can be processed in a single run by repeating per-SAE arguments:
usage: run-precomp-acts [-h] --data-path DATA_PATH
--layer LAYER --hidden-dim HIDDEN_DIM --k K
--sae-ckpt-path SAE_CKPT_PATH
[--outdir-path OUTDIR_PATH]
[--chunk-size CHUNK_SIZE] [--out-prefix OUT_PREFIX]
[--model-name MODEL_NAME] [--use-molformer]
[--device {auto,cpu,cuda}]
options:
-h, --help show this help message and exit.
--data-path DATA_PATH Path to .txt or one-column .smi file.
--layer LAYER Base model layer (repeat per SAE).
--hidden-dim HIDDEN_DIM SAE latent dimension (repeat per SAE).
--k K Number of top-k SAE latents (repeat per SAE).
--sae-ckpt-path SAE_CKPT_PATH Path to trained SAE checkpoint (repeat per SAE).
--outdir-path OUTDIR_PATH Output directory (repeat per SAE). Default: current directory.
--chunk-size CHUNK_SIZE Number of samples per chunk. Default: 8192.
--out-prefix OUT_PREFIX Output filename prefix. Default: current timestamp.
--model-name MODEL_NAME Hugging Face model name.
--use-molformer Enable MoLFormer-specific setting.
--device {auto,cpu,cuda} Inference device. Default: auto.
For evaluating association between specific SAE latents and atom-level molecular concepts, we use a two-step approach: first filtering latents by standardized mean difference (SMD) using the --is-prefilter flag, then concept presence classification with binarized activations, evaluated by F1 score. Currently, this method only supports SAEs trained on MoLFormer, as MoLFormer uses atom-wise tokenization. We provide run-eval-concepts with Numba-accelerated computation:
usage: run-eval-concepts [-h] --data-path DATA_PATH --acts-h5-path ACTS_H5_PATH
--label-path LABEL_PATH --outdir-path OUTDIR_PATH
--outfn OUTFN --sample-colname SAMPLE_COLNAME
--concept-colname CONCEPT_COLNAME
--label-colname LABEL_COLNAME
--index-colname INDEX_COLNAME
[--fpc-path FPC_PATH] [--desc-colname DESC_COLNAME]
[--thresholds THRESHOLDS] [--use-pooling]
[--is-prefilter] [--batch-size BATCH_SIZE]
[--score-colname SCORE_COLNAME]
[--score-threshold SCORE_THRESHOLD] [--k K]
[--is-sampling] [--fraction-sampling FRACTION_SAMPLING]
[--seed-sampling SEED_SAMPLING]
options:
-h, --help show this help message and exit
path options:
--data-path DATA_PATH Path to input .parquet file
--acts-h5-path ACTS_H5_PATH Path to precomputed activations .h5 file
--label-path LABEL_PATH Path to concept label .tsv file
--outdir-path OUTDIR_PATH Output directory
--outfn OUTFN Output filename (without extension)
--fpc-path FPC_PATH Path to prefiltering output. If not provided, concepts are evaluated across all SAE latents
column name options:
--sample-colname SAMPLE_COLNAME Column name for samples
--concept-colname CONCEPT_COLNAME Column name for concepts (must match in
both data and label files)
--label-colname LABEL_COLNAME Column name for labels
--index-colname INDEX_COLNAME Column name for concept indices in label file
--desc-colname DESC_COLNAME Column name for concept descriptions (optional)
evaluation options:
--thresholds THRESHOLDS Thresholds for evaluation (pass multiple). Default: 0
--use-pooling Use pooling-based evaluation for concepts
spanning multiple tokens
--is-prefilter Run SAE latent prefiltering with SMD
instead of full evaluation
--batch-size BATCH_SIZE Batch size for evaluation. Default: 8192
post-prefiltering options:
--score-colname SCORE_COLNAME Score column name for filtering. Set to 'smd' when using --is-prefilter output.
--score-threshold SCORE_THRESHOLD Minimum SMD score threshold. Default: 0
--k K Top-k features per concept. Default: 64
sampling options:
--is-sampling Enable molecule sampling
--fraction-sampling FRACTION_SAMPLING
Fraction of data to sample. Default: 0.20
--seed-sampling SEED_SAMPLING Random seed for sampling. Default: 42
Below is an example of prefiltering single-token concepts with a sampled dataset. For multi-token concepts, simply set the --use-pooling flag.
run-eval-concepts \
--data-path valid_dataset.parquet \
--acts-h5-path valid_acts.h5 \
--label-path SMARTS_ignChi_f10.tsv \
--outdir-path ../results_layer1/ \
--outfn valid_smd \
--sample-colname smiles \
--concept-colname concept \
--index-colname id \
--label-colname token_idxs \
--is-prefilter \
--is-sampling \
--fraction-sampling 0.2 \
--seed-sampling 42To run association analysis on single-token concepts, remove the --is-prefilter flag. If --fpc-path is supplied with prefiltering output, set --score-colname to specify the score column (e.g., 'smd'), --score-threshold to filter by SMD score, and --k for the number of top-k latents. If --fpc-path is supplied with the ConceptEvaluator output, omit these args. As before, set --use-pooling for multi-token concepts.
run-eval-concepts \
--data-path valid_dataset.parquet \
--acts-h5-path valid_acts.h5 \
--label-path SMARTS_ignChi_f10.tsv \
--fpc-path valid_smd.tsv \
--outdir-path ../results_layer1/ \
--outfn valid_eval \
--sample-colname smiles \
--concept-colname concept \
--index-colname id \
--label-colname token_idxs \
--thresholds 0 0.1 0.2 0.35 0.5 0.6 0.8 \
--score-colname smd \
--score-threshold 0 \
--k 64-
activation_dist.ipynb— Visualizing activation distribution of a specific SAE latent. Supports coloring by specific tokens and SMARTS patterns. -
probe_latents.ipynb— Extracting SAE activations and other chemical features into a.h5file for interpretable linear probing experiments.
This project is licensed under the MIT License.