mechbench-compute

The compute engine for mechbench — composable mechanistic-interpretability primitives built on MLX.

This repo provides:

  • Hook-aware forward pass. One canonical path through the model; instrumentation via named hook points and TransformerLens-style callbacks.
  • Declarative interventions. Ablate, Capture, Patch primitives composable into a single model.run(..., interventions=[...]) call.
  • Activation cache. ActivationCache container for collected activations; bf16 throughout, float32 only at the analysis boundary.
  • Architecture adapter. Arch dataclass that handles per-variant differences (layer count, global-attention pattern, RoPE parameters, etc.). Currently supports Gemma 4 E4B and E2B; the adapter pattern follows TransformerLens 3's TransformerBridge.
  • Analysis helpers. Logit lens, direct logit attribution (accumulated_resid, decompose_resid, head_results, logit_attrs), fact vectors, centroid decoding, probe primitives, head-weight static analysis, geometry metrics.
  • Plot helpers. Matplotlib conventions baked in for quick diagnostic figures — not the full visualization surface (that lives in mechbench-ui).

See PACKAGE_README.md for the full API tour and worked examples.

Install

pip install mechbench-compute

From source:

git clone https://github.com/mechbench/mechbench-compute.git
cd mechbench-compute
pip install -e '.[dev]'

Apple Silicon required (MLX is the only supported backend today). A PyTorch backend would live as mechbench_compute.backends.torch alongside the MLX one if/when the need arises; splitting repos by backend is explicitly not planned.

Quick start

from mechbench_compute import Model, Ablate, Capture

model = Model.load()
ids = model.tokenize("Complete this sentence with one word: The Eiffel Tower is in")

result = model.run(ids)
for tok, p in result.top_k(model.tokenizer, k=5):
    print(f"{tok!r:20s} p={p:.4f}")

Distributional-target training (distill + lora)

Primitives for training a model toward a specified distribution over responses rather than toward example responses (task 000114): soft-target cross-entropy at decision tokens has gradient P − T, so the adapter learns to emit the distribution.

import mlx.nn as nn
import mlx.optimizers as optim
import numpy as np
from mechbench_compute import Model, distill, lora
from mechbench_compute.distill import TargetMap

model = Model.load()
tok = model.tokenizer

# A target is a Map<String, Double> — hardcoded, from JSON, or uniform —
# with whole-map transforms that each return a new map:
target = TargetMap.from_json("weights.json").sqrt().normalize()
target = TargetMap.uniform([str(i) for i in range(1, 7)])   # fair d6

# Compile it against the rendered prompt: items become token paths
# (multi-token items share trie nodes; a closer appends continuation
# anchors so the flattening can't leak past the envelope):
prompt = distill.render_chat(tok, system, "Please roll the die.",
                             prefill='{ "roll": ')
trie = target.tokenize(tok, prompt, closer=" }")

n = lora.apply_lora(model.lm)                 # freeze + wrap q/v projections
step = nn.value_and_grad(model.lm, distill.soft_ce)
opt = optim.Adam(learning_rate=1e-4)
rng = np.random.default_rng(7)
for _ in range(steps):
    batch = [trie.hard_example(trie.sample(rng)) for _ in range(3)]
    batch.append(trie.marginal_example())     # exact first-token marginal
    batch.append(sharp_anchor)                # keeps confident tasks sharp
    loss, grads = step(model.lm, batch)
    opt.update(model.lm, grads)

lora.save_adapter(model.lm, "adapter.safetensors")
# Later, on a fresh model: merge + exact undo
handle = lora.fuse(model.lm, lora.load_adapter("adapter.safetensors"),
                   scale=16 / 8)              # alpha / rank from training
lora.restore(model.lm, handle)

Calibration is measured at item level (trie.score, distill.item_metrics — captured mass, entropy, KL from target) and at the decision token (distill.first_token_metrics). python -m mechbench_compute._smoke_distill runs the full lifecycle on E2B.

Two forward paths. Training and scoring call Model.lm (the text decoder, uniform across families) directly — plain module calls, differentiable, no instrumentation. Model.run remains the hook-aware forward for capture/patch/lens work. Adapters bridge the two: fuse an adapter into the weights and every instrumented run sees the adapted model; restore flips it back, so base-vs-adapted comparisons run in one script.

Scoring tiers (task 000227): score_items is the sequential reference oracle; score_items_batched adds length-bucketed batching (~1.5×); score_items_fast additionally splits the forward via Model.trunk_hidden / Model.head_logits and unembeds only the supervised rows (~1.5–2.1× vs oracle, family-dependent — best when items share no prefix, e.g. cross-document scoring); score_items_cached encodes a shared prompt once into a KV cache (Model.prompt_cache) and scores each item's 1–4 suffix tokens against a per-item copy — the tier for shared-prompt batteries. Measured (flat name batteries, cached vs oracle): E2B 2.8×, Gemma-3-4B 3.4×, Qwen-3B 2.6×, Llama-8B 3.0×; positions/attention exact by construction, bf16 envelope from decomposed attention: mass-region |ΔlogP| ≤ 0.45, renormalized KL ≤ 2.4e-2 bits, idempotent (0.0 across repeat calls). Fast-tier envelope: max |ΔlogP| ≤ 0.99 (deep tail) / ≤ 0.24 (mass region), renorm KL ≤ 6.5e-3 bits. All of it is bf16 matmul tiling — same rows, same math, verified bit-exact where shapes match. Flat-target KL diagnostics weight the tail, so switch tiers only between comparisons, never mid-experiment. Known upstream limit (mlx 0.31.2 / mlx-lm 0.31.3 / mlx-vlm 0.6.1): batched cached decoding corrupts every batch row after the first on both stacks (reproduced with natively built B=4 caches and identical rows), which is why the cached tier is per-item; batched suffix scoring behind an upstream fix is the remaining ~5–10× path.

Status

The Arch adapter supports Gemma 4 E4B and E2B; generalization to other architecture families is ongoing.

The substrate epic that will define how intermediate results are cached and shared across experiments is 000162 (DAG solver + content-addressed memoization). It consumes the canonical-serialization guarantee from 000161 (binary formats) and the path grammar from 000163 (identity scheme).

Relationship to other mechbench repos

  • mechbench-schema — the typed emission contract. mechbench-compute emits records shaped by schema types.
  • mechbench-experiments — research scripts and findings that consume this package. Uses mechbench-compute as its primary dependency.
  • mechbench-runner — exposes these primitives as agent-callable tools. Imports mechbench-compute.
  • mechbench-ui — TypeScript frontend. Does not import mechbench-compute directly; reads bundles produced by it through the mechbench-schema contract.

See mechbench.ai for the family overview and the design principles.

License

MIT.

Project links

Key dates

PyPI data
Data sourced directly from PyPI's database.
  • Released:
Newer release available (0.79.0)

1 maintainer

PyPI data
Data sourced directly from PyPI's database.
Avatar for benjismith from gravatar.com benjismith

Credits

Author: Benji Smith

License

MIT License (MIT)

Requires

Python >=3.11

Provides Extra

eval dev

Tags

interpretability mechanistic-interpretability mlx transformers gemma

Classifiers

Report project as malware

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

mechbench_compute-0.20.0.tar.gz (197.4 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

mechbench_compute-0.20.0-py3-none-any.whl (189.9 kB view details)

Uploaded Python 3

File details

Details for the file mechbench_compute-0.20.0.tar.gz.

File metadata

  • Download URL: mechbench_compute-0.20.0.tar.gz
  • Upload date:
  • Size: 197.4 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/7.0.0 CPython/3.14.7

File hashes

Hashes for mechbench_compute-0.20.0.tar.gz
Algorithm Hash digest
SHA256 784252134d51ad437ac08a3fd8bebf2724d57e6b329a09f99b55f0f1321feb3a
MD5 979dbfae969bc275d3ca733bce44c6f0
BLAKE2b-256 1abebc8d841aab0ea42d36a30f3a43ea3c6e469006477316a18cbb246cc2600b

See more details on using hashes here.

File details

Details for the file mechbench_compute-0.20.0-py3-none-any.whl.

File metadata

File hashes

Hashes for mechbench_compute-0.20.0-py3-none-any.whl
Algorithm Hash digest
SHA256 8907ab278d5fb63193baa69d83b946e7d84e812ac866e8a4ec89ea4ca758bd5a
MD5 ab17f9f94db7c05f9507b12a199eb5c9
BLAKE2b-256 eb5bdd90121a12db9f0aca6b34a422c79db39a7689503c93771267de0c932bc4

See more details on using hashes here.

Release history Release notifications | RSS feed

0.79.0

2 files

0.78.1

2 files

0.78.0

2 files

0.77.1

2 files

0.77.0

2 files

0.76.2

2 files

0.76.1

2 files

0.76.0

2 files

0.75.1

2 files

0.75.0

2 files

0.74.0

2 files

0.73.0

2 files

0.72.0

2 files

0.71.0

2 files

0.70.0

2 files

0.69.0

2 files

0.68.0

2 files

0.67.0

2 files

0.66.0

2 files

0.65.0

2 files

0.64.0

2 files

0.63.0

2 files

0.62.0

2 files

0.61.0

2 files

0.60.0

2 files

0.59.0

2 files

0.58.0

2 files

0.56.0

2 files

0.55.0

2 files

0.54.0

2 files

0.53.0

2 files

0.52.1

2 files

0.52.0

2 files

0.51.0

2 files

0.50.0

2 files

0.45.0

2 files

0.44.0

2 files

0.43.0

2 files

0.41.0

2 files

0.40.0

2 files

0.36.0

2 files

0.29.0

2 files

0.28.0

2 files

0.27.0

2 files

0.26.0

2 files

This release

0.20.0 This release

2 files

0.18.2

2 files

0.18.1

2 files

0.18.0

2 files

0.17.1

2 files

0.17.0

2 files

0.16.4

2 files

0.16.3

2 files

0.16.2

2 files

0.16.1

2 files

0.16.0

2 files

0.15.9

2 files

0.15.8

2 files

0.15.7

2 files

0.15.6

2 files

0.15.5

2 files

0.15.4

2 files

0.15.3

2 files

0.15.2

2 files

0.15.1

2 files

0.15.0

2 files

0.14.4

2 files

0.14.3

2 files

0.14.2

2 files

0.14.1

2 files

0.14.0

2 files

0.13.0

2 files

0.12.0

2 files

0.11.1

2 files

0.11.0

2 files

0.10.0

2 files

Anthropic, PBC Visionary sponsor Bloomberg Visionary sponsor Hudson River Trading Visionary sponsor Meta Visionary sponsor NVIDIA Visionary sponsor Microsoft Sustainability sponsor Depot Continuous Integration AWS Cloud computing and Security Sponsor Datadog Monitoring Fastly CDN Google Download Analytics Sentry Error logging StatusPage Status page