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.40.0.tar.gz (317.8 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.40.0-py3-none-any.whl (286.9 kB view details)

Uploaded Python 3

File details

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

File metadata

  • Download URL: mechbench_compute-0.40.0.tar.gz
  • Upload date:
  • Size: 317.8 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.40.0.tar.gz
Algorithm Hash digest
SHA256 aea227d707e217340c9da822187a1aa9ae8147bcf8c5a0372206c7f3de88b278
MD5 f374cfb2402874d268959d0485d7d7bc
BLAKE2b-256 e7f2a28fcfc9044a97c06298955ef2530962e1618956c6270d6c3c44bf94d7fb

See more details on using hashes here.

File details

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

File metadata

File hashes

Hashes for mechbench_compute-0.40.0-py3-none-any.whl
Algorithm Hash digest
SHA256 2a3cc6d3ac624c9cb560adc9111ab8267b31df3319bac3c4a76052b9a88f7948
MD5 8ff1d006431f8fe13d5db3d4784ddc41
BLAKE2b-256 169f7d6a07ed86dbdc0bdf1290bd04606e4ecf1defcb46b9237239bd1ec52769

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

This release

0.40.0 This release

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

0.20.0

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