Add Logit Lens and Tuned Lens methods by bernasraphael · Pull Request #161 · FOR-sight-ai/interpreto

bernasraphael · 2026-08-20T14:10:56Z

Description

This PR adds Logit Lens and Tuned Lens methods for inspecting predictions decoded from intermediate transformer representations.

Both methods integrate with ModelWithSplitPoints: the lens operates on the wrapped model's configured split point and reuses the model's prediction path.
Main changes:

Add LogitLens for projecting an intermediate activation through the model's prediction path.
Add TunedLens with an affine translator tied to the configured split point.
Support causal language modeling, masked language modeling, and compatible single-label sequence classification.
Add conservative automatic projection-path resolution and explicit head_name, pre_head_name, and pooling_strategy overrides.
Add plot_lens, following the existing Interpreto visualization conventions.
Export the new public API from interpreto.
Add API documentation and an executed notebook covering language modeling, classification, fitting, metrics, and visualization.
Add focused tests for input validation, batching, padding, projection fidelity, device handling, metrics, checkpoint validation, and visualization alignment.

Related Issue

#85

Type of Change

📚 Examples / docs / tutorials / dependencies update
🚀 New feature (non-breaking change which adds functionality)

Checklist

I've read the CODE_OF_CONDUCT.md document.
I've read the CONTRIBUTING.md guide.
I've successfully run the style checks using make lint.
I've written tests for all new methods and classes that I created and successfully ran make test.
I've written the docstring in Google format for all the methods and classes that I used.

- add AllLayersSplitter to capture every residual-stream boundary - project all layer activations together through the native model head - skip redundant transformer computation while preserving output contracts - add zero-initialized residual translators for Tuned Lens - provide a simple single-text API for lens explanations and training - add lens visualization, documentation, notebook, exports, and tests

AntoninPoche

Really nice work, much smoother than the previous version!

My comments are mainly on the documentation.

Also, I do not think we need a BaseLens, if LogitLens correspond to it exactly.

Copilot

🟡 Changes recommended

Layer discovery is unsafe for some architectures, and several advertised APIs and executed documentation artifacts are missing.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

Adds Logit Lens and Tuned Lens support using a new all-layer transformer splitter, public visualization, tests, and documentation.

Changes:

Adds all-layer activation extraction and lens projection/training.
Adds normalized lens outputs and HTML visualization.
Exports and documents the new public APIs.

File summaries

File	Description
`tests/lens/test_lens.py`	Tests lens behavior and visualization.
`tests/concepts/splitters/test_all_layers_splitter.py`	Tests activation extraction and head replay.
`README.md`	Advertises lens methods and tutorial.
`mkdocs.yml`	Adds lens documentation navigation.
`interpreto/visualizations/lens.py`	Implements lens HTML rendering.
`interpreto/visualizations/commons.py`	Makes visualization JavaScript optional.
`interpreto/visualizations/__init__.py`	Exports `plot_lens`.
`interpreto/typing.py`	Defines lens result types.
`interpreto/lens/tuned_lens.py`	Implements Tuned Lens fitting.
`interpreto/lens/logit_lens.py`	Defines Logit Lens.
`interpreto/lens/_lens_base.py`	Provides shared lens inference logic.
`interpreto/lens/__init__.py`	Exports lens classes.
`interpreto/concepts/splitters/all_layers_splitter.py`	Adds all-layer extraction and prediction replay.
`interpreto/concepts/splitters/__init__.py`	Exports the new splitter.
`interpreto/concepts/__init__.py`	Re-exports the splitter.
`interpreto/__init__.py`	Adds top-level public exports.
`docs/notebooks/lens_notebook.ipynb`	Demonstrates lens usage.
`docs/api/visualizations.md`	Documents lens visualization.
`docs/api/lens/overview.md`	Introduces the lens API.
`docs/api/lens/methods/tuned_lens.md`	Documents Tuned Lens.
`docs/api/lens/methods/logit_lens.md`	Documents Logit Lens.
`docs/api/concepts/splitters/all_layers_splitter.md`	Documents the new splitter.

Review details

Files reviewed: 22/22 changed files
Comments generated: 5
Review effort level: Balanced

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

- Add to capture every transformer block boundary and reuse the model's native prediction path. - Add , decoding every residual state together through one prediction-head call. - Add with zero-initialized affine translators trained across all model depths. - Support causal language modeling and compatible sequence-classification models. - Add following Interpreto's visualization conventions. - Export the Lens API from . - Add consolidated API documentation and an executed notebook covering classification, generation, fitting, and visualization. - Add tests for projection fidelity, layer discovery, device handling, training, notebook execution, and visualization compatibility.

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

- accept model repository IDs directly in LogitLens and TunedLens - replace per-layer tables with compact confidence grids - preserve top-k details through interactive hover information - expand the executed notebook with dataset-size and epoch comparisons - update documentation, dependencies, and regression tests

- keep next-token predictions as the default generation behavior - add align=True for target-aligned visualizations - improve whitespace rendering and update lens documentation and tests

- align causal lens predictions by default - distinguish input, embeddings, and layer outputs - highlight correct predictions and document unaligned views

- adapt lens colors to the active documentation theme - improve prediction text contrast - refresh notebook outputs and add regression coverage

AntoninPoche

Looks perfect now, thank you very much for your contribution and your patience!!

Add Logit Lens and Tuned Lens methods

b902f65

bernasraphael requested a review from AntoninPoche August 20, 2026 14:11

AntoninPoche and others added 6 commits August 26, 2026 09:51

splitters: initialize the all layers splitter

08ca4ce

all layers splitters: add apply_head

9c8b0f1

all layers splitter: fix

dde35c1

fixup! all layers splitter: fix

eee8564

chore: merge fixes into lenses

f24bb91

AntoninPoche changed the base branch from main to fixes September 15, 2026 15:26

AntoninPoche requested changes Sep 15, 2026

View reviewed changes

AntoninPoche linked an issue Sep 15, 2026 that may be closed by this pull request

Include LogitLens #85

Closed

AntoninPoche requested a balanced review from Copilot September 15, 2026 16:47

Copilot started reviewing on behalf of AntoninPoche September 15, 2026 16:47 View session

Copilot AI reviewed Sep 15, 2026

View reviewed changes

Comment thread interpreto/concepts/splitters/all_layers_splitter.py Outdated

Comment thread interpreto/concepts/splitters/all_layers_splitter.py

Comment thread interpreto/lens/tuned_lens.py

Comment thread tests/lens/test_lens.py Outdated

Comment thread docs/api/visualizations.md Outdated

AntoninPoche changed the base branch from fixes to main September 18, 2026 11:53

AntoninPoche and others added 7 commits September 18, 2026 13:54

Merge branch 'main' of github.com:FOR-sight-ai/interpreto into lenses

543fb12

AllLayersSplitter: catch missing tokenizer

94c3377

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

feat(lens): add optional alignment for generated token predictions

ac11ed5

- keep next-token predictions as the default generation behavior - add align=True for target-aligned visualizations - improve whitespace rendering and update lens documentation and tests

feat(lens): improve token alignment and visualization

310f739

- align causal lens predictions by default - distinguish input, embeddings, and layer outputs - highlight correct predictions and document unaligned views

fix(lens): support light and dark visualization themes

cdee532

- adapt lens colors to the active documentation theme - improve prediction text contrast - refresh notebook outputs and add regression coverage

AntoninPoche reviewed Sep 29, 2026

View reviewed changes

Comment thread interpreto/lens/logit_lens.py Outdated

AntoninPoche reviewed Sep 29, 2026

View reviewed changes

Comment thread tests/lens/test_lens.py

fix(lens): use nnsight generation and expand model coverage

874d4ae

AntoninPoche approved these changes Oct 5, 2026

View reviewed changes

AntoninPoche added 3 commits October 6, 2026 10:31

test AllLayersSplitters: fix device

11ae09d

fixup! test AllLayersSplitters: fix device

5ac32e3

0.5.1 -> 0.5.2

3331d7b

AntoninPoche merged commit b3b9a0f into main Oct 6, 2026
5 checks passed

AntoninPoche deleted the lenses branch October 7, 2026 14:31

Navigation Menu

Repository navigation

Add Logit Lens and Tuned Lens methods - #161

Add Logit Lens and Tuned Lens methods#161
AntoninPoche merged 18 commits into
mainfrom
lenses

bernasraphael commented Aug 20, 2026 •

edited by AntoninPoche

Loading

Uh oh!

AntoninPoche left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

AntoninPoche left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

bernasraphael commented Aug 20, 2026 • edited by AntoninPoche Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Related Issue

Type of Change

Checklist

Uh oh!

AntoninPoche left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

AntoninPoche left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

bernasraphael commented Aug 20, 2026 •

edited by AntoninPoche

Loading