Skip to content

spec.md

Latest commit

 

History

History
880 lines (681 loc) · 43.6 KB

spec.md

File metadata and controls

880 lines (681 loc) · 43.6 KB
Raw

Neuroscope: SAE-Instrumented LLM Inference Server

Overview

Neuroscope is a CLI tool that runs local LLM inference with real-time Sparse Autoencoder (SAE) feature extraction. It exposes an OpenAI-compatible chat API alongside a Server-Sent Events (SSE) stream of SAE feature activations, enabling external UIs to visualize what concepts the model is "thinking about" as it generates each token.

Goals

  1. Run a local LLM and serve it via an OpenAI-compatible /v1/chat/completions endpoint.
  2. Instrument the model's forward pass to extract residual stream activations at a configurable layer.
  3. Run a pre-trained SAE encoder on those activations per token.
  4. Stream the top-K activated features (with human-readable labels) via a separate SSE endpoint in real-time, synchronized to token generation.
  5. Architect the system so that adding new model architectures and new SAEs is straightforward.

Non-goals (for v1)

  • Training SAEs.
  • Supporting every model architecture. Start with one, expand later.
  • A frontend UI. The SSE stream is the contract; the UI is a separate project.
  • Model steering (clamping/modifying SAE features to change output). Future work.
  • Multi-user / batched inference. Single-user local tool for now.

Model and SAE Selection

Decision rationale

The choice of starting model is driven by one constraint: a high-quality, pre-trained SAE must be freely downloadable for the exact model we run. Training an SAE is a multi-week, compute-intensive process and would block the project. Downloading one takes minutes.

Two realistic pairings exist as of March 2026:

Model SAE Source SAE Format SAE Coverage Notes
Gemma 2 2B Google Gemma Scope NPZ (numpy) All layers, all sub-layers (res, mlp, attn). 16K and 65K width variants. Best SAE coverage of any open model. 400+ SAEs freely available. Small enough to run on any machine.
Llama 3.1 8B Instruct Goodfire safetensors Layer 19 residual stream only. Single SAE. Larger model, only one layer instrumented, and Goodfire's API/demo has been deprecated.

Primary recommendation: Gemma 2 2B (instruction-tuned)

  • Model: google/gemma-2-2b-it from HuggingFace
  • SAE: google/gemma-scope-2b-pt-res — specifically the layer 20 residual stream SAE at 16K width (layer_20/width_16k/average_l0_71)
  • Why layer 20: Layer 20 of 26 is deep enough to capture high-level concepts (rather than surface syntax), and this specific SAE has an L0 of ~71 (meaning ~71 features fire per token on average), which provides a rich but not overwhelming signal.
  • Why 16K width: The 16K-width SAE has 16,384 latent features. This is large enough to capture fine-grained concepts but small enough that the encoder matrix (2304 × 16384 for Gemma 2 2B's hidden dim of 2304) is only ~144MB — trivially fast to compute.

Fallback options

If the implementer encounters blocking issues with Gemma 2 (e.g., mistral.rs Gemma 2 support has bugs, or the SAE format is harder to parse than expected), these alternatives are acceptable:

  1. Gemma 2 9B with Gemma Scope SAEs (google/gemma-scope-9b-pt-res). Same SAE ecosystem, larger model. Requires more RAM.
  2. Llama 3.1 8B Instruct with Goodfire's SAE (Goodfire/Llama-3.1-8B-Instruct-SAE-l19). Only layer 19 is available. The SAE weights are in safetensors format on HuggingFace. mistral.rs has mature Llama support. Note: Llama SAEs use safetensors format, not NPZ — would need a different loading path.

The implementer should not spend more than 2 days fighting a model/SAE integration before switching to a fallback.

Architecture

Component overview

┌──────────────────────────────────────────────────┐
│                   CLI (clap)                      │
│  neuroscope serve --model google/gemma-2-2b-it    │
│                   --sae google/gemma-scope-2b-... │
│                   --layer 20                      │
│                   --port 8080                     │
│                   --features-port 8081            │
└──────────┬───────────────────────┬───────────────┘
           │                       │
           ▼                       ▼
┌──────────────────┐   ┌──────────────────────────┐
│  Chat API Server │   │  Features SSE Server     │
│  (axum)          │   │  (axum)                  │
│  :8080           │   │  :8081                   │
│                  │   │                          │
│ POST /v1/chat/   │   │ GET /v1/features/stream  │
│   completions    │   │   → SSE event stream     │
│                  │   │                          │
│ GET /v1/models   │   │ GET /v1/features/labels  │
│                  │   │   → full label map JSON  │
└────────┬─────────┘   └────────┬─────────────────┘
         │                      │
         ▼                      │
┌────────────────────────────┐  │
│     Inference Engine       │  │
│     (mistral.rs core)      │  │
│                            │  │
│  ┌──────────────────────┐  │  │
│  │ Model Forward Pass   │  │  │
│  │ (Candle/Gemma2)      │  │  │
│  │                      │  │  │
│  │  layer 20 ──hook──►──│──│──│──► SAE Encoder
│  │                      │  │  │      │
│  └──────────────────────┘  │  │      ▼
│                            │  │  Top-K features
│  Token output ─────────────│──│──► broadcast
└────────────────────────────┘  │      channel
                                │      │
                                └──────┘

Key design decisions

Two ports, not one. The chat API server runs on :8080 and the features SSE server on :8081. Rationale: the chat API must be OpenAI-compatible for drop-in use with existing clients (Continue, Open WebUI, etc.). Mixing custom SSE streams into that API would break compatibility. A separate port keeps concerns clean and lets the features stream be consumed independently.

Broadcast channel for feature events. The inference engine publishes feature activations to a tokio::sync::broadcast channel. The SSE server subscribes to it. This decouples inference from network I/O and allows multiple SSE clients to connect simultaneously (useful if you want both a terminal logger and a web UI).

SAE encoder runs on same device as model. The SAE encoder weights are loaded as Candle tensors on the exact same Device instance as the model pipeline. This is critical: Candle uses pointer equality on Device::Metal variants, so two separately-constructed Device::Metal instances pointing at the same physical GPU are treated as different devices and matmul will fail. The engine must build the model pipeline first, extract its Device, then load SAE weights onto that device. The encoder forward pass happens synchronously during the model's forward pass, inside the hook callback.

Trait design for activation hooks

The ActivationHook trait is defined in the vendored mistralrs-core crate at vendor/mistral.rs/mistralrs-core/src/models/activation_hook.rs:

/// Trait for receiving intermediate activations during inference.
/// Implementations must be Send + Sync for use across async boundaries.
pub trait ActivationHook: Send + Sync {
    /// Called after each transformer layer's forward pass.
    /// `layer_idx`: 0-indexed layer number.
    /// `hidden_states`: the residual stream tensor at this point.
    fn on_residual_post(
        &self,
        layer_idx: usize,
        hidden_states: &candle_core::Tensor,
    ) -> candle_core::Result<()>;
}

/// No-op hook for non-instrumented use.
pub struct NoOpHook;
impl ActivationHook for NoOpHook {
    #[inline(always)]
    fn on_residual_post(&self, _: usize, _: &candle_core::Tensor) -> candle_core::Result<()> {
        Ok(())
    }
}

/// Type alias for the hook stored in models.
pub type SharedActivationHook = Arc<dyn ActivationHook>;

The hook is stored as Arc<dyn ActivationHook> (trait object with dynamic dispatch), not a generic parameter. This means there is a small vtable overhead per hook call, but it simplifies the pipeline API — the hook can be swapped at runtime without recompiling the model. The Gemma 2 model's forward pass calls self.activation_hook.on_residual_post(i, &xs) after each transformer layer.

The SaeHook implementation lives in neuroscope-engine/src/hooks.rs (not in neuroscope-core).

SAE module

pub struct SparseAutoencoder {
    /// Encoder weights: [hidden_dim, sae_width]
    encoder_weight: Tensor,
    /// Encoder bias: [sae_width]
    encoder_bias: Tensor,
    /// Per-feature JumpReLU thresholds: [sae_width]
    threshold: Tensor,
    /// Number of top activations to retain per token
    k: usize,
    /// SAE width (number of latent features)
    width: usize,
    /// Hidden dimension size
    hidden_dim: usize,
    /// Human-readable labels for each latent index
    labels: Vec<String>,
}

pub struct FeatureActivation {
    pub token_index: usize,
    pub token: String,
    pub layer: usize,
    pub top_features: Vec<ActiveFeature>,
    pub filtered_count: usize,  // features suppressed by calibration filter
}

pub struct ActiveFeature {
    pub index: usize,
    pub label: String,
    pub activation: f32,
    pub surprise: Option<f32>,  // z-score vs. calibration stats
}

pub enum FeatureEvent {
    Activation(FeatureActivation),
    GenerationComplete { total_tokens: usize, model: String },
}

SAE weight format: Gemma Scope SAEs are distributed as NPZ (numpy) files, not safetensors. The NPZ file contains three tensors with capitalized keys: W_enc ([hidden_dim, sae_width]), b_enc ([sae_width]), and threshold ([sae_width]). These are loaded using candle_core::npy::NpzTensors.

The SAE encoder forward pass:

  1. Receive hidden state tensor of shape [1, hidden_dim] (single token during autoregressive generation).
  2. Normalize the hidden state to unit RMS: x_normalized = x / sqrt(mean(x^2)). Gemma Scope SAEs were trained on normalized inputs.
  3. Compute z = x_normalized @ W_enc + b_enc.
  4. Apply JumpReLU activation: output = z * H(z - threshold) where H is the Heaviside step function. Values below the per-feature threshold are zeroed; values above pass through at their original magnitude.
  5. If a FeatureFilterer is present, apply calibration-based filtering (frequency/surprise/combined) before top-K selection. Otherwise, extract top-K activated features by raw value.
  6. Publish FeatureEvent::Activation to the broadcast channel.

Note: We only need the encoder half of the SAE. The decoder is not used for observation (it's only needed for reconstruction or steering). This halves memory requirements.

Configuration

pub struct NeuroscopeConfig {
    pub model_id: String,               // "google/gemma-2-2b-it"
    pub sae_repo: String,               // "google/gemma-scope-2b-pt-res"
    pub sae_path: String,               // "layer_20/width_16k/average_l0_71/params.npz"
    pub sae_local_path: Option<String>,  // Optional local path (bypasses HF download)
    pub target_layer: usize,            // 20
    pub top_k: usize,                   // 10
    pub neuronpedia_model_id: String,   // "gemma-2-2b"
    pub neuronpedia_sae_id: String,     // "20-gemmascope-res-16k"
    pub broadcast_buffer_size: usize,   // 1024
    pub feature_filter: String,         // "none", "frequency", "surprise", "combined"
    pub filter_threshold: f32,          // 0.5
    pub calibration_path: Option<String>, // Optional explicit calibration file path
}

All fields have sensible defaults via Default impl, so the simplest invocation requires no arguments.

SSE Protocol

Endpoint: GET /v1/features/stream

Returns an SSE stream. Each event corresponds to one generated token.

event: feature_activation
data: {
  "token_index": 0,
  "token": "Hello",
  "layer": 20,
  "top_features": [
    {"index": 4521, "label": "greeting or salutation", "activation": 3.82},
    {"index": 12033, "label": "politeness and social norms", "activation": 2.14},
    {"index": 891, "label": "beginning of conversation", "activation": 1.97}
  ]
}

event: feature_activation
data: {
  "token_index": 1,
  "token": "!",
  "layer": 20,
  "top_features": [
    {"index": 7744, "label": "exclamation or emphasis", "activation": 4.11},
    {"index": 4521, "label": "greeting or salutation", "activation": 1.23}
  ]
}

event: generation_complete
data: {"total_tokens": 42, "model": "google/gemma-2-2b-it"}

Endpoint: GET /v1/features/labels

Returns the complete label map as JSON. Useful for UIs that want to pre-fetch all labels.

{
  "sae_id": "google/gemma-scope-2b-pt-res/layer_20/width_16k/average_l0_71",
  "width": 16384,
  "labels": {
    "0": "unknown/unlabeled",
    "1": "articles and determiners",
    "4521": "greeting or salutation",
    ...
  }
}

Endpoint: POST /v1/chat/completions

Standard OpenAI-compatible chat completions. Supports streaming ("stream": true). When streaming is enabled, SSE feature events on the features port are synchronized with the token stream on the chat port — each token's feature activation event is published before or simultaneously with the corresponding token SSE chunk.

Feature Labels

Gemma Scope SAEs do not ship with human-readable labels. Labels must be sourced separately.

Label sources (in priority order)

  1. Local auto-interp labels (generated by neuroscope labels generate, see below). These are the highest quality because they use max-activating examples from the actual model+SAE pairing.
  2. Neuronpedia auto-interp labels. Neuronpedia has labeled many Gemma Scope SAEs. If available, download and cache the label JSON. These are decent but can be noisy — some labels are misleading or overly generic.
  3. Numeric fallback (e.g., "feature_4521"). The system must work without labels.

The label file is a simple JSON mapping from integer index to string. It is loaded at startup and held in memory.

Auto-interp label generation: neuroscope labels generate

Generate high-quality human-readable labels for SAE features by collecting max-activating examples and sending them to Claude for description.

Pipeline overview

Calibration corpus (text samples)
        ↓
Run model + SAE hook on each sample, record all feature activations
        ↓
For each feature, collect the top N max-activating examples
        ↓
Send each feature's examples to Claude API for description
        ↓
Score each label with detection scoring
        ↓
Write labels JSON to disk cache

Step 1: Collect max-activating examples

Run the model with SAE instrumentation over a calibration corpus. For each of the 16,384 features, maintain a max-heap of the top N tokens that most strongly activated that feature.

For each max-activating example, store:

  • The token text that triggered the activation
  • 24 tokens of surrounding context (12 before, 12 after) so the labeling LLM can see the token in context
  • The activation value

Corpus selection: Use a diverse text corpus — a mix of Wikipedia, code, conversation, and general web text. The corpus should be large enough that most features fire at least a few times. Based on Gemma Scope's L0 of ~71 (71 features fire per token on average), processing ~50K text samples should give reasonable coverage of 16K features. If --corpus-path is omitted, the system auto-downloads WikiText-103 via the HuggingFace Datasets API.

Number of examples per feature: Collect the top 20 max-activating examples per feature (following Bills et al. / OpenAI's methodology). Also collect 5 random-activation examples (sampled from across the activation distribution, not just the top) to catch polysemantic features that behave differently at different activation levels (per Anthropic's recommendation).

Step 2: Generate labels with Claude

For each feature, send a prompt to Claude containing:

  • The top 20 max-activating examples, each showing the token highlighted within its surrounding context and the activation magnitude
  • The 5 random-activation examples for contrast

The prompt should ask Claude to:

  1. Identify what pattern these tokens/contexts have in common
  2. Produce a concise label (under 15 words) describing what the feature detects
  3. Rate confidence (high/medium/low) based on how consistent the examples are

Use claude-haiku-4-5-20251001 for cost efficiency — labeling 16K features with Opus would be expensive and Haiku is sufficient for pattern description.

Parallelism: Run up to 50 concurrent API calls to keep total time under 30 minutes.

Step 3: Score labels with detection scoring

Use EleutherAI's detection scoring method to validate each label:

  • Show Claude the generated label, then present 20 non-activating text sequences and 50 activating text sequences
  • Ask it to classify each sequence as "contains this feature" or "does not contain this feature"
  • Compute balanced accuracy as the score

Labels scoring below 60% balanced accuracy should be flagged as low-confidence and can optionally be replaced with numeric fallback labels.

This scoring step can be run as a separate pass (neuroscope labels score) so labels can be regenerated without re-scoring.

Step 4: Write labels

Save to ~/.cache/neuroscope/labels/{model}_{sae}_autointerp.json. This file takes priority over Neuronpedia labels when both exist.

CLI interface

neuroscope labels generate \
  --model google/gemma-2-2b-it \
  --sae google/gemma-scope-2b-pt-res \
  --sae-path layer_20/width_16k/average_l0_71/params.npz \
  --corpus-path /path/to/texts/ \   # optional: auto-downloads WikiText if omitted
  --samples 50000 \                 # number of text samples to process
  --examples-per-feature 20 \       # max-activating examples to collect
  --labeler claude-haiku-4-5-20251001 \  # Claude model for label generation
  --concurrency 50 \                # parallel API calls
  --score                           # also run detection scoring

neuroscope labels score \           # score existing labels
  --labels ~/.cache/neuroscope/labels/gemma-2-2b_20-gemmascope-res-16k_autointerp.json

neuroscope labels show \            # preview labels for specific features
  --features 4521,3022,11612

Runtime estimates

  • Corpus pass: 50K samples through Gemma 2 2B with SAE hook on Metal ≈ 2-4 hours
  • Label generation: 16K Claude Haiku API calls at 50 concurrency ≈ 15-30 minutes, ~$3-8
  • Detection scoring (optional): 16K × 70 classifications ≈ 1-2 hours, ~$5-15
  • Total: ~3-6 hours for a full run

Design notes

  • The calibration corpus pass doubles as input for the feature frequency calibration (see below) — run both from the same pass to avoid redundant computation.
  • In the future, this pipeline could also extract the SAE decoder weights to compute top logit weights (which tokens a feature boosts in the output distribution). Neuronpedia's np_max-act-logits method shows this significantly improves label quality. This requires loading the decoder half of the SAE, which we currently skip.

Feature Frequency Calibration

Problem

Some SAE features fire on nearly every token regardless of content. These "always-on" features dominate the top-K list but carry no per-token information. In our testing, features like index 11612 ("references to legal actions...") fire at 40-60 activation on every single token — the label is almost certainly wrong, and the feature likely represents something generic like "instruction-following mode" or a baseline bias in the residual stream.

Anthropic's research categorizes features by log density (the fraction of tokens they fire on):

  • Dead features: Not active over 10^7 tokens (~2% of features for well-trained SAEs)
  • Ultralow-frequency features (density < 10^-4): Often uninterpretable, take up SAE capacity without contributing useful signal
  • High-frequency features: Fire on most tokens, represent generic patterns rather than specific concepts
  • Sweet spot: Features with moderate density that fire selectively

Solution: Calibration pass + frequency-aware filtering

Step 1: Calibration pass

At server startup (or as a separate CLI command), run a small calibration corpus through the model+SAE and compute per-feature statistics:

For each of the 16,384 features, record:

  • Firing rate: Fraction of tokens on which activation > 0
  • Mean activation: Average activation value when the feature fires
  • Activation variance: How much the activation varies across tokens

Store these statistics in ~/.cache/neuroscope/calibration/{model}_{sae}_stats.json.

Corpus size for calibration: ~1,000 text samples is sufficient to get stable frequency estimates. This is much smaller than the auto-interp corpus and should take 5-10 minutes on Metal.

Corpus acquisition: If --corpus-path is not provided, the system automatically downloads WikiText-103 from the HuggingFace Datasets REST API (no Python dependency required) and caches it at ~/.cache/neuroscope/corpus/. The neuroscope corpus build command can also be used to pre-download or customize the corpus dataset.

Step 2: Feature filtering modes

The server should support configurable filtering via a --feature-filter flag:

neuroscope serve \
  --feature-filter frequency \    # default: filter by firing rate
  --filter-threshold 0.5 \        # suppress features firing on >50% of tokens
  ...

Filtering strategies (configurable):

  1. Frequency filter (default): Suppress features with firing rate above a threshold (default: 0.5). A feature that fires on >50% of tokens is unlikely to carry token-specific information.

  2. Surprise filter: Instead of reporting features by raw activation, report by surprise — how much the current activation deviates from the feature's mean activation. A feature that always fires at 50 but suddenly fires at 200 is interesting; a feature that always fires at 50 is not.

    surprise_score = (activation - mean_activation) / sqrt(variance + epsilon)

    This is effectively a z-score. Features are ranked by surprise_score instead of raw activation for the top-K selection.

  3. Combined: Apply frequency filter first (remove always-on features), then rank remaining features by raw activation. This is the simplest approach that works well.

Step 3: Runtime behavior

During inference, when the SAE hook produces top-K features:

  1. Look up each feature's calibration stats
  2. Apply the configured filter
  3. If a feature is filtered out, skip it and take the next-highest activation
  4. Include a filtered count in the SSE event so the UI knows features were suppressed

Updated SSE event format:

{
  "token_index": 0,
  "token": "Hello",
  "layer": 20,
  "top_features": [
    {"index": 4521, "label": "greeting or salutation", "activation": 3.82, "surprise": 2.1},
    ...
  ],
  "filtered_count": 3
}

CLI interface

# Corpus auto-downloads if --corpus-path is omitted
neuroscope calibrate run \
  --model google/gemma-2-2b-it \
  --sae google/gemma-scope-2b-pt-res \
  --sae-path layer_20/width_16k/average_l0_71/params.npz \
  --samples 1000                    # text samples for calibration

# Or bring your own corpus
neuroscope calibrate run \
  --corpus-path /path/to/texts/ \
  --samples 1000

neuroscope calibrate show \         # inspect calibration results
  --stats ~/.cache/neuroscope/calibration/stats.json \
  --top 20                          # show 20 highest-frequency features

# Pre-download a corpus manually
neuroscope corpus build --samples 2000
neuroscope corpus show              # check what's cached

Shared calibration corpus

The calibration pass and the auto-interp max-activating example collection should share the same corpus run. When running neuroscope labels generate, the calibration stats are computed as a free byproduct — every token's activations are already being processed. The labels generate command should automatically produce the calibration stats file alongside the labels.

Input normalization

Gemma Scope's paper recommends normalizing input activations to unit mean squared norm before feeding them to the SAE encoder. Specifically: scale the hidden state vector so its average squared L2 norm equals the residual stream dimension (2304 for Gemma 2 2B). The Gemma Scope SAEs were trained on normalized inputs, so feeding unnormalized activations may produce inflated or distorted activation values.

This normalization should be applied in SparseAutoencoder::encode() before the matmul. The normalization factor can be precomputed during the calibration pass or computed per-token (the per-token overhead is trivial — one norm computation on a 2304-dim vector).

CLI Interface

# === Primary command: start the server ===
neuroscope serve \
  --model google/gemma-2-2b-it \
  --sae google/gemma-scope-2b-pt-res \
  --sae-layer 20 \
  --sae-width 16k \
  --top-k 10 \
  --port 8080 \
  --features-port 8081 \
  --device auto \
  --quantization q4k \
  --feature-filter frequency \     # filter out always-on features
  --filter-threshold 0.5           # suppress features firing >50% of tokens

# === Model/SAE management ===
neuroscope models list          # list downloaded models
neuroscope models pull <id>     # download model from HuggingFace
neuroscope sae list             # list downloaded SAEs
neuroscope sae pull <id>        # download SAE from HuggingFace

# === Corpus management ===
neuroscope corpus build         # download WikiText-103 (auto-done by calibrate/labels if needed)
neuroscope corpus build --samples 5000 --dataset wikitext --config wikitext-103-raw-v1
neuroscope corpus show          # show cached corpus info

# === Feature calibration (compute per-feature firing statistics) ===
neuroscope calibrate run \      # corpus auto-downloads if --corpus-path omitted
  --model google/gemma-2-2b-it \
  --sae google/gemma-scope-2b-pt-res \
  --samples 1000
neuroscope calibrate show --stats path/to/stats.json  # inspect calibration stats

# === Label generation (auto-interp via Claude) ===
neuroscope labels generate \    # corpus auto-downloads if --corpus-path omitted
  --model google/gemma-2-2b-it \
  --sae google/gemma-scope-2b-pt-res \
  --samples 50000 \
  --labeler claude-haiku-4-5-20251001 \
  --score                       # also run detection scoring
neuroscope labels score --labels path/to/labels.json  # score existing labels
neuroscope labels show --features 4521,3022  # preview specific labels

The --device auto flag should detect available hardware: prefer CUDA if available, then Metal, then CPU.

Test Plan

Philosophy

Write tests first. The project starts by building an end-to-end test suite that defines the expected behavior. Then build the tool. Then get the tests passing. Tests are the spec made executable.

Test categories

1. Unit tests: SAE module

These tests verify the SAE encoder logic in isolation, independent of any LLM.

test_sae_load_weights
  - Load a real Gemma Scope SAE NPZ file.
  - Assert encoder weight shape is [hidden_dim, sae_width] (e.g., [2304, 16384]).
  - Assert encoder bias shape is [sae_width].
  - Assert threshold shape is [sae_width].

test_sae_encode_known_input
  - Construct a synthetic hidden state tensor of shape [1, 2304] with known values.
  - Run the SAE encoder forward pass.
  - Assert output shape is [1, sae_width].
  - Assert the output is sparse (most values are zero or near-zero after JumpReLU).

test_sae_topk_extraction
  - Given a sparse encoded vector, extract top-10 features.
  - Assert exactly 10 features are returned.
  - Assert they are sorted by activation value descending.
  - Assert all returned activations are positive.

test_sae_jumprelu_activation
  - Verify JumpReLU behavior: values below the threshold are zeroed,
    values above are passed through (not just binary gated).
  - Use the threshold value from the SAE's config/weights.

test_sae_deterministic
  - Run the same input through the encoder twice.
  - Assert identical output both times.

2. Unit tests: Activation hook

test_hook_receives_correct_layer
  - Create a mock model forward pass that calls the hook at layers 0, 1, 2.
  - Register a hook for layer 1 only.
  - Assert the hook is called exactly once, with layer_idx=1.

test_hook_tensor_shape
  - Assert the tensor passed to the hook has shape [seq_len, hidden_dim]
    where hidden_dim matches the model config.

test_noop_hook_is_zero_cost
  - Compile with NoOpHook and verify (via benchmarking or inspection)
    that hook call sites are eliminated.
  - (This can be a benchmark test rather than a unit test.)

3. Integration tests: Inference + SAE

test_inference_produces_features
  - Load the real model (Gemma 2 2B) and real SAE.
  - Run a single chat completion: "Hello, how are you?"
  - Assert that feature activations are produced for each generated token.
  - Assert each activation event has the expected structure
    (token_index, token text, layer, list of features).

test_feature_activations_are_plausible
  - Send the prompt "What is the capital of France?"
  - In the generated response (which should mention "Paris"),
    collect all feature activations.
  - Assert that at least some feature labels (if available) relate
    to geography, countries, or proper nouns.
  - (This is a soft/heuristic test — it verifies the system is
    producing meaningful signal, not random noise.)

test_features_sync_with_tokens
  - Enable streaming on the chat endpoint.
  - Collect tokens and feature events.
  - Assert that for every token emitted on the chat stream,
    there is a corresponding feature activation event
    with a matching token_index.
  - Assert token_index values are monotonically increasing.

4. API tests: HTTP endpoints

test_chat_completions_openai_compat
  - POST /v1/chat/completions with a standard OpenAI-format request body.
  - Assert response matches OpenAI response schema
    (id, object, created, model, choices, usage).

test_chat_completions_streaming
  - POST /v1/chat/completions with "stream": true.
  - Assert response is SSE stream with "data: {...}" chunks.
  - Assert final chunk contains [DONE].

test_models_list
  - GET /v1/models.
  - Assert response lists the loaded model.

test_features_stream_sse
  - Connect to GET /v1/features/stream.
  - Trigger a chat completion on the other port.
  - Assert SSE events arrive with event type "feature_activation".
  - Assert events are valid JSON matching the FeatureActivation schema.

test_features_labels_endpoint
  - GET /v1/features/labels.
  - Assert response is JSON with sae_id, width, and labels fields.
  - Assert labels is a map from string keys (parseable as integers) to strings.

test_features_stream_no_generation
  - Connect to GET /v1/features/stream with no active generation.
  - Assert connection stays open (no immediate close or error).
  - Assert no events are sent until a generation begins.

5. Unit tests: Feature frequency calibration

test_calibration_computes_firing_rates
  - Run 100 synthetic inputs through a small SAE.
  - Assert firing rate for each feature is between 0.0 and 1.0.
  - Assert firing rates sum to roughly L0 * num_tokens / sae_width.

test_calibration_frequency_filter
  - Create calibration stats with known firing rates.
  - Set threshold to 0.5.
  - Assert features with firing rate > 0.5 are filtered from top-K.
  - Assert features with firing rate < 0.5 pass through.

test_calibration_surprise_scoring
  - Create calibration stats with known means and variances.
  - Given an activation vector, compute surprise scores.
  - Assert a feature with activation far above its mean scores higher
    than a feature with higher raw activation but near its mean.

test_calibration_persistence
  - Compute calibration stats and save to JSON.
  - Load from JSON and assert identical values.

6. Tests: Auto-interp label generation

test_max_activating_collection
  - Run a corpus through the model+SAE.
  - For each feature, assert max-activating examples are sorted
    by activation value descending.
  - Assert context windows are 24 tokens wide.
  - Assert at most N examples are stored per feature.

test_label_generation_prompt_format
  - Construct the prompt for a feature with known max-activating examples.
  - Assert the prompt includes token text, context, and activation values.
  - Assert the prompt asks for a concise label under 15 words.

test_label_scoring_detection
  - Given a label and a set of activating/non-activating examples,
    compute balanced accuracy.
  - Assert a good label (that correctly predicts activating examples)
    scores above 0.7.
  - Assert a random/bad label scores near 0.5.

test_generated_labels_override_neuronpedia
  - Save both Neuronpedia labels and auto-interp labels to cache.
  - Assert load_labels() returns the auto-interp labels (higher priority).

7. End-to-end smoke test

test_full_workflow
  - Start the server with a real model and SAE.
  - Connect an SSE client to the features stream.
  - Send a chat completion request.
  - Collect both the chat response and the feature events.
  - Assert the chat response is coherent text.
  - Assert feature events were received for every generated token.
  - Assert the server can handle a second request after the first completes.

Test infrastructure

  • Use cargo test with a #[cfg(test)] module in each crate.
  • Integration tests that require model weights should be gated behind a feature flag or env var (NEUROSCOPE_INTEGRATION_TESTS=1) so CI can skip them.
  • Unit tests for the SAE module should use small synthetic weights (e.g., random [64, 256] encoder) so they run in milliseconds without downloading anything.
  • API tests should spin up the server on a random port using tokio::test and make real HTTP requests.

Crate Structure

neuroscope/
├── Cargo.toml                 # workspace root (includes vendor/mistral.rs members)
├── spec.md                    # this file
├── crates/
│   ├── neuroscope-core/       # SAE module, feature types, label loading, calibration
│   │   ├── src/
│   │   │   ├── lib.rs
│   │   │   ├── sae.rs         # SparseAutoencoder struct, encode (with RMS normalization), topk
│   │   │   ├── types.rs       # FeatureActivation, ActiveFeature (with surprise), FeatureEvent
│   │   │   ├── labels.rs      # Label loading (auto-interp → Neuronpedia → fallback)
│   │   │   ├── calibration.rs # FeatureStats, CalibrationData, OnlineStatsCollector
│   │   │   ├── filter.rs      # FilterMode, FeatureFilterer (frequency/surprise/combined)
│   │   │   ├── corpus.rs      # Corpus download (HF Datasets API), loading, cache management
│   │   │   ├── max_activating.rs # MaxActivatingCollector (min-heap + reservoir sampling)
│   │   │   ├── claude_api.rs  # ClaudeClient (Anthropic Messages API)
│   │   │   ├── autointerp.rs  # Label prompt/response, concurrent generation, caching
│   │   │   └── scoring.rs     # Detection scoring (balanced accuracy)
│   │   └── tests/
│   │       └── sae_real_weights.rs  # Tests with real Gemma Scope weights
│   ├── neuroscope-engine/     # Inference engine wrapping mistral.rs
│   │   ├── src/
│   │   │   ├── lib.rs
│   │   │   ├── engine.rs      # InferenceEngine: load model, SAE, wire hook, load calibration
│   │   │   ├── hooks.rs       # SaeHook: implements ActivationHook with optional FeatureFilterer
│   │   │   ├── config.rs      # NeuroscopeConfig with defaults (incl. filter settings)
│   │   │   └── corpus.rs      # CorpusHook, CorpusRunner for batch corpus processing
│   │   └── Cargo.toml
│   ├── neuroscope-server/     # HTTP servers (chat API + features SSE)
│   │   ├── src/
│   │   │   ├── lib.rs         # run_servers() entrypoint
│   │   │   ├── chat_api.rs    # OpenAI-compatible endpoints
│   │   │   ├── features_api.rs # SSE stream + labels endpoint
│   │   │   └── state.rs       # AppState (shared engine + broadcast channel)
│   │   └── tests/
│   │       └── integration.rs # End-to-end integration tests
│   └── neuroscope-cli/        # CLI binary
│       ├── src/
│       │   └── main.rs        # clap CLI, subcommands
│       └── Cargo.toml
└── vendor/
    └── mistral.rs/            # Vendored mistral.rs with activation hook patches
        └── mistralrs-core/
            └── src/models/
                ├── activation_hook.rs  # ActivationHook trait, NoOpHook, SharedActivationHook
                └── gemma2.rs           # Gemma 2 model with hook call in forward pass

Note: The ActivationHook trait lives in the vendored mistral.rs, not in neuroscope-core. This keeps the trait definition close to the model code that calls it. The SaeHook implementation (which uses the trait) lives in neuroscope-engine.

Dependencies

Crate Purpose
mistralrs / mistralrs-core LLM inference engine (vendored, with activation hook patches)
candle-core, candle-nn Tensor operations, SAE encoder
axum, axum-extra HTTP servers with SSE support
tokio Async runtime
tokio-stream BroadcastStream wrapper for SSE fan-out
clap CLI argument parsing
serde, serde_json JSON serialization
hf-hub Downloading models/SAEs from HuggingFace
tracing, tracing-subscriber Structured logging
reqwest HTTP client for Neuronpedia labels, Datasets API corpus download, Claude API
dirs Platform-appropriate cache directory (~/.cache/neuroscope/)
indicatif Progress bars for calibration and label generation
fastrand Reservoir sampling in max-activating example collection

Implementation Order

  1. Set up workspace and crate structure. Get cargo build working with empty crates.
  2. Write all test stubs. Every test from the test plan above, with #[ignore] or todo!() bodies. This is the executable spec.
  3. Implement neuroscope-core: SAE loading, JumpReLU, top-K extraction, label loading. Get SAE unit tests passing.
  4. Implement neuroscope-engine: Wrap mistral.rs, add the activation hook to the Gemma 2 forward pass, wire SAE into the hook. Get integration tests passing.
  5. Implement neuroscope-server: Chat API and features SSE endpoints. Get API tests passing.
  6. Implement neuroscope-cli: Wire everything together with clap. Get the end-to-end smoke test passing.
  7. Polish: Error handling, graceful shutdown, --help text, README.

Modifying mistral.rs for Activation Hooks

This was the most invasive part of the project. The approach taken:

  1. Vendored mistral.rs into vendor/mistral.rs/ as a workspace member. The workspace root Cargo.toml includes the necessary mistral.rs crates (mistralrs, mistralrs-core, mistralrs-quant, etc.) and excludes unneeded ones (mistralrs-cli, mistralrs-server, etc.).
  2. Added activation_hook.rs in mistralrs-core/src/models/ defining the ActivationHook trait, NoOpHook, and SharedActivationHook type alias.
  3. Added set_activation_hook() method to the Pipeline trait (pipeline/mod.rs) and the NormalPipeline implementation (pipeline/normal.rs), plus the NormalModel trait (pipeline/loaders/normal_loaders.rs).
  4. Added the hook call inside the Gemma 2 model's forward pass (models/gemma2.rs): self.activation_hook.on_residual_post(i, &xs)?; — one line after each layer's forward.
  5. Re-exported ActivationHook from the top-level mistralrs crate for ergonomic imports.

The Qwen 2 model was also patched with hook support as a secondary architecture option.

If upstreaming to mistral.rs is desirable later, the hook points should be submitted as a separate PR with documentation explaining the use case.

Known Risks and Mitigations

Risk Mitigation
Gemma Scope SAE uses JumpReLU, which has a per-feature threshold parameter. We might load it wrong. The SAE NPZ file contains a threshold tensor. Loaded successfully — Gemma Scope uses capitalized keys (W_enc, b_enc, threshold).
Quantized model activations may not match the FP16/BF16 activations the SAE was trained on. Start with unquantized (BF16) inference. Only add quantization support after confirming features look reasonable.
Feature labels may not be available for our specific SAE variant on Neuronpedia. The system must work without labels. Numeric feature IDs are acceptable for v1.
mistral.rs's Gemma 2 implementation may have bugs or missing features. Fall back to Llama 3.1 8B + Goodfire SAE. This pairing is simpler (standard Llama architecture, single SAE file).
The SAE encoder adds latency to each token. For Gemma 2 2B with a 16K SAE, the encoder is a [2304 × 16384] matmul — about 75M FLOPs. This is negligible compared to the ~6B FLOPs per token for the model itself. Measure and confirm.
tokio::sync::broadcast may drop events under backpressure if the SSE client is slow. Use a large enough buffer (1024 events). Log dropped events. For v1, this is acceptable.

Implementation Order (v2 features)

After v1 is stable and the core serve pipeline works end-to-end:

  1. Input normalization: Add unit mean squared norm normalization to SparseAutoencoder::encode(). This is a one-line change that may significantly improve activation quality.
  2. Feature frequency calibration: Implement neuroscope calibrate command and frequency-based filtering in the serve pipeline. This is the fastest path to cleaner output.
  3. Auto-interp label generation: Implement neuroscope labels generate command. This shares the calibration corpus infrastructure, so do it after calibration is working.
  4. Detection scoring: Implement neuroscope labels score as a validation pass on generated labels.

Future Work (not in scope for v1 or v2)

  • Multi-layer instrumentation: Instrument several layers simultaneously and show how features evolve through the network.
  • Feature steering: Allow the UI to clamp feature activations to specific values, modifying the model's behavior in real-time.
  • SAE training pipeline: Train SAEs on new models directly from Neuroscope.
  • Quantization-aware SAEs: Investigate whether SAEs trained on FP16 activations work with Q4/Q8 quantized inference.
  • Decoder-assisted labeling: Load the SAE decoder weights to compute top logit weights (which output tokens a feature boosts). Neuronpedia's research shows this significantly improves auto-interp label quality.
  • Contrastive analysis: Compare feature activations between two similar prompts that differ in one key respect (per Anthropic's Circuit Tracing methodology) to identify causally relevant features.
  • Diff mode: Compare feature activations between two prompts or two models side by side.
  • Frontend UI: A web UI that consumes the SSE stream and visualizes features as a real-time heatmap, token-by-token timeline, or UMAP projection.