Add Logit Lens and Tuned Lens methods

README.md

@@ -27,10 +27,11 @@ Checkout the tutorials to get started:
 - [Attributions walkthrough](https://github.com/FOR-sight-ai/interpreto/tree/main/docs/notebooks/attribution_tutorial.ipynb) (both classification and generation)
 - [Classification concept-based explanations](https://github.com/FOR-sight-ai/interpreto/tree/main/docs/notebooks/classification_concept_tutorial.ipynb)
 - [Generation concept-based explanations](https://github.com/FOR-sight-ai/interpreto/tree/main/docs/notebooks/generation_concept_tutorial.ipynb)
+- [Logit Lens and Tuned Lens](https://github.com/FOR-sight-ai/interpreto/tree/main/docs/notebooks/lens_notebook.ipynb)
 
 ## 📦 What's Included
 
-Interpreto 🪄 provides a modular framework encompassing Attribution Methods, Concept-Based Methods, and Evaluation Metrics.
+Interpreto 🪄 provides a modular framework encompassing Attribution Methods, Lens Methods, Concept-Based Methods, and Evaluation Metrics.
 
 ### 🔥 Attribution Methods
 
@@ -55,6 +56,13 @@ They all work seamlessly for both classification (`...ForSequenceClassification`
 - [`SquareGrad`](https://for-sight-ai.github.io/interpreto/api/attributions/methods/squaregrad/) — [Hooker et al., 2019](https://arxiv.org/abs/1806.10758)
 - [`VarGrad`](https://for-sight-ai.github.io/interpreto/api/attributions/methods/vargrad/) — [Richter et al., 2020](https://proceedings.neurips.cc/paper/2020/hash/9c22c0b51b3202246463e986c7e205df-Abstract.html)
 
+### 🔍 Lens Methods
+
+- [`LogitLens`](https://for-sight-ai.github.io/interpreto/api/lens/) — [nostalgebraist, 2020](https://www.lesswrong.com/posts/AcKRB8wDpdaN6v6ru/interpreting-gpt-the-logit-lens)
+- [`TunedLens`](https://for-sight-ai.github.io/interpreto/api/lens/) — [Belrose et al., 2023](https://arxiv.org/abs/2303.08112)
+
+Both methods accept a model repository ID directly or a configured `AllLayersSplitter` for advanced use cases.
+
 ### 💡 Concept-Based Methods or Mechanistic Interpretability
 
 Concept-based explanations aim to provide high-level interpretations of latent model representations.

docs/api/concepts/splitters/all_layers_splitter.md

@@ -0,0 +1,10 @@
+# All Layers Splitter
+
+`AllLayersSplitter` captures the residual stream before the first transformer block and after every block. It also
+applies the wrapped model's native normalization, pooling, and prediction head to one or more residual states.
+
+::: interpreto.AllLayersSplitter
+    handler: python
+    options:
+      show_root_heading: true
+      show_source: true

docs/api/lens.md

@@ -0,0 +1,91 @@
+# Lens
+
+Lens methods decode the residual stream at every transformer block boundary. Pass a model repository ID directly for
+the common case, or provide a configured [`AllLayersSplitter`](concepts/splitters/all_layers_splitter.md) when you need
+to select a model class, tokenizer, device, or transformer block path.
+
+Inference accepts one text at a time. To explain several texts, iterate over them. `TunedLens.fit()` accepts either one
+text or an iterable and trains on each text sequentially.
+
+## Classification with Logit Lens
+
+```python
+from transformers import AutoModelForSequenceClassification
+
+from interpreto import AllLayersSplitter, LogitLens, plot_lens
+
+text = "Interpreto makes model decisions easier to inspect."
+splitter = AllLayersSplitter(
+    "distilbert-base-uncased-finetuned-sst-2-english",
+    automodel=AutoModelForSequenceClassification,
+)
+lens = LogitLens(splitter, top_k=2)
+results = lens(text)
+
+plot_lens(results, text, tokenizer=splitter.tokenizer, label_names=["negative", "positive"])
+```
+
+## Generation with Tuned Lens
+
+```python
+from interpreto import TunedLens, plot_lens
+
+text = "Paris is the capital of"
+lens = TunedLens("distilgpt2", top_k=3)
+lens.fit(["Paris is the capital of France.", "Rome is the capital of Italy."], epochs=1)
+results = lens(text)
+
+plot_lens(results, text, tokenizer=lens.splitter.tokenizer)
+```
+
+For causal language models, predictions are aligned with their observed target tokens by default. Each column therefore
+shows a target token, while the bottom `Input` row shows the preceding token that produced its prediction. Correct top
+predictions are outlined in green.
+
+`generate()` uses the same convention. Include the prompt when plotting so the visualization can display the input to
+the first generated prediction:
+
+```python
+prompt = "Although the committee initially rejected the proposal,"
+generated_text, generated_results = lens.generate(prompt, max_new_tokens=10)
+
+plot_lens(generated_results, prompt + generated_text, tokenizer=lens.splitter.tokenizer)
+```
+
+Set `align=False` on both calls for the conventional next-token view, where each column contains the prediction made
+after its displayed input token:
+
+```python
+generated_text, generated_results = lens.generate(prompt, max_new_tokens=10, align=False)
+plot_lens(
+    generated_results,
+    prompt + generated_text,
+    tokenizer=lens.splitter.tokenizer,
+    align=False,
+)
+```
+
+## `LogitLens`
+
+::: interpreto.LogitLens
+    handler: python
+    options:
+      show_root_heading: true
+      show_source: true
+
+## `TunedLens`
+
+::: interpreto.TunedLens
+    handler: python
+    options:
+      show_root_heading: true
+      show_source: true
+
+## `plot_lens`
+
+`plot_lens` renders model outputs from the final layer down to the embeddings, followed by the actual input tokens. The
+`Embeddings` row is the prediction head applied before the first transformer block; it is not the raw model input. The
+visible cells contain only the top prediction and use color intensity to show relative confidence. A green outline marks
+a correct prediction. Hover over a cell to inspect its numerical score and remaining top-k results.
+
+::: interpreto.visualizations.plot_lens

docs/index.md

@@ -31,10 +31,11 @@ Checkout the tutorials to get started:
 - [Attributions walkthrough](https://github.com/FOR-sight-ai/interpreto/tree/main/docs/notebooks/attribution_tutorial.ipynb) (both classification and generation)
 - [Classification concept-based explanations](https://github.com/FOR-sight-ai/interpreto/tree/main/docs/notebooks/classification_concept_tutorial.ipynb)
 - [Generation concept-based explanations](https://github.com/FOR-sight-ai/interpreto/tree/main/docs/notebooks/generation_concept_tutorial.ipynb)
+- [Logit Lens and Tuned Lens](https://github.com/FOR-sight-ai/interpreto/tree/main/docs/notebooks/lens_notebook.ipynb)
 
 ## 📦 What's Included
 
-Interpreto 🪄 provides a modular framework encompassing Attribution Methods, Concept-Based Methods, and Evaluation Metrics.
+Interpreto 🪄 provides a modular framework encompassing Attribution Methods, Lens Methods, Concept-Based Methods, and Evaluation Metrics.
 
 ### 🔥 Attribution Methods
 
@@ -59,6 +60,13 @@ They all work seamlessly for both classification (`...ForSequenceClassification`
 - [`SquareGrad`](https://for-sight-ai.github.io/interpreto/api/attributions/methods/squaregrad/) — [Hooker et al., 2019](https://arxiv.org/abs/1806.10758)
 - [`VarGrad`](https://for-sight-ai.github.io/interpreto/api/attributions/methods/vargrad/) — [Richter et al., 2020](https://proceedings.neurips.cc/paper/2020/hash/9c22c0b51b3202246463e986c7e205df-Abstract.html)
 
+### 🔍 Lens Methods
+
+- [`LogitLens`](api/lens.md) — [nostalgebraist, 2020](https://www.lesswrong.com/posts/AcKRB8wDpdaN6v6ru/interpreting-gpt-the-logit-lens)
+- [`TunedLens`](api/lens.md) — [Belrose et al., 2023](https://arxiv.org/abs/2303.08112)
+
+Both methods accept a model repository ID directly or a configured `AllLayersSplitter` for advanced use cases.
+
 ### 💡 Concept-Based Methods or Mechanistic Interpretability
 
 Concept-based explanations aim to provide high-level interpretations of latent model representations.