Logit Lens and Tuned Lens¶

Author: Raphaël Bernas

Lens methods show how a model's prediction develops across transformer blocks. Pass a model repository ID directly for the common case, or provide an AllLayersSplitter when the model needs custom configuration.

In [1]:

import warnings

warnings.filterwarnings("ignore", message="IProgress not found.*")

import warnings warnings.filterwarnings("ignore", message="IProgress not found.*")

In [2]:

from itertools import islice

from datasets import load_dataset
from transformers import AutoModelForSequenceClassification
from transformers.utils import logging as transformers_logging

from interpreto import AllLayersSplitter, LogitLens, TunedLens, plot_lens

transformers_logging.set_verbosity_error()
transformers_logging.disable_progress_bar()

from itertools import islice from datasets import load_dataset from transformers import AutoModelForSequenceClassification from transformers.utils import logging as transformers_logging from interpreto import AllLayersSplitter, LogitLens, TunedLens, plot_lens transformers_logging.set_verbosity_error() transformers_logging.disable_progress_bar()

Classification with Logit Lens¶

Logit Lens can expose how class rankings develop through a sequence-classification model. Its activation_names list describes the order of the returned results.

In [3]:

classification_splitter = AllLayersSplitter(
    "distilbert-base-uncased-finetuned-sst-2-english",
    automodel=AutoModelForSequenceClassification,
)
classification_text = "The explanations are clear and useful."

logit_lens = LogitLens(classification_splitter, top_k=2)
logit_results = logit_lens(classification_text)
list(logit_results)

classification_splitter = AllLayersSplitter( "distilbert-base-uncased-finetuned-sst-2-english", automodel=AutoModelForSequenceClassification, ) classification_text = "The explanations are clear and useful." logit_lens = LogitLens(classification_splitter, top_k=2) logit_results = logit_lens(classification_text) list(logit_results)

Warning: You are sending unauthenticated requests to the HF Hub. Please set a HF_TOKEN to enable higher rate limits and faster downloads.

Out[3]:

['model.distilbert.transformer.layer.0.input',
 'model.distilbert.transformer.layer.0',
 'model.distilbert.transformer.layer.1',
 'model.distilbert.transformer.layer.2',
 'model.distilbert.transformer.layer.3',
 'model.distilbert.transformer.layer.4',
 'model.distilbert.transformer.layer.5']

In [4]:

plot_lens(
    logit_results,
    classification_text,
    tokenizer=classification_splitter.tokenizer,
    label_names={0: "negative", 1: "positive"},
)

plot_lens( logit_results, classification_text, tokenizer=classification_splitter.tokenizer, label_names={0: "negative", 1: "positive"}, )

The explanations are clear and useful.

Relative confidencelowhigh

Layer

Prediction

5-out

positive

4-out

positive

3-out

negative

2-out

negative

1-out

negative

0-out

positive

Embeddings

positive

Generation with Tuned Lens¶

A Tuned Lens learns one residual affine translator for each non-final state. The translators are trained jointly against the language model's final prediction distribution. Each training text provides more target distributions for the translators, so a larger and more varied corpus generally produces a better lens. This example uses a modest sample to remain quick to run; real analyses should use a much larger representative corpus and separate evaluation text. Training texts are processed sequentially, while all depths from one text share a prediction-head call. Causal lens results are aligned with their observed target tokens by default. In the plots, Embeddings is the prediction head applied before the first transformer block, while Input contains the actual preceding tokens.

In [5]:

training_texts = [
    "Paris is the capital of France.",
    "Rome is the capital of Italy.",
    "Berlin is the capital of Germany.",
    "Madrid is the capital of Spain.",
    "Ottawa is the capital of Canada.",
    "Tokyo is the capital of Japan.",
    "Water freezes when its temperature reaches zero degrees Celsius.",
    "The Earth travels around the Sun once each year.",
    "Plants use sunlight to produce energy through photosynthesis.",
    "A triangle has three sides and three angles.",
    "Musicians combine rhythm, melody, and harmony to create songs.",
    "Computers execute instructions to process and store information.",
]

tuned_lens = TunedLens("distilgpt2", top_k=3)
losses = tuned_lens.fit(training_texts, epochs=1)
losses

training_texts = [ "Paris is the capital of France.", "Rome is the capital of Italy.", "Berlin is the capital of Germany.", "Madrid is the capital of Spain.", "Ottawa is the capital of Canada.", "Tokyo is the capital of Japan.", "Water freezes when its temperature reaches zero degrees Celsius.", "The Earth travels around the Sun once each year.", "Plants use sunlight to produce energy through photosynthesis.", "A triangle has three sides and three angles.", "Musicians combine rhythm, melody, and harmony to create songs.", "Computers execute instructions to process and store information.", ] tuned_lens = TunedLens("distilgpt2", top_k=3) losses = tuned_lens.fit(training_texts, epochs=1) losses

Out[5]:

[5.869888007640839]

Small corpus, one epoch¶

In [6]:

generation_text = "Although the committee initially rejected the proposal, new evidence persuaded several members to"
tuned_results = tuned_lens(generation_text)
plot_lens(tuned_results, generation_text, tokenizer=tuned_lens.splitter.tokenizer)

generation_text = "Although the committee initially rejected the proposal, new evidence persuaded several members to" tuned_results = tuned_lens(generation_text) plot_lens(tuned_results, generation_text, tokenizer=tuned_lens.splitter.tokenizer)

Relative confidencelowhighcorrect prediction

Layer

the

committee

initially

rejected

the

proposal

,

new

evidence

persuaded

several

members

to

5-out

The

first

has

considered

the

idea

,

but

documents

emerged

the

members

of

4-out

The

first

has

asked

the

idea

,

but

research

emerged

the

lawmakers

of

3-out

The

police

has

began

requests

idea

to

which

developments

suggests

him

members

of

2-out

The

data

committee

started

the

idea

for

which

bies

evidence

me

dozen

members

1-out

The

dust

committee

launched

the

new

for

which

additions

evidence

the

hundred

members

0-out

The

dust

committee

broadcast

the

first

to

and

additions

evidence

to

hundred

members

Embeddings

When

first

committee

initially

rejected

first

proposal

,

new

evidence

persuaded

several

members

Input

Although

the

committee

initially

rejected

the

proposal

,

new

evidence

persuaded

several

members

Greedy continuation¶

The preceding example applies the lens to fixed human-written text with teacher forcing. With the default alignment, each column is an observed target token and the bottom row shows the preceding input token. generate() instead lets the wrapped model choose the continuation before applying the same aligned view.

In [7]:

generation_prompt = "Although the committee initially rejected the proposal,"
generated_text, generated_results = tuned_lens.generate(generation_prompt, max_new_tokens=10)
print(generation_prompt + generated_text)
plot_lens(
    generated_results,
    generation_prompt + generated_text,
    tokenizer=tuned_lens.splitter.tokenizer,
)

generation_prompt = "Although the committee initially rejected the proposal," generated_text, generated_results = tuned_lens.generate(generation_prompt, max_new_tokens=10) print(generation_prompt + generated_text) plot_lens( generated_results, generation_prompt + generated_text, tokenizer=tuned_lens.splitter.tokenizer, )

Although the committee initially rejected the proposal, but the committee later voted to approve the proposal.

Relative confidencelowhighcorrect prediction

Layer

but

the

committee

later

voted

to

approve

the

proposal

.

5-out

but

the

committee

later

voted

to

approve

the

proposal

.

4-out

but

then

committee

eventually

decided

unanimously

approve

it

proposal

.

3-out

which

it

appeals

's

found

unanimously

postpone

it

proposal

.

2-out

which

it

latter

members

found

unanimously

vote

approval

new

.

1-out

which

it

only

members

later

unanimously

create

approval

new

to

0-out

and

mostly

most

committee

later

to

be

approve

main

to

Embeddings

,

but

first

committee

later

voted

to

approve

first

proposal

Input

,

but

the

committee

later

voted

to

approve

the

proposal

Next-token view with align=False¶

Without alignment, each column is the token currently seen by the model and its cells contain predictions made after that token. A correct prediction therefore matches the next column, while the final column predicts a token beyond the displayed continuation.

In [8]:

unaligned_text, unaligned_results = tuned_lens.generate(generation_prompt, max_new_tokens=10, align=False)
plot_lens(
    unaligned_results,
    generation_prompt + unaligned_text,
    tokenizer=tuned_lens.splitter.tokenizer,
    align=False,
)

unaligned_text, unaligned_results = tuned_lens.generate(generation_prompt, max_new_tokens=10, align=False) plot_lens( unaligned_results, generation_prompt + unaligned_text, tokenizer=tuned_lens.splitter.tokenizer, align=False, )

Relative confidencelowhighcorrect prediction

Layer

but

the

committee

later

voted

to

approve

the

proposal

.

5-out

the

committee

later

voted

to

approve

the

proposal

.

\n

4-out

then

committee

eventually

decided

unanimously

approve

it

proposal

.

\n

3-out

it

appeals

's

found

unanimously

postpone

it

proposal

.

\n

2-out

it

latter

members

found

unanimously

vote

approval

new

.

It

1-out

it

only

members

later

unanimously

create

approval

new

to

It

0-out

mostly

most

committee

later

to

be

approve

main

to

\n

Embeddings

but

first

committee

later

voted

to

approve

first

proposal

When

Input

but

the

committee

later

voted

to

approve

the

proposal

.

Dataset size versus training epochs¶

More epochs and more data improve a lens in different ways. Repeating the same texts gives the optimizer more opportunities to fit those examples, but it can also overfit their limited distribution. More unique texts usually provide broader coverage.

The comparison below holds the number of text updates fixed: the small 12-sentence corpus is repeated for four epochs, while 48 distinct examples from the WikiText-2 dataset are seen once. WikiText paragraphs are truncated to keep this documentation example lightweight. The number of tokens still differs and is reported below. Training losses from different corpora are not directly comparable, so all lenses are also inspected on the same held-out prompt used above.

In [9]:

comparison_epochs = 4
larger_dataset_size = len(training_texts) * comparison_epochs
wikitext = load_dataset(
    "Salesforce/wikitext",
    "wikitext-2-raw-v1",
    split="train",
    streaming=True,
)
larger_training_texts = list(
    islice(
        (row["text"].strip()[:256] for row in wikitext if len(row["text"].split()) >= 8),
        larger_dataset_size,
    )
)
len(larger_training_texts)

comparison_epochs = 4 larger_dataset_size = len(training_texts) * comparison_epochs wikitext = load_dataset( "Salesforce/wikitext", "wikitext-2-raw-v1", split="train", streaming=True, ) larger_training_texts = list( islice( (row["text"].strip()[:256] for row in wikitext if len(row["text"].split()) >= 8), larger_dataset_size, ) ) len(larger_training_texts)

Out[9]:

48

In [10]:

more_epochs_lens = TunedLens(tuned_lens.splitter, top_k=3)
larger_dataset_lens = TunedLens(tuned_lens.splitter, top_k=3)

more_epochs_losses = more_epochs_lens.fit(training_texts, epochs=comparison_epochs)
larger_dataset_losses = larger_dataset_lens.fit(larger_training_texts, epochs=1)

more_epochs_lens = TunedLens(tuned_lens.splitter, top_k=3) larger_dataset_lens = TunedLens(tuned_lens.splitter, top_k=3) more_epochs_losses = more_epochs_lens.fit(training_texts, epochs=comparison_epochs) larger_dataset_losses = larger_dataset_lens.fit(larger_training_texts, epochs=1)

In [11]:

tokenizer = tuned_lens.splitter.tokenizer
small_dataset_tokens = sum(len(tokenizer.encode(text)) for text in training_texts)
larger_dataset_tokens = sum(len(tokenizer.encode(text)) for text in larger_training_texts)

{
    "small corpus, 1 epoch": {
        "unique texts": len(training_texts),
        "text updates": len(training_texts),
        "tokens seen": small_dataset_tokens,
        "final training loss": losses[-1],
    },
    "small corpus, 4 epochs": {
        "unique texts": len(training_texts),
        "text updates": len(training_texts) * comparison_epochs,
        "tokens seen": small_dataset_tokens * comparison_epochs,
        "final training loss": more_epochs_losses[-1],
    },
    "larger corpus, 1 epoch": {
        "unique texts": len(larger_training_texts),
        "text updates": len(larger_training_texts),
        "tokens seen": larger_dataset_tokens,
        "final training loss": larger_dataset_losses[-1],
    },
}

tokenizer = tuned_lens.splitter.tokenizer small_dataset_tokens = sum(len(tokenizer.encode(text)) for text in training_texts) larger_dataset_tokens = sum(len(tokenizer.encode(text)) for text in larger_training_texts) { "small corpus, 1 epoch": { "unique texts": len(training_texts), "text updates": len(training_texts), "tokens seen": small_dataset_tokens, "final training loss": losses[-1], }, "small corpus, 4 epochs": { "unique texts": len(training_texts), "text updates": len(training_texts) * comparison_epochs, "tokens seen": small_dataset_tokens * comparison_epochs, "final training loss": more_epochs_losses[-1], }, "larger corpus, 1 epoch": { "unique texts": len(larger_training_texts), "text updates": len(larger_training_texts), "tokens seen": larger_dataset_tokens, "final training loss": larger_dataset_losses[-1], }, }

Out[11]:

{'small corpus, 1 epoch': {'unique texts': 12,
  'text updates': 12,
  'tokens seen': 110,
  'final training loss': 5.869888007640839},
 'small corpus, 4 epochs': {'unique texts': 12,
  'text updates': 48,
  'tokens seen': 440,
  'final training loss': 1.474266807238261},
 'larger corpus, 1 epoch': {'unique texts': 48,
  'text updates': 48,
  'tokens seen': 2207,
  'final training loss': 4.1064460799098015}}

Small corpus, four epochs¶

In [12]:

more_epochs_results = more_epochs_lens(generation_text)
plot_lens(more_epochs_results, generation_text, tokenizer=tokenizer)

more_epochs_results = more_epochs_lens(generation_text) plot_lens(more_epochs_results, generation_text, tokenizer=tokenizer)

Relative confidencelowhighcorrect prediction

Layer

the

committee

initially

rejected

the

proposal

,

new

evidence

persuaded

several

members

to

5-out

The

first

has

considered

the

idea

,

but

documents

emerged

the

members

of

4-out

The

US

chair

announced

the

proposal

because

but

research

suggests

the

lawmakers

to

3-out

The

first

has

announced

the

idea

to

which

arrivals

suggests

the

members

of

2-out

The

extent

members

announced

the

idea

for

a

entrants

that

him

members

members

1-out

The

focus

members

broadcast

the

state

for

a

additions

that

the

elements

members

0-out

The

Sun

committee

broadcast

the

Sun

proposal

and

new

that

enough

parts

members

Embeddings

Although

first

committee

initially

rejected

world

proposal

,

new

evidence

persuaded

several

members

Input

Although

the

committee

initially

rejected

the

proposal

,

new

evidence

persuaded

several

members

Larger corpus, one epoch¶

In [13]:

larger_dataset_results = larger_dataset_lens(generation_text)
plot_lens(larger_dataset_results, generation_text, tokenizer=tokenizer)

larger_dataset_results = larger_dataset_lens(generation_text) plot_lens(larger_dataset_results, generation_text, tokenizer=tokenizer)

Relative confidencelowhighcorrect prediction

Layer

the

committee

initially

rejected

the

proposal

,

new

evidence

persuaded

several

members

to

5-out

The

first

has

considered

the

idea

,

but

documents

emerged

the

members

of

4-out

The

first

will

recommended

the

idea

.

the

research

shows

the

lawmakers

of

3-out

The

majority

was

had

the

proposal

for

which

research

suggests

the

states

of

2-out

The

majority

was

had

the

idea

for

which

additions

that

the

people

of

1-out

The

majority

committee

posted

the

state

for

which

additions

that

it

weapons

of

0-out

The

state

committee

ordered

the

state

of

and

weapons

of

the

hundred

of

Embeddings

You

identity

committee

to

to

identity

,

,

new

,

to

few

members

Input

Although

the

committee

initially

rejected

the

proposal

,

new

evidence

persuaded

several

members

Neither lever is universally better: enough optimization is needed to fit the translators, while enough diverse data is needed for them to generalize. Choose the corpus size and epoch count together, monitor a held-out metric, and stop increasing epochs when held-out performance no longer improves.

TunedLens is a regular PyTorch module. Save and restore its translators with state_dict() and load_state_dict().