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

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, 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, 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, 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. 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, the TransformerWithValueHead wrapper leverages this method to share the backbone between the actor (policy) and critic (value) networks.

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.

• 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 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 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, 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.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}