representation-geometry

representation-geometry is a Python toolkit for measuring representation geometry inside transformer models. It collects internal activations with PyTorch hooks, updates streaming covariance statistics, computes spectral and subspace measurements, and saves reproducible artifact bundles.

The package follows the project vision in VISION.md: it provides measurement infrastructure, not paper-specific interpretation. It is model-agnostic where possible, stores reusable quantitative outputs, and stays isolated from manuscript code, figures, tables, and LaTeX files.

Features (Current)

• Analyze loaded torch.nn.Module models or Hugging Face causal language model ids.
• Collect residual block inputs, residual block outputs, or named submodule inputs/outputs.
• Accumulate covariance statistics with a streaming Welford-style backend.
• Compute spectral geometry metrics, sample-SVD metrics, and between-layer novelty metrics.
• Optionally run a normalization ablation diagnostic.
• Optionally collect simple MoE router load statistics when model outputs expose router logits.
• Save JSON, CSV, or bundle artifacts with versioned metadata, metrics, diagnostics, router stats, and eigenvalues.

Architecture Overview

The package is organized into four layers that can evolve independently:

1. Activation collection through PyTorch hooks.
2. Streaming statistical backends for covariance and sampled activations.
3. Geometry metrics and diagnostics that consume those statistical objects.
4. Artifact generation for DataFrames, JSON, CSV, NumPy arrays, and reproducibility metadata.

This separation is intentional. New measurements should be implemented against the statistical backend rather than directly inside model-specific collection code.

Why Streaming?

Many activation-analysis scripts store all hidden states before computing metrics. representation-geometry instead updates statistics incrementally whenever possible. The current covariance backend keeps streaming means, covariance numerators, and a bounded activation sample, which allows larger token budgets without keeping every activation in memory.

Dense covariance is still hidden_dim x hidden_dim, so very wide models require care, but the package is designed around bounded-memory measurement rather than one-off activation dumps.

Reproducibility

Every saved artifact bundle records the measurement protocol, not just the final numbers. Metadata includes the artifact schema version, model information, hook targets, selected layers, token budget, observed tokens, metric and diagnostic configuration, tokenizer and dataset labels when supplied, software versions, runtime, and optional user metadata.

The goal is for downstream research projects to reproduce the measurement independently of the original paper scripts.

Install

From this directory:

python -m pip install -e ".[dev]"

For Hugging Face loading and examples:

python -m pip install -e ".[dev,hf]"

The core package depends on numpy, pandas, and torch. The hf extra installs transformers, accelerate, datasets, and safetensors.

Development tooling is included in the dev extra: pytest, ruff, and build.

Quick Start

from representation_geometry import analyze_model

results = analyze_model(
    model_name=model,
    dataloader=dataloader,
    max_tokens=100_000,
    layers="all",
    hook_point="residual_input",
    metrics="default",
    save_mode="bundle",
    save_dir="runs/my_model",
    tokenizer_name="my-tokenizer",
    dataset_name="my-dataset",
)

print(results.metrics.head())
print(results.novelty.head())
print(results.metadata)

model_name may be:

• a loaded torch.nn.Module,
• a Hugging Face model id string, if transformers is installed,
• or omitted when using the explicit model= argument.

For large models, prefer passing a preloaded model so you control dtype, device placement, quantization, and device_map.

Dataloader Format

dataloader must be an iterable of model batches. Dictionary batches with input_ids are preferred:

yield {
    "input_ids": input_ids,
    "attention_mask": attention_mask,
}

The package moves tensors to the selected input device recursively. It calls the model under torch.no_grad() and tries use_cache=False for transformer-style forward methods.

max_tokens is a stopping budget based on observed batch tokens. The current implementation processes whole batches, so tokens_observed can exceed max_tokens on the final batch.

Hook Points

The default hook is the input to each transformer block:

analyze_model(model_name=model, dataloader=dataloader, hook_point="residual_input")

Supported hook points are:

• residual_input: forward pre-hook on transformer blocks.
• residual_output: forward hook on transformer blocks.
• module_input:<module_path>: forward pre-hook on a named submodule.
• module_output:<module_path>: forward hook on a named submodule.

Examples:

analyze_model(model_name=model, dataloader=dataloader, layers=[0, 3, 7])

analyze_model(
    model_name=model,
    dataloader=dataloader,
    hook_point="module_output:blocks.0.proj",
)

For residual hooks, transformer blocks are auto-detected from common containers such as .transformer.h, .model.layers, .gpt_neox.layers, .decoder.layers, .layers, .blocks, or .h.

For named module hooks, use paths from dict(model.named_modules()). Named module hooks currently collect a single target and should be called with the default layers="all".

Metrics

metrics="default", metrics=None, or metrics=[] computes all currently implemented metric families:

• spectrum
• sample
• novelty

You can request a narrower set:

analyze_model(model_name=model, dataloader=dataloader, metrics=["spectrum"])

Spectrum Metrics

Computed from covariance eigenvalues:

• d_pr
• d_pr_norm
• d_entropy
• d_entropy_norm
• stable_rank
• top1_share
• top10_share
• top32_share
• k90
• k95
• log_spectrum_slope_10_200

The near-rank-1 criterion used by the original paper workflow can be computed downstream:

near_rank1 = (results.metrics["top1_share"] >= 0.99) & (results.metrics["k90"] <= 1)

Activation Sample Metrics

Computed from a bounded activation sample stored by the streaming accumulator:

• sv_effective_rank
• sv_stable_rank
• sv_participation_ratio
• cos_sim_mean
• cos_sim_std
• ve_4
• ve_8
• ve_16
• ve_32

The sample size is controlled by sample_limit, default 512.

Novelty Metrics

Computed between consecutive collected layers:

• feature_subspace_novelty
• linear_cka
• layer_cosine_sim

novelty_k controls the number of top covariance eigenvectors used for feature-subspace novelty.

Diagnostics

Currently supported diagnostics:

diagnostics=["normalization_ablation"]

This diagnostic recomputes spectrum metrics on sampled activations under:

• raw
• token_l2
• token_rms
• feature_standardized

Diagnostics are stored in results.diagnostics and saved under diagnostics/ when save_mode is enabled.

Extending The Package

The current public API does not yet accept arbitrary user-defined metrics at runtime. Internally, however, metrics are already separated from activation collection: api.py collects activations into streaming statistics, and metrics.py computes measurements from those statistics.

To add a new built-in metric today, implement it as a focused function in metrics.py, thread it through the metric-family selection in api.py, add tests, and document the new output columns. Diagnostics follow the same pattern through the small DIAGNOSTIC_REGISTRY in diagnostics.py.

This is a stepping stone toward the longer-term vision of a registry-style metric API.

Outputs

analyze_model returns an AnalysisResults object:

results.metrics      # pandas.DataFrame
results.novelty      # pandas.DataFrame
results.router       # pandas.DataFrame, empty unless router stats are available
results.diagnostics  # dict[str, pandas.DataFrame]
results.eigenvalues  # dict[str, numpy.ndarray]
results.metadata     # dict
results.artifact_dir # pathlib.Path after save()

Supported save modes:

• save_mode=None: no files are written.
• save_mode="json": writes JSON tables.
• save_mode="csv": writes CSV tables.
• save_mode="bundle": writes CSV tables plus bundle metadata and eigenvalues.

Saved bundle contents:

metrics.csv or metrics.json
novelty.csv or novelty.json
router.csv or router.json          # only when router stats are available
diagnostics/*                      # only when diagnostics are requested
eigenvalues.npz
metadata.json

Metadata includes artifact_schema_version, model id, model class, model config summary, token budget, observed tokens, batch count, hook point, hook targets, selected layers, metric names, diagnostic names, tokenizer label, dataset label, input device, software versions, runtime, save mode, and optional run_metadata.

Hugging Face Usage

Small runnable example:

from pathlib import Path

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

from representation_geometry import analyze_model

model_id = "hf-internal-testing/tiny-random-gpt2"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id)
model.eval()

tokens = tokenizer(
    "Residual geometry measures covariance structure inside transformer layers.",
    return_tensors="pt",
)

results = analyze_model(
    model_name=model,
    dataloader=[tokens],
    max_tokens=16,
    layers=[0],
    save_mode="json",
    save_dir=Path(".tmp") / "hf_usage_example",
    tokenizer_name=model_id,
    dataset_name="inline example sentence",
)

print(results.metrics[["layer", "d_pr_norm", "top1_share", "k90"]])

The same pattern is available as:

python examples/hf_tiny_random_gpt2.py

Examples

Tiny in-memory model with no Hugging Face dependency:

python examples/tiny_in_memory.py

GPT-2 Hugging Face example:

python examples/minimal_hf_gpt2.py

Tiny Hugging Face model example:

python examples/hf_tiny_random_gpt2.py

Package smoke test:

python -m representation_geometry.smoke --out-dir .tmp/smoke

After installation, the console script is:

representation-geometry-smoke --out-dir .tmp/smoke

Tests

Run the test suite:

python -m pytest

Run linting:

python -m ruff check .

Build a wheel and source distribution:

python -m build

The tests cover covariance streaming, metric formulas, diagnostics, JSON/bundle saving, named module hooks, residual output hooks, and metric subset selection.

Current Limitations

• analyze_model currently loads Hugging Face ids with AutoModelForCausalLM.
• Residual block discovery uses common transformer container names; uncommon architectures may need named module hooks.
• Named module hooks collect one target per call.
• Covariance matrices are stored on CPU as dense hidden_dim x hidden_dim tensors, so very large hidden dimensions still need care.
• Distributed collection, sketching, randomized eigensolvers, dashboards, multimodal-specific helpers, and causal intervention APIs are vision items, not implemented yet.

Design Boundary

This package owns reusable measurement infrastructure:

• activation collection,
• streaming statistics,
• geometry metrics,
• diagnostics,
• artifact serialization.

It deliberately does not own:

• paper prose,
• LaTeX generation,
• publication tables,
• manuscript figures,
• benchmark leaderboards,
• research conclusions.

Downstream projects should interpret the measurements; representation-geometry should make the measurements reproducible.