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.
Add batched fitting, evaluation metrics, and versioned tensor-only Tuned Lens checkpoints.
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.

AntoninPoche · 2026-09-15T15:46:46Z

+`LogitLens` follows the method introduced by
+[nostalgebraist](https://www.lesswrong.com/posts/AcKRB8wDpdaN6v6ru/interpreting-gpt-the-logit-lens). It projects the residual
+stream at every model depth through the model's native prediction path. The method has no learned parameters.
+
+The prediction head was trained on final states, so early-layer scores should be interpreted as rankings rather than
+calibrated probabilities. Tuned Lens addresses part of this distribution mismatch by learning a translator for every
+non-final state.


This could be in the class docstring directly.

AntoninPoche · 2026-09-15T15:47:25Z

+`TunedLens` follows [Belrose et al. (2023)](https://arxiv.org/abs/2303.08112). It learns one affine residual
+translator for each non-final model state and trains all of them to match the model's final prediction distribution.
+
+The translators start at zero, making an unfitted Tuned Lens exactly equivalent to a Logit Lens. Fitting processes one
+text at a time while projecting all model depths together. Use separate training and evaluation texts when assessing
+the fitted lens.


I think this should be in the docstring directly. (as the docstring is shown in the doc).

AntoninPoche · 2026-09-15T15:49:00Z

+`results` is ordered like `splitter.activation_names`. Each entry contains `top_indices` and `top_scores`. Language
+model outputs have shape `(1, sequence_length, k)`; sequence-classification outputs have shape `(1, k)`, where `k` is
+`top_k` capped at the model's output size. The singleton dimension represents the one input text. Input batches and
+pre-tokenized inputs are intentionally not part of this initial API.


Does it support a list of strings by computing them iteratively, or is it the user's role to do so?

AntoninPoche · 2026-09-15T15:55:03Z

+ones, do not follow the final prediction head's training distribution, so these values are useful for rankings and
+within-model comparisons rather than as calibrated probabilities.
+
+## Tuned Lens


If you describe both methods here, you could put their api description from mkdocs at the end or in the middle.

Therefore, making it a single file for all the lens modules and simplifying the doc. Your call.

AntoninPoche · 2026-09-15T15:56:16Z

+## Visualization
+
+```python
+from interpreto import plot_lens
+
+plot_lens(results, "Interpreto is useful.", tokenizer=splitter.tokenizer)
+```
+
+For sequence classification, pass `label_names` to display readable class names.


You should add the plot line directly using the code snippet above.

AntoninPoche · 2026-09-15T16:13:07Z

+.lens-score { color: #526172; }
+"""
+
+Tokenizer = PreTrainedTokenizer | PreTrainedTokenizerFast


You can use the class PreTrainedTokenizerBase, which is the parent of both.

AntoninPoche · 2026-09-15T16:18:15Z

+    - Lens:
+      - Overview: api/lens/overview.md
+      - Methods:
+        - Logit Lens: api/lens/methods/logit_lens.md
+        - Tuned Lens: api/lens/methods/tuned_lens.md


I suggested above to merge everything in a single file, no folders. You tell me.

AntoninPoche · 2026-09-15T16:19:19Z

When you modify the readme.md, do not forget the index.md

AntoninPoche · 2026-09-15T16:29:55Z

I will be better able to comment after the notebook is compiled.

AntoninPoche · 2026-09-15T16:31:34Z

The notebook should be compiled for people to check its results.

You can include your name as the author.

You can add classification.

Do not hesitate to use real models, even if small.

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.

+        layer_name, layers = max(module_lists, key=lambda item: len(item[1]))
+        self.split_points = [f"model.{layer_name}.{index}" for index in range(len(layers))]


+        if self.tokenizer.pad_token is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token


+    def fit(
+        self,
+        inputs: str | Iterable[str],
+        epochs: int = 1,
+        learning_rate: float = 1e-3,
+        weight_decay: float = 0.0,
+    ) -> list[float]:


+def test_lens_notebook_is_portable_and_has_no_error_outputs():
+    notebook = json.loads((REPOSITORY_ROOT / "docs" / "notebooks" / "lens_notebook.ipynb").read_text())
+
+    assert notebook["metadata"]["kernelspec"]["name"] == "python3"
+    assert all(
+        output.get("output_type") != "error" for cell in notebook["cells"] for output in cell.get("outputs", [])
+    )


+the explanation and the splitter tokenizer. For sequence classification, `label_names` can map label ids to readable
+names. Like the other Interpreto plotting functions, it accepts `custom_css` and an optional `save_path`.
+
+The displayed values are intermediate softmax scores, not calibrated probabilities. Language-model tooltips show the numerical top-k scores at each position. For classification, bar lengths show the corresponding class scores.


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

Open

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

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

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

543fb12

Navigation Menu

Add Logit Lens and Tuned Lens methods - #161

Add Logit Lens and Tuned Lens methods#161
bernasraphael wants to merge 8 commits into
mainfrom
lenses

bernasraphael commented Aug 20, 2026 •

edited by AntoninPoche

Loading

Uh oh!

AntoninPoche left a comment

Uh oh!

AntoninPoche Sep 15, 2026

Uh oh!

AntoninPoche Sep 15, 2026

Uh oh!

AntoninPoche Sep 15, 2026

Uh oh!

AntoninPoche Sep 15, 2026

Uh oh!

AntoninPoche Sep 15, 2026

Uh oh!

AntoninPoche Sep 15, 2026

Uh oh!

AntoninPoche Sep 15, 2026

Uh oh!

AntoninPoche Sep 15, 2026

Uh oh!

AntoninPoche Sep 15, 2026

Uh oh!

AntoninPoche Sep 15, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

		layer_name, layers = max(module_lists, key=lambda item: len(item[1]))
		self.split_points = [f"model.{layer_name}.{index}" for index in range(len(layers))]

		if self.tokenizer.pad_token is None:
		self.tokenizer.pad_token = self.tokenizer.eos_token

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!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

bernasraphael commented Aug 20, 2026 •

edited by AntoninPoche

Loading