Last layer hidden state: GPT2

🤗Transformers
1

Hi, I am trying to understand exactly where the last last hidden output in GPT2 stems from. The following code outputs hidden states from the embedding layer and all layers (1+12=13 layers):

tokenizer = GPT2Tokenizer.from_pretrained(‘gpt2’)
model = GPT2Model.from_pretrained(‘gpt2’)

input_ids = torch.tensor(tokenizer.encode(sent))
outputs = model(input_ids, output_hidden_states=True, return_dict=True)
hidden_states = result_model[‘hidden_states’]

As GPT2 has a LayerNorm layer after the very last decoder block (ln_f: LayerNorm(768)), are the last hidden states extracted before this final normalization layer (after the last decoder block) or after this final normalization layer?

I am running Transformers 3.1.0.

Thanks so much!

2

From what I read, it seemed like GPT-2’s last hidden state would not have gone through a layer norm and that I would have to hook ln_f to have a comparable highest layer. However, when I did that, I found that every value from ln_f was within precision limitations of hidden state 12 (over 1290x16x768). Can anyone give a definitive answer to what is being reported as hidden state 12?

3

I took a quick look:

Yes — in the PyTorch GPT-2 implementation you referenced, hidden_states[12] is the post-ln_f value.

More precisely, the final part of the data flow is:

raw output of transformer block 11
        ↓
      ln_f
        ↓
hidden_states[12]
last_hidden_state

A normal PyTorch forward hook on ln_f receives the output of that module, so its values are expected to match hidden_states[12]. The agreement across the full 1290 × 16 × 768 tensor is therefore not a numerical coincidence; both observations refer to the same point in the forward data flow.

The indexing is slightly easy to misread because the tuple has 13 entries for a 12-block model:

Returned entry What it represents
hidden_states[0] Embedding sum after embedding dropout; input to block 0
hidden_states[1] Output of block 0; input to block 1
... ...
hidden_states[11] Output of block 10; input to block 11
not retained separately Raw output of block 11; input to ln_f
hidden_states[12] ln_f(raw output of block 11)

So an important distinction is:

hidden_states[11] != raw output of the final block

hidden_states[11] is the input to the final block. The raw output of the final block exists immediately before ln_f, but it is not stored as a separate entry in the returned hidden_states tuple.

Which value should you use?

  • If you want the normal final representation returned by Hugging Face GPT-2, use:

    outputs.hidden_states[-1]

    or, with GPT2Model:

    outputs.last_hidden_state

    These are post-ln_f.

  • If you specifically want the raw output of the last transformer block, before the final LayerNorm, capture the input to ln_f with a forward pre-hook:

    captured = {}

def save_pre_ln_f(module, args):
    captured["pre_ln_f"] = args[0].detach().clone()

handle = model.ln_f.register_forward_pre_hook(save_pre_ln_f)

try:
    model.eval()
    with torch.no_grad():
        outputs = model(
            **inputs,
            output_hidden_states=True,
            return_dict=True,
        )
finally:
    handle.remove()

raw_final_block_output = captured["pre_ln_f"]
final_model_hidden_state = outputs.hidden_states[-1]

If the object is a GPT2LMHeadModel rather than a bare GPT2Model, the module is normally under:

model.transformer.ln_f

PyTorch’s hook documentation distinguishes these two points directly: a forward pre-hook runs before the module’s forward(), while a normal forward hook runs after its output has been computed.

Why the v3.1.0 indexing works this way

The relevant part of the Transformers v3.1.0 GPT-2 implementation can be reduced to this sequence:

hidden_states = inputs_embeds + position_embeds + token_type_embeds
hidden_states = embedding_dropout(hidden_states)

for block in transformer_blocks:
    if output_hidden_states:
        all_hidden_states += (hidden_states,)

    hidden_states = block(hidden_states)

hidden_states = ln_f(hidden_states)

if output_hidden_states:
    all_hidden_states += (hidden_states,)

The timing of the append operation is the key:

  1. Before each block runs, its input is appended.
  2. There are 12 blocks, so that produces 12 entries.
  3. After all blocks run, ln_f is applied.
  4. That post-ln_f value is appended as the 13th entry.

For a 12-block GPT-2:

entry 0  = input to block 0
entry 1  = input to block 1 = output of block 0
...
entry 11 = input to block 11 = output of block 10
entry 12 = output of ln_f

This also explains why the general documentation phrase “embedding output plus the output of each layer” can be a little ambiguous here. At the final boundary, GPT-2 has a model-level normalization after the transformer block stack, so it is useful to distinguish:

raw final-block output

from:

final model hidden state

The returned final entry is the latter.

Small hook-based verification

A small 12-block GPT-2 sanity check reproduced the following relationships:

number of transformer blocks: 12
number of returned hidden states: 13

final block input
== hidden_states[-2]
maximum absolute difference: 0.0

final block output
== ln_f input
maximum absolute difference: 0.0

ln_f output
== hidden_states[-1]
maximum absolute difference: 0.0

ln_f output
== last_hidden_state
maximum absolute difference: 0.0

The pre- and post-ln_f values themselves were not equal:

ln_f input vs ln_f output
maximum absolute difference: 2.880181312561035

That is useful because it separates two possibilities:

  • ln_f happened to behave approximately like an identity operation; or
  • the hook result and hidden_states[-1] represent the same post-ln_f boundary.

The observation supports the second explanation, which is also what the source code says.

Here is a fuller check that captures all relevant boundaries in one forward pass:

import torch

# GPT2Model -> core is the model itself
# GPT2LMHeadModel -> core is model.transformer
core = model.transformer if hasattr(model, "transformer") else model

captured = {}

def save_final_block_input(module, args):
    captured["final_block_input"] = args[0].detach().clone()

def save_final_block_output(module, args, output):
    # Older Transformers versions may return a tuple/list from a block,
    # while newer implementations may return the tensor directly.
    tensor = output[0] if isinstance(output, (tuple, list)) else output
    captured["final_block_output"] = tensor.detach().clone()

def save_pre_ln_f(module, args):
    captured["pre_ln_f"] = args[0].detach().clone()

def save_post_ln_f(module, args, output):
    captured["post_ln_f"] = output.detach().clone()

handles = [
    core.h[-1].register_forward_pre_hook(save_final_block_input),
    core.h[-1].register_forward_hook(save_final_block_output),
    core.ln_f.register_forward_pre_hook(save_pre_ln_f),
    core.ln_f.register_forward_hook(save_post_ln_f),
]

try:
    model.eval()

    with torch.no_grad():
        outputs = model(
            **inputs,
            output_hidden_states=True,
            return_dict=True,
        )
finally:
    for handle in handles:
        handle.remove()

assert len(outputs.hidden_states) == len(core.h) + 1

torch.testing.assert_close(
    captured["final_block_input"],
    outputs.hidden_states[-2],
)

torch.testing.assert_close(
    captured["final_block_output"],
    captured["pre_ln_f"],
)

torch.testing.assert_close(
    captured["post_ln_f"],
    outputs.hidden_states[-1],
)

# Present on GPT2Model outputs. A GPT2LMHeadModel output does not normally
# expose last_hidden_state directly at its top level.
if hasattr(outputs, "last_hidden_state"):
    torch.testing.assert_close(
        captured["post_ln_f"],
        outputs.last_hidden_state,
    )

print(
    "pre/post ln_f maximum absolute difference:",
    (
        captured["pre_ln_f"]
        - captured["post_ln_f"]
    ).abs().max().item(),
)

Keeping all comparisons inside the same forward pass avoids dropout or input differences between separate calls. model.eval() also disables the model’s dropout layers.

The ln_f hook output is the output of the complete LayerNorm module, including its learned affine scale and bias when enabled—not merely an intermediate mean/variance-normalized tensor. See the PyTorch LayerNorm documentation.

Version and backend notes

There is also a useful historical confirmation in Transformers issue #13102.

That issue explicitly described the backend behavior at the time:

  • PyTorch GPT-2 added the last hidden state after ln_f.
  • TensorFlow GPT-2 did the same.
  • Flax GPT-2 was then adding it before ln_f.

The Flax difference was treated as an inconsistency and was subsequently aligned through PR #13109.

That is independent confirmation that the post-ln_f interpretation was not merely an accidental consequence of this particular comparison.

The internal machinery used to collect hidden states has changed in newer Transformers releases, so the old tuple-building code should not be assumed to be the present implementation. However, the tagged Transformers 5.14.1 GPT-2 source still performs the same final model-level sequence:

for block in self.h:
    hidden_states = block(hidden_states, ...)

hidden_states = self.ln_f(hidden_states)

return BaseModelOutputWithPastAndCrossAttentions(
    last_hidden_state=hidden_states,
    ...
)

A small check with that version also produced:

hidden_states[-1] == last_hidden_state == post-ln_f

The version-matched v3.1.0 source remains the primary evidence for the code in the question; the newer check just shows that the same GPT-2 boundary is still represented that way.

One final scope note: this describes Hugging Face’s PyTorch GPT-2 forward path. It should not be generalized automatically to every Transformers architecture. The generic model-output documentation specifically warns that hidden_states[-1] and last_hidden_state are not guaranteed to match for every model, because some architectures apply additional normalization or processing at different points.