From e3160f0a850548bf2601aeb48b07964833e9f614 Mon Sep 17 00:00:00 2001 From: phusroyal Date: Wed, 5 Aug 2026 10:21:55 +0200 Subject: [PATCH 1/3] Feat: make fega native --- NOTICE.txt | 12 +- THIRD_PARTY_LICENSES/spherecluster-MIT.txt | 21 + .../reproductions/feature_effect_geometry.mdx | 131 + docs/src/pages/reproductions/index.astro | 17 +- .../artifacts/hoang2026_fega/atlas.json | 20703 ++++++++++++++++ .../rank_01_f14513/geometry.npz | Bin 0 -> 3749 bytes .../rank_01_f14513/metadata.json | 49 + .../rank_01_f14513/projection_2d.png | Bin 0 -> 39095 bytes .../rank_01_f14513/sphere_surface.png | Bin 0 -> 450192 bytes .../directed_ray/rank_01_f33760/geometry.npz | Bin 0 -> 24197 bytes .../directed_ray/rank_01_f33760/metadata.json | 49 + .../rank_01_f33760/projection_2d.png | Bin 0 -> 60126 bytes .../rank_01_f33760/sphere_surface.png | Bin 0 -> 460569 bytes .../rank_01_f54361/geometry.npz | Bin 0 -> 2520 bytes .../rank_01_f54361/metadata.json | 48 + .../rank_01_f54361/projection_2d.png | Bin 0 -> 93259 bytes .../rank_01_f54361/sphere_surface.png | Bin 0 -> 439251 bytes .../rank_01_f32542/geometry.npz | Bin 0 -> 24413 bytes .../rank_01_f32542/metadata.json | 49 + .../rank_01_f32542/projection_2d.png | Bin 0 -> 60124 bytes .../rank_01_f32542/sphere_surface.png | Bin 0 -> 488803 bytes .../rank_01_f1425/geometry.npz | Bin 0 -> 1883 bytes .../rank_01_f1425/metadata.json | 62 + .../rank_01_f1425/projection_2d.png | Bin 0 -> 37097 bytes .../rank_01_f1425/sphere_surface.png | Bin 0 -> 424146 bytes .../oneD_diffuse/rank_01_f59154/geometry.npz | Bin 0 -> 24316 bytes .../oneD_diffuse/rank_01_f59154/metadata.json | 54 + .../rank_01_f59154/projection_2d.png | Bin 0 -> 61683 bytes .../rank_01_f59154/sphere_surface.png | Bin 0 -> 462474 bytes .../rank_01_f19224/geometry.npz | Bin 0 -> 25471 bytes .../rank_01_f19224/metadata.json | 49 + .../rank_01_f19224/projection_2d.png | Bin 0 -> 88445 bytes .../rank_01_f19224/sphere_surface.png | Bin 0 -> 529642 bytes .../rank_01_f34636/geometry.npz | Bin 0 -> 24534 bytes .../rank_01_f34636/metadata.json | 48 + .../rank_01_f34636/projection_2d.png | Bin 0 -> 89810 bytes .../rank_01_f34636/sphere_surface.png | Bin 0 -> 657390 bytes .../hoang2026_fega/geometry_atlas.png | Bin 0 -> 917000 bytes .../artifacts/hoang2026_fega/index.json | 63 + .../curation/audit_fega_current_source.py | 1250 + .../curation/audit_fega_effect_call_path.py | 458 + .../curation/audit_fega_same_input.py | 883 + .../curation/curate_fega_reference.py | 484 + .../hoang2026_feature_effect_geometry.ipynb | 1995 ++ pyproject.toml | 16 +- src/murano/__init__.py | 47 +- src/murano/_optional.py | 1 + src/murano/backend.py | 5 + src/murano/fega/__init__.py | 33 + src/murano/fega/artifacts.py | 99 + src/murano/fega/checkpoints.py | 278 + src/murano/fega/config.py | 66 + src/murano/fega/contexts.py | 203 + src/murano/fega/dictionary_sae.py | 158 + src/murano/fega/effects.py | 265 + src/murano/fega/geometry.py | 308 + src/murano/fega/reporting.py | 450 + src/murano/fega/stability.py | 637 + src/murano/fega/visual_detail.py | 277 + src/murano/fega/visualization.py | 590 + src/murano/fega/vmf.py | 591 + src/murano/io.py | 38 + src/murano/keys.py | 7 + src/murano/model.py | 11 + src/murano/steps/__init__.py | 36 + src/murano/steps/fega.py | 809 + src/murano/steps/fega_analysis.py | 604 + src/murano/steps/sae.py | 20 +- tests/fixtures/fega/real_feature_33760.json | 48 + tests/fixtures/fega/real_feature_33760.npz | Bin 0 -> 22717 bytes tests/test_fega_analysis.py | 725 + tests/test_fega_pipeline.py | 727 + tests/test_fega_visualization.py | 155 + tests/test_imports.py | 28 + tests/test_io.py | 19 + tests/test_sae.py | 22 + uv.lock | 500 +- 77 files changed, 34185 insertions(+), 13 deletions(-) create mode 100644 THIRD_PARTY_LICENSES/spherecluster-MIT.txt create mode 100644 docs/src/pages/reproductions/feature_effect_geometry.mdx create mode 100644 notebooks/reproductions/artifacts/hoang2026_fega/atlas.json create mode 100644 notebooks/reproductions/artifacts/hoang2026_fega/candidates/axis_or_antipodal/rank_01_f14513/geometry.npz create mode 100644 notebooks/reproductions/artifacts/hoang2026_fega/candidates/axis_or_antipodal/rank_01_f14513/metadata.json create mode 100644 notebooks/reproductions/artifacts/hoang2026_fega/candidates/axis_or_antipodal/rank_01_f14513/projection_2d.png create mode 100644 notebooks/reproductions/artifacts/hoang2026_fega/candidates/axis_or_antipodal/rank_01_f14513/sphere_surface.png create mode 100644 notebooks/reproductions/artifacts/hoang2026_fega/candidates/directed_ray/rank_01_f33760/geometry.npz create mode 100644 notebooks/reproductions/artifacts/hoang2026_fega/candidates/directed_ray/rank_01_f33760/metadata.json create mode 100644 notebooks/reproductions/artifacts/hoang2026_fega/candidates/directed_ray/rank_01_f33760/projection_2d.png create mode 100644 notebooks/reproductions/artifacts/hoang2026_fega/candidates/directed_ray/rank_01_f33760/sphere_surface.png create mode 100644 notebooks/reproductions/artifacts/hoang2026_fega/candidates/global_2D_directional_subspace/rank_01_f54361/geometry.npz create mode 100644 notebooks/reproductions/artifacts/hoang2026_fega/candidates/global_2D_directional_subspace/rank_01_f54361/metadata.json create mode 100644 notebooks/reproductions/artifacts/hoang2026_fega/candidates/global_2D_directional_subspace/rank_01_f54361/projection_2d.png create mode 100644 notebooks/reproductions/artifacts/hoang2026_fega/candidates/global_2D_directional_subspace/rank_01_f54361/sphere_surface.png create mode 100644 notebooks/reproductions/artifacts/hoang2026_fega/candidates/global_kD_directional_subspace/rank_01_f32542/geometry.npz create mode 100644 notebooks/reproductions/artifacts/hoang2026_fega/candidates/global_kD_directional_subspace/rank_01_f32542/metadata.json create mode 100644 notebooks/reproductions/artifacts/hoang2026_fega/candidates/global_kD_directional_subspace/rank_01_f32542/projection_2d.png create mode 100644 notebooks/reproductions/artifacts/hoang2026_fega/candidates/global_kD_directional_subspace/rank_01_f32542/sphere_surface.png create mode 100644 notebooks/reproductions/artifacts/hoang2026_fega/candidates/multi_mode_directional_geometry/rank_01_f1425/geometry.npz create mode 100644 notebooks/reproductions/artifacts/hoang2026_fega/candidates/multi_mode_directional_geometry/rank_01_f1425/metadata.json create mode 100644 notebooks/reproductions/artifacts/hoang2026_fega/candidates/multi_mode_directional_geometry/rank_01_f1425/projection_2d.png create mode 100644 notebooks/reproductions/artifacts/hoang2026_fega/candidates/multi_mode_directional_geometry/rank_01_f1425/sphere_surface.png create mode 100644 notebooks/reproductions/artifacts/hoang2026_fega/candidates/oneD_diffuse/rank_01_f59154/geometry.npz create mode 100644 notebooks/reproductions/artifacts/hoang2026_fega/candidates/oneD_diffuse/rank_01_f59154/metadata.json create mode 100644 notebooks/reproductions/artifacts/hoang2026_fega/candidates/oneD_diffuse/rank_01_f59154/projection_2d.png create mode 100644 notebooks/reproductions/artifacts/hoang2026_fega/candidates/oneD_diffuse/rank_01_f59154/sphere_surface.png create mode 100644 notebooks/reproductions/artifacts/hoang2026_fega/candidates/residual_lowD_k/rank_01_f19224/geometry.npz create mode 100644 notebooks/reproductions/artifacts/hoang2026_fega/candidates/residual_lowD_k/rank_01_f19224/metadata.json create mode 100644 notebooks/reproductions/artifacts/hoang2026_fega/candidates/residual_lowD_k/rank_01_f19224/projection_2d.png create mode 100644 notebooks/reproductions/artifacts/hoang2026_fega/candidates/residual_lowD_k/rank_01_f19224/sphere_surface.png create mode 100644 notebooks/reproductions/artifacts/hoang2026_fega/candidates/unresolved_high_dimensional_or_diffuse/rank_01_f34636/geometry.npz create mode 100644 notebooks/reproductions/artifacts/hoang2026_fega/candidates/unresolved_high_dimensional_or_diffuse/rank_01_f34636/metadata.json create mode 100644 notebooks/reproductions/artifacts/hoang2026_fega/candidates/unresolved_high_dimensional_or_diffuse/rank_01_f34636/projection_2d.png create mode 100644 notebooks/reproductions/artifacts/hoang2026_fega/candidates/unresolved_high_dimensional_or_diffuse/rank_01_f34636/sphere_surface.png create mode 100644 notebooks/reproductions/artifacts/hoang2026_fega/geometry_atlas.png create mode 100644 notebooks/reproductions/artifacts/hoang2026_fega/index.json create mode 100644 notebooks/reproductions/curation/audit_fega_current_source.py create mode 100644 notebooks/reproductions/curation/audit_fega_effect_call_path.py create mode 100644 notebooks/reproductions/curation/audit_fega_same_input.py create mode 100644 notebooks/reproductions/curation/curate_fega_reference.py create mode 100644 notebooks/reproductions/hoang2026_feature_effect_geometry.ipynb create mode 100644 src/murano/fega/__init__.py create mode 100644 src/murano/fega/artifacts.py create mode 100644 src/murano/fega/checkpoints.py create mode 100644 src/murano/fega/config.py create mode 100644 src/murano/fega/contexts.py create mode 100644 src/murano/fega/dictionary_sae.py create mode 100644 src/murano/fega/effects.py create mode 100644 src/murano/fega/geometry.py create mode 100644 src/murano/fega/reporting.py create mode 100644 src/murano/fega/stability.py create mode 100644 src/murano/fega/visual_detail.py create mode 100644 src/murano/fega/visualization.py create mode 100644 src/murano/fega/vmf.py create mode 100644 src/murano/steps/fega.py create mode 100644 src/murano/steps/fega_analysis.py create mode 100644 tests/fixtures/fega/real_feature_33760.json create mode 100644 tests/fixtures/fega/real_feature_33760.npz create mode 100644 tests/test_fega_analysis.py create mode 100644 tests/test_fega_pipeline.py create mode 100644 tests/test_fega_visualization.py diff --git a/NOTICE.txt b/NOTICE.txt index f0f7b14..5911de5 100644 --- a/NOTICE.txt +++ b/NOTICE.txt @@ -9,10 +9,13 @@ file for the full license text. ------------------------------------------------------------------------------- Third party legal information -Murano does not bundle or redistribute third-party source code. It depends on -the following open-source projects, which are installed separately (via pip) and -remain under their own licenses. Refer to each project's own distribution for -the full license text: +Murano's dense vMF implementation adapts portions of spherecluster's +von-Mises--Fisher mixture implementation. See +THIRD_PARTY_LICENSES/spherecluster-MIT.txt for its license and attribution. + +Murano also depends on the following open-source projects, which are installed +separately (via pip) and remain under their own licenses. Refer to each +project's own distribution for the full license text: - nnsight (https://github.com/ndif-team/nnsight) - nnterp (https://github.com/Butanium/nnterp) @@ -26,4 +29,3 @@ the full license text: - seaborn (https://github.com/mwaskom/seaborn) - plotly (https://github.com/plotly/plotly.py) - sae-lens (https://github.com/jbloomAus/SAELens) - diff --git a/THIRD_PARTY_LICENSES/spherecluster-MIT.txt b/THIRD_PARTY_LICENSES/spherecluster-MIT.txt new file mode 100644 index 0000000..57439dc --- /dev/null +++ b/THIRD_PARTY_LICENSES/spherecluster-MIT.txt @@ -0,0 +1,21 @@ +The MIT License (MIT) + +Copyright (c) 2016 Clara Labs + +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. diff --git a/docs/src/pages/reproductions/feature_effect_geometry.mdx b/docs/src/pages/reproductions/feature_effect_geometry.mdx new file mode 100644 index 0000000..80a5c15 --- /dev/null +++ b/docs/src/pages/reproductions/feature_effect_geometry.mdx @@ -0,0 +1,131 @@ +--- +layout: ../../layouts/ReproLayout.astro +title: "Sparse Autoencoders Encode Both Concepts and Functions" +description: "An executable Murano demonstration of Feature-Effect Geometry Analysis, with active rerenders of compact source-derived geometry." +--- + +
+ ← Back to gallery + +
+ arXiv 2026 + feature effects + SAE + geometry +
+ +

+ Sparse Autoencoders Encode Both Concepts and Functions: The Downstream Geometry of Feature Effects +

+ +

+ Phu Gia Hoang, Anwoy Chatterjee, Tanmoy Chakraborty, Iryna Gurevych, Subhabrata Dutta +

+ + +
+ +
+ +
+ +## TL;DR + +Feature-Effect Geometry Analysis (FEGA) asks whether removing the same active sparse-autoencoder feature has a consistent output effect across contexts. The notebook runs all six FEGA analysis phases plus Murano visualization for Gemma 2 2B feature 19403 on 24 authored city–country prompts. It then actively rerenders compact source-derived geometry for every reported family. The live run is a small demonstration, not a reproduction of the paper's experiments. + +## Abstract + +Activation descriptions do not guarantee that an SAE feature has one reusable causal effect. Hoang et al. instead remove the same active latent in many contexts, collect the changes in output-logit space, and study the geometry of that effect cloud. FEGA distinguishes rays, axes, low-dimensional subspaces, multiple directional modes, and diffuse geometry. The paper reports that clean one-dimensional effects are uncommon: value-like features tend to be more structured, while context-dependent pointer-like features are usually diffuse. + +
+ +## Running FEGA with Murano + +The notebook uses `google/gemma-2-2b`, the standard width-2pow16 dictionary-learning SAE at residual layer 12, and one feature that is active on all 24 authored prompts. Feature 19403 is not presented as a paper-selected feature. + +### Step 1: Define the intervention + +For each prompt, Murano records the model readout after replacing the target residual with the complete SAE reconstruction. It then decodes the same latent vector with feature 19403 set to zero. The FEGA effect is `ablated - baseline` at the input to the output embedding. Every authored example targets the final token of a prompt such as `"city: Paris, country:"`. + +### Step 2: Run the analysis and visualization phases + +`fega_steps(phases="all")` runs data preparation, effect collection, geometry metrics, dense CPU vMF fitting, stability checks, geometry reporting, and visualization in dependency order. A list of phases can be supplied instead; Murano loads required checkpoints explicitly and does not silently rerun omitted phases. + +```python +results = Pipeline([ + LoadPrompts(prompts), + *fega_steps( + model, + release=SAE_REPO, + sae_id=SAE_ID, + sae_model=sae, + feature_ids=[FEATURE_ID], + phases="all", + batch_size=len(prompts), + max_contexts=len(prompts), + n_jobs=1, + top_k_per_family=1, + figures=("atlas", "sphere_surface", "projection_2d"), + output_dir=live_directory.name, + dpi=120, + ), +]).run() +``` + +### Step 3: Inspect the live result + +The executed notebook retained all 24 effects and assigned the authored population to `directed_ray`. The selected-family stability decision is `stable`, while the overall confidence remains `exploratory` because this is a 24-context run. + +| Live check | Result | +| --- | --- | +| Feature activation, min / median / max | 16.5 / 17.94 / 19.25 | +| Effect magnitude, min / median / max | 144.7 / 330.9 / 895.3 | +| Retained effects | 24 / 24 | +| Ray concentration, `c_ray` | 0.9329 | +| Leading-span share, `s_span(1)` | 0.9364 | +| Selected vMF components | 3 | +| Assignment stability | 0.9254 | +| Family / stability / confidence | `directed_ray` / `stable` / `exploratory` | +| Secondary flags | `lowD_candidate_blocked`, `span_selected_k` | +| Completed stages | six analysis phases plus visualization | +| Rendered figures | atlas, unit-sphere surface, 2D projection | +| Final notebook check | `FEGA_NOTEBOOK_OK` | + +### Step 4: Compare compact source geometry + +One feature cannot exercise every visual family. The second notebook section therefore loads compact exported arrays for eight representatives, rerenders them through Murano, and shows each copied source PNG as a labeled target. The primary standard-SAE run supplies seven families; a separately labeled Matryoshka run supplies the multi-mode example. + +The notebook produces one fresh atlas plus a sphere-surface and family-specific projection for every representative. The compact bundle begins after source geometry export and contains only per-feature effect-space kernels, derived coordinates, labels, metrics, and reference PNGs. It excludes prompts, model or SAE weights, full residual-space effect and Gram artifacts, and vocabulary coordinates. It cannot rerun prompt preparation, effect collection, vMF fitting, paper datasets, feature selection, or aggregate analyses. + +
+ +## Scope + +The live section demonstrates that native Murano FEGA executes end to end on a small authored population and exposes typed results for every phase. The compact section demonstrates Murano's atlas, unit-sphere surface, and family-specific projection renderers across the source families. + +It does **not** reproduce the paper's RAVEL or ICL datasets, paper-selected features, model and SAE sweep, aggregate family frequencies, or value-like versus pointer-like comparison. The live `directed_ray` result applies only to feature 19403 on these 24 prompts. Source-faithfulness checks are maintained separately from this public notebook; copied source figures are visual targets, not parity evidence. + +
+ +## Citation + +```bibtex +@misc{hoang2026sparseautoencodersencodeconcepts, + title = {Sparse Autoencoders Encode Both Concepts and Functions: The Downstream Geometry of Feature Effects}, + author = {Phu Gia Hoang and Anwoy Chatterjee and Tanmoy Chakraborty and Iryna Gurevych and Subhabrata Dutta}, + year = {2026}, + eprint = {2607.24645}, + archivePrefix = {arXiv}, + primaryClass = {cs.LG}, + url = {https://arxiv.org/abs/2607.24645}, +} +``` + +
diff --git a/docs/src/pages/reproductions/index.astro b/docs/src/pages/reproductions/index.astro index bd1d282..9d9546c 100644 --- a/docs/src/pages/reproductions/index.astro +++ b/docs/src/pages/reproductions/index.astro @@ -45,16 +45,27 @@ const reproductions = [ paper: 'https://arxiv.org/abs/2510.01025', code: 'https://github.com/UKPLab/supervised-multidimensional-scaling', }, + { + slug: 'feature_effect_geometry', + title: 'Sparse Autoencoders Encode Both Concepts and Functions: The Downstream Geometry of Feature Effects', + venue: 'arXiv 2026', + authors: 'Hoang, Chatterjee, Chakraborty, Gurevych, Dutta', + description: + 'Runs native Feature-Effect Geometry Analysis on an authored prompt set and actively rerenders compact source-derived geometry for every reported family.', + tags: ['feature effects', 'geometry', 'SAE'], + paper: 'https://arxiv.org/abs/2607.24645', + code: 'https://github.com/UKPLab/FEGA', + }, ]; --- - +