representation-geometry 0.1.0__tar.gz

representation_geometry-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Vidit Gupta
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

representation_geometry-0.1.0/PKG-INFO

@@ -0,0 +1,417 @@
+Metadata-Version: 2.4
+Name: representation-geometry
+Version: 0.1.0
+Summary: Reproducible representation-geometry analysis for transformer language models.
+Author: Vidit Gupta
+Maintainer: Vidit Gupta
+License: MIT
+Project-URL: Homepage, https://github.com/Vidit-01/representation-geometry
+Project-URL: Repository, https://github.com/Vidit-01/representation-geometry
+Project-URL: Issues, https://github.com/Vidit-01/representation-geometry/issues
+Project-URL: Documentation, https://github.com/Vidit-01/representation-geometry#readme
+Keywords: transformers,residual-stream,representation-geometry,effective-dimension,mechanistic-interpretability
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy>=1.23
+Requires-Dist: pandas>=1.5
+Requires-Dist: torch>=2.0
+Provides-Extra: hf
+Requires-Dist: accelerate>=0.26; extra == "hf"
+Requires-Dist: datasets>=2.14; extra == "hf"
+Requires-Dist: safetensors>=0.4; extra == "hf"
+Requires-Dist: transformers>=4.40; extra == "hf"
+Provides-Extra: plot
+Requires-Dist: matplotlib>=3.7; extra == "plot"
+Provides-Extra: dev
+Requires-Dist: build>=1.2; extra == "dev"
+Requires-Dist: pytest>=8.0; extra == "dev"
+Requires-Dist: ruff>=0.8; extra == "dev"
+Dynamic: license-file
+
+# 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:
+
+```bash
+python -m pip install -e ".[dev]"
+```
+
+For Hugging Face loading and examples:
+
+```bash
+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
+
+```python
+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:
+
+```python
+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:
+
+```python
+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:
+
+```python
+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:
+
+```python
+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:
+
+```python
+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:
+
+```python
+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:
+
+```python
+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:
+
+```text
+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:
+
+```python
+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:
+
+```bash
+python examples/hf_tiny_random_gpt2.py
+```
+
+## Examples
+
+Tiny in-memory model with no Hugging Face dependency:
+
+```bash
+python examples/tiny_in_memory.py
+```
+
+GPT-2 Hugging Face example:
+
+```bash
+python examples/minimal_hf_gpt2.py
+```
+
+Tiny Hugging Face model example:
+
+```bash
+python examples/hf_tiny_random_gpt2.py
+```
+
+Package smoke test:
+
+```bash
+python -m representation_geometry.smoke --out-dir .tmp/smoke
+```
+
+After installation, the console script is:
+
+```bash
+representation-geometry-smoke --out-dir .tmp/smoke
+```
+
+## Tests
+
+Run the test suite:
+
+```bash
+python -m pytest
+```
+
+Run linting:
+
+```bash
+python -m ruff check .
+```
+
+Build a wheel and source distribution:
+
+```bash
+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.