# What Is the Purpose of the `forward_hidden` Method in the Transformer Model?

> Discover the purpose of the forward_hidden method in Transformer models. Learn how it efficiently computes intermediate token representations for various downstream tasks.

- Repository: [Fareed Khan/train-llm-from-scratch](https://github.com/FareedKhan-dev/train-llm-from-scratch)
- Tags: internals
- Published: 2026-06-11

---

**The `forward_hidden` method computes intermediate token-level representations by running the Transformer’s embedding layers, positional encodings, and attention blocks while skipping the final language model head, enabling efficient reuse of the backbone for multiple downstream tasks such as value estimation and policy generation.**

In the `FareedKhan-dev/train-llm-from-scratch` repository, the `forward_hidden` method serves as a critical abstraction within the Transformer architecture. Located in [`src/models/transformer.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/models/transformer.py), this method isolates the representation-learning portion of the model to support advanced training paradigms like actor-critic reinforcement learning while avoiding redundant computation across different prediction heads.

## Understanding the `forward_hidden` Method

### Core Responsibilities

According to the source implementation in [`src/models/transformer.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/models/transformer.py), the method fulfills three essential functions:

1. **Compute token-level hidden representations** – It processes input token IDs through the embedding layer, positional encoding, and all self-attention and feed-forward blocks, returning a tensor of shape `(batch, seq_len, hidden_dim)` containing the final hidden states for each token.

2. **Enable reuse of the shared backbone** – By exposing just the hidden states, the same Transformer backbone can power multiple downstream heads. In the PPO-style actor-critic implementation found in [`src/post_training/value_head.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/value_head.py), these hidden states feed both the language-model head (producing policy logits) and a dedicated value-head that predicts per-token state-values, eliminating the need to duplicate expensive Transformer computations.

3. **Provide a stable interface for evaluation** – The method supplies hidden representations to reward models and other utilities that require encoded sequence information without the computational cost of projecting to the full vocabulary, which is particularly useful when only scalar values or similarity scores are needed.

### Implementation Location

The method is defined within the `Transformer` class in [`src/models/transformer.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/models/transformer.py). Unlike the standard `forward` method, which internally calls `forward_hidden` and then applies `self.lm_head` to produce vocabulary logits, `forward_hidden` terminates immediately after the final Transformer block, returning the raw hidden states.

## Actor-Critic Implementation and Multi-Head Reuse

The primary use case for `forward_hidden` appears in reinforcement learning workflows where a single model must produce both policy distributions and value estimates. In [`src/post_training/value_head.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/value_head.py), the `TransformerWithValueHead` wrapper leverages this method to share the backbone between the actor (policy) and critic (value) networks.

```python
import torch
from src.models.transformer import Transformer
from src.post_training.value_head import TransformerWithValueHead

# Build a plain Transformer backbone

model = Transformer(
    vocab_size=50257,
    d_model=768,
    n_head=12,
    n_layer=12,
    max_seq_len=1024,
)

# Example input: batch of token IDs (batch_size, seq_len)

tokens = torch.randint(0, model.vocab_size, (2, 128))

# Get hidden states only: shape (2, 128, 768)

hidden = model.forward_hidden(tokens)

# Standard forward pass produces logits: shape (2, 128, vocab_size)

logits = model(tokens)

# Actor-critic wrapper reuses hidden states for both heads

actor_critic = TransformerWithValueHead(model)
policy_logits, values = actor_critic(tokens)  # values: (2, 128)

```

In this architecture, `TransformerWithValueHead` calls `forward_hidden` once and feeds the resulting tensor to both the `lm_head` (for policy logits) and a small value head MLP (for state-value estimation). This design pattern prevents the memory overhead and latency that would occur if the Transformer layers were computed separately for each output type.

## Performance Benefits and Memory Efficiency

Separating hidden state computation from vocabulary projection yields significant performance advantages during training and inference. When only value estimates are required—such as during rollout phases in PPO—the `value_only` convenience method calls `forward_hidden` and skips the costly projection to vocabulary space entirely.

This approach reduces memory usage by avoiding the storage of large `(batch, seq_len, vocab_size)` logits tensors when they are unnecessary, and it eliminates redundant forward passes through the Transformer blocks that would otherwise occur if the policy and value networks were maintained as separate models.

## Summary

- **`forward_hidden`** isolates the Transformer’s representation-learning layers from the output head, returning hidden states of shape `(batch, seq_len, hidden_dim)`.

- The method enables **multi-head architectures** by allowing shared backbones to feed both language-model heads and auxiliary heads like value estimators, as implemented in [`src/post_training/value_head.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/value_head.py).

- By exposing intermediate representations, the method supports **efficient reinforcement learning** workflows where both policy logits and value estimates are computed from a single forward pass through the Transformer blocks.

- The implementation avoids **redundant computation** and reduces memory overhead compared to running separate Transformer instances for different prediction tasks.

## Frequently Asked Questions

### What is the difference between `forward_hidden` and the standard `forward` method?

The `forward` method in [`src/models/transformer.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/models/transformer.py) internally calls `forward_hidden` to obtain the hidden states and then applies `self.lm_head` to project those states into vocabulary logits. In contrast, `forward_hidden` stops after the final Transformer block, returning the raw hidden representations without the vocabulary projection, making it suitable for reuse across multiple downstream heads.

### Why is `forward_hidden` necessary for reinforcement learning?

Reinforcement learning algorithms like PPO require both policy distributions (next-token probabilities) and value estimates (state-values) from the same input sequence. The `forward_hidden` method allows the implementation in [`src/post_training/value_head.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/src/post_training/value_head.py) to compute the Transformer’s hidden states once and feed them to both the language-model head and a value head, avoiding the computational cost of running the full Transformer stack twice.

### What shape does the `forward_hidden` output tensor have?

The method returns a tensor of shape `(batch_size, sequence_length, hidden_dimension)`, where `hidden_dimension` corresponds to the `d_model` parameter (typically 768 in the examples). This tensor contains the final hidden state for each token position after processing through all attention and feed-forward layers.

### Where is `forward_hidden` tested in the codebase?

The method is validated in [`tests/test_post_training_smoke.py`](https://github.com/FareedKhan-dev/train-llm-from-scratch/blob/main/tests/test_post_training_smoke.py), which includes unit tests verifying that `forward_hidden` produces the correct intermediate representations and that these representations, when passed through the language model head, reproduce the logits of the standard full forward pass.