ShortKit-ML — Detect and mitigate shortcuts and biases in machine learning embedding spaces. 20+ detection and mitigation methods with a unified API. Multi-attribute support tests multiple sensitive attributes simultaneously. Model Comparison mode for benchmarking multiple embedding models.

• Overview
• Installation
• Quick Start
• Detection Methods
• Overall Assessment Conditions
• MCP Server
• Paper Benchmarks
• Reproducing Paper Results
• GPU Support
• Interactive Dashboard
• Testing
• Contributing
• Citation

ShortKit-ML provides a comprehensive toolkit for detecting and mitigating shortcuts (unwanted biases) in embedding spaces:

• 20+ detection methods: HBAC, Probe, Statistical, Geometric, Bias Direction PCA, Equalized Odds, Demographic Parity, Intersectional, GroupDRO, GCE, Causal Effect, SSA, SIS, CAV, VAE, Early-Epoch Clustering, and more
• 6 mitigation methods: Shortcut Masking, Background Randomization, Adversarial Debiasing, Explanation Regularization, Last Layer Retraining, Contrastive Debiasing
• 5 pluggable risk conditions: indicator_count, majority_vote, weighted_risk, multi_attribute, meta_classifier

Key Features:

• Unified ShortcutDetector API for all methods
• Interactive Gradio dashboard with real-time analysis
• PDF/HTML/Markdown reports with visualizations
• Embedding-only mode (no model access needed)
• Multi-attribute support: test race, gender, age simultaneously
• Model Comparison mode: compare multiple embedding models side-by-side

Available on PyPI at pypi.org/project/shortkit-ml.

pip install shortkit-ml

For all optional extras (dashboard, reporting, VAE, HuggingFace, MCP, etc.):

pip install "shortkit-ml[all]"

git clone https://github.com/criticaldata/ShortKit-ML.git
cd ShortKit-ML
pip install -e ".[all]"

Or with uv:

uv venv --python 3.10
source .venv/bin/activate  # Windows: .venv\Scripts\activate
uv pip install -e ".[all]"

# macOS
brew install pango gdk-pixbuf libffi
# Ubuntu/Debian
sudo apt-get install libpango-1.0-0 libpangocairo-1.0-0 libgdk-pixbuf2.0-0

HTML and Markdown reports work without these. PDF export is optional.

from shortcut_detect import ShortcutDetector
import numpy as np

embeddings = np.load("embeddings.npy")  # (n_samples, embedding_dim)
labels = np.load("labels.npy")          # (n_samples,)

detector = ShortcutDetector(methods=['hbac', 'probe', 'statistical', 'geometric', 'equalized_odds'])
detector.fit(embeddings, labels)

detector.generate_report("report.html", format="html")
print(detector.summary())

For closed-source models or systems that only expose embeddings:

from shortcut_detect import ShortcutDetector, HuggingFaceEmbeddingSource

hf_source = HuggingFaceEmbeddingSource(model_name="sentence-transformers/all-MiniLM-L6-v2")
detector = ShortcutDetector(methods=["probe", "statistical"])
detector.fit(embeddings=None, labels=labels, group_labels=groups,
             raw_inputs=texts, embedding_source=hf_source)

See Embedding-Only Guide for CallableEmbeddingSource and caching options.

See Detection Methods Overview for per-method usage, interpretation guides, and code examples.

ShortcutDetector supports pluggable risk aggregation conditions that control how method-level results map to the final HIGH/MODERATE/LOW summary.

detector = ShortcutDetector(
    methods=["probe", "statistical"],
    condition_name="weighted_risk",
    condition_kwargs={"high_threshold": 0.6, "moderate_threshold": 0.3},
)

Custom conditions can be registered via @register_condition("name"). See Conditions API for details.

ShortKit-ML ships an MCP server so AI assistants (Claude, Cursor, etc.) can call detection tools directly from chat — no Python script required.

pip install -e ".[mcp]"

# via entry point (after install)
shortkit-ml-mcp

# or directly
python -m shortcut_detect.mcp_server

Add the following to ~/Library/Application Support/Claude/claude_desktop_config.json (macOS):

{
  "mcpServers": {
    "shortkit-ml": {
      "command": "python",
      "args": ["-m", "shortcut_detect.mcp_server"],
      "cwd": "/path/to/ShortKit-ML"
    }
  }
}

A ready-to-edit template is included at claude_desktop_config.json.

Configure examples/paper_benchmark_config.json to control effect sizes, sample sizes, imbalance ratios, and embedding dimensionalities. A smoke profile (examples/paper_benchmark_config_smoke.json) is provided for quick sanity checks.

python -m shortcut_detect.benchmark.paper_run --config examples/paper_benchmark_config.json

Outputs CSVs, figures, and summary markdown into output/paper_benchmark/.

Requires a CheXpert manifest (data/chexpert_manifest.csv) plus model-specific embedding pickles. Supported models: medclip, biomedclip, cxr-foundation.

python3 scripts/run_dataset2_benchmark.py \
  --manifest data/chexpert_manifest.csv \
  --model medclip \
  --root . \
  --artifacts-dir output/paper_benchmark/chexpert_embeddings \
  --config examples/paper_benchmark_config.json

See scripts/reproduce_paper.sh and the Dockerfile for full reproducibility.

All paper results are fully reproducible with fixed seeds (seed=42). Every table and figure in the paper can be regenerated from the scripts and data in this repository.

13 benchmark methods are evaluated across all datasets: hbac, probe, statistical, geometric, frequency, bias_direction_pca, sis, demographic_parity, equalized_odds, intersectional, groupdro, gce, ssa. These span 5 paradigms: embedding-level analysis, representation geometry, fairness evaluation, explainability, and training dynamics.

docker build -t shortcut-detect .
docker run --rm -v $(pwd)/output:/app/output shortcut-detect full

Important: The data/ folder in this repository is empty. All embeddings and metadata are hosted on HuggingFace and must be downloaded separately. Raw CheXpert and MIMIC-CXR images and labels are not redistributed — access requires accepting the respective dataset licenses (PhysioNet for MIMIC-CXR, Stanford for CheXpert).

All embeddings are hosted on HuggingFace (gated, PhysioNet-restricted access):

MITCriticalData/ShortKit-ML-data

# Download all embeddings into data/
huggingface-cli download MITCriticalData/ShortKit-ML-data --repo-type dataset --local-dir data/

See docs/reproducibility.md for full details.

The library auto-selects the best available device. PyTorch components (probes, VAE, GroupDRO, adversarial debiasing, etc.) use the standard torch.device fallback:

Note: Most detection methods (HBAC, statistical, geometric, etc.) run on CPU via NumPy/scikit-learn and do not require GPU. GPU acceleration benefits the torch-based probe, VAE, GroupDRO, and mitigation methods. MPS support depends on PyTorch operator coverage; if you encounter errors on Apple Silicon, fall back to device="cpu".

python app.py
# Opens at http://127.0.0.1:7860

Features: sample CheXpert data, custom CSV upload, PDF/HTML reports, model comparison tab, multi-attribute analysis.

CSV Format:

embedding_0,embedding_1,...,task_label,group_label,attr_race,attr_gender
0.123,0.456,...,1,group_a,Black,Male

See Dashboard Guide for detailed usage.

pytest tests/ -v
pytest --cov=shortcut_detect --cov-report=html

638 tests passing across all detection and mitigation methods.

pip install -e ".[dev]"
pre-commit install

• Black for formatting (line length: 100), Ruff for linting, MyPy for types
• Pre-commit hooks run automatically; CI tests on Python 3.10, 3.11, 3.12
• New detectors must implement DetectorBase. See docs/contributing.md and shortcut_detect/detector_template.py

shortcut_detect/
├── probes/           # Probe-based detection (sklearn + torch)
├── clustering/       # HBAC detector
├── statistical/      # Statistical testing
├── geometric/        # Geometric & bias direction analysis
├── fairness/         # Equalized Odds, Demographic Parity, Intersectional
├── frequency/        # Frequency shortcut detector
├── causal/           # Causal effect detector
├── gce/              # Generalized cross-entropy detector
├── training/         # Early epoch clustering (SPARE)
├── vae/              # VAE latent disentanglement
├── xai/              # CAV, SpRAy, GradCAM mask overlap, SIS
├── ssa/              # Semi-supervised spectral analysis
├── groupdro/         # GroupDRO worst-group robustness
├── conditions/       # Pluggable risk aggregation conditions
│   ├── base.py, registry.py, indicator_count.py, majority_vote.py
│   ├── weighted_risk.py, multi_attribute.py, meta_classifier.py
│   └── meta_model.joblib  # Trained meta-classifier (bundled)
├── benchmark/        # Paper benchmark infrastructure
│   ├── runner.py, paper_runner.py, synthetic_generator.py
│   ├── measurement.py, fp_analysis.py, sensitivity.py
│   ├── convergence_viz.py, baseline_comparison.py, figures.py
├── comparison/       # Model comparison runner
├── mitigation/       # Debiasing & masking methods (M01-M07)
├── reporting/        # HTML/PDF/CSV reports & visualizations
├── unified.py        # ShortcutDetector unified API
└── detector_base.py  # DetectorBase ABC with results_ schema

docs/                 # MkDocs documentation site
examples/             # Notebooks and benchmark configs
app.py                # Gradio dashboard
Dockerfile            # Reproducible environment
scripts/              # Paper reproduction scripts
tests/                # Test suite (475+ tests)

pip install mkdocs mkdocs-material "mkdocstrings[python]" pymdown-extensions
mkdocs serve  # http://127.0.0.1:8000

• Getting Started
• Detection Methods -- all 20+ methods with guides
• API Reference
• Contributing

@software{shortkit_ml2025,
  title={ShortKit-ML: Tools for Identifying Biases in Embedding Spaces},
  author={Sebastian Cajas, Aldo Marzullo, Sahil Kapadia, Qingpeng Kong, Filipe Santos, Alessandro Quarta, Leo Celi},
  year={2025},
  url={https://github.com/criticaldata/ShortKit-ML}
}

MIT License - see LICENSE file

• GitHub: criticaldata/ShortKit-ML
• Issues: GitHub Issues