Skip to content

Model Heads

About

The d9d.module.block.head package handles the model heads.

Features

Causal Language Modelling

SplitLanguageModellingHead provides a causal language modelling head that computes per-token logprobs.

It uses efficient fused Linear-Cross-Entropy kernel from the Cut-Cross-Entropy project and avoids full logit tensor materialization.

Supports vocab split to multiple independent splits following the SplitTokenEmbeddings embedding implementation.

d9d.module.block.head

Task heads that turn backbone hidden states into a typed output, plus their IO contracts.

LM_IGNORE_INDEX = -100 module-attribute

Index ignored by LM head while calculating logps

ClassificationHead

Bases: TaskHead[SequencePoolingHeadShared, SequenceClassificationOutput]

A classification head module that is typically used on top of model hidden states.

It applies dropout followed by a linear projection to produce logits for a specified number of classes. It supports optional pooling via a mask, allowing for selection of specific tokens (e.g., [CLS] tokens or specific sequence positions) before projection.

__init__(hidden_size, num_labels, dropout)

Constructs the ClassificationHead object.

Parameters:

Name Type Description Default
hidden_size int

The input dimensionality (hidden state size).

 required
num_labels int

The number of output classes.

 required
dropout float

The dropout probability.

 required

forward(hidden_states, shared)

Computes class logits from hidden states.

Parameters:

Name Type Description Default
hidden_states Tensor

Input tensor of hidden states.

 required
shared SequencePoolingHeadShared

The head shared input. Its optional pooling_mask selects specific hidden states: the input is indexed as hidden_states[pooling_mask == 1], flattening the batch and sequence dimensions into a single dimension of selected tokens.

 required

Returns:

Type Description
SequenceClassificationOutput

The classification output holding the unnormalized logits.

reset_parameters()

Resets module parameters.

EmbeddingHead

Bases: TaskHead[SequencePoolingHeadShared, SequenceEmbeddingOutput]

A head module for extracting dense representations from hidden states.

It optionally applies a linear projection and L2 normalization to produce embeddings for contrastive learning or retrieval tasks. It supports boolean masking to select specific tokens (e.g., the last token) before extraction.

__init__(hidden_size, embedding_dim, normalize)

Constructs the EmbeddingHead object.

Parameters:

Name Type Description Default
hidden_size int

The input dimensionality (hidden state size).

 required
embedding_dim int | None

The dimensionality of the output embedding. If None, additional linear projection won't be applied.

 required
normalize bool

Whether to apply L2 normalization to the final embeddings.

 required

forward(hidden_states, shared)

Computes dense embeddings from hidden states.

Parameters:

Name Type Description Default
hidden_states Tensor

Input tensor of hidden states.

 required
shared SequencePoolingHeadShared

The head shared input. Its optional pooling_mask selects specific hidden states: the input is indexed as hidden_states[pooling_mask == 1], flattening the batch and sequence dimensions into a single dimension of selected tokens.

 required

Returns:

Type Description
SequenceEmbeddingOutput

The embedding output holding the extracted embeddings.

reset_parameters()

Resets module parameters.

SequenceCausalLMHeadShared dataclass

The shared input a causal language modeling head consumes.

Attributes:

Name Type Description
labels Tensor

Target tokens for the loss computation.

SequenceCausalLMOutput dataclass

The output of a causal language modeling head.

Attributes:

Name Type Description
logps Tensor

Per-token log-probabilities / loss, shape [batch, seq].

SequenceClassificationOutput dataclass

The output of a classification head.

Attributes:

Name Type Description
scores Tensor

Classification logits, shape [num_pooled_tokens, num_labels] when a pooling mask selects tokens, or [batch, seq, num_labels] when no pooling mask is used.

SequenceEmbeddingOutput dataclass

The output of an embedding head.

Attributes:

Name Type Description
embeddings Tensor

Pooled embeddings, shape [num_pooled_tokens, embedding_dim] when a pooling mask selects tokens, or [batch, seq, embedding_dim] when no pooling mask is used.

SequencePoolingHeadShared dataclass

The shared input a pooled head (classification/embedding) consumes.

Attributes:

Name Type Description
pooling_mask Tensor | None

Binary mask indicating which token(s) to pool. You can use d9d.dataset.token_pooling_mask_from_attention_mask to build it from an attention mask.

SplitLanguageModellingHead

Bases: TaskHead[SequenceCausalLMHeadShared, SequenceCausalLMOutput]

A segmented language modeling head computing per-token cross-entropy loss.

Computes per-token cross-entropy loss values using a composed weight matrix.

This class maintains separate linear layers for different segments of the vocabulary (e.g., regular vs. special tokens). During the forward pass, it concatenates the weights to form a unified projection matrix and computes the cross-entropy loss efficiently, typically using a fused kernel to avoid materializing full logits.

The concatenation order of the weights is determined by split_order, which ensures consistency with the global vocabulary indices.

__init__(split_vocab_size, split_order, hidden_size)

Constructs the SplitLanguageModellingHead object.

Parameters:

Name Type Description Default
split_vocab_size Mapping[str, int]

A dictionary mapping split names to their output vocabulary sizes.

 required
split_order Sequence[str]

A sequence defining the order in which vocabulary segments should be concatenated. This determines the mapping of global indices to specific heads.

 required
hidden_size int

The input dimensionality (hidden state size).

 required

forward(hidden_states, shared)

Computes the cross-entropy loss for the given hidden states and labels.

Parameters:

Name Type Description Default
hidden_states Tensor

Input tensor of shape (B, S, H).

 required
shared SequenceCausalLMHeadShared

The head shared input. Its labels of shape (B, S) must correspond to the global vocabulary formed by concatenating splits in split_order.

 required

Returns:

Type Description
SequenceCausalLMOutput

The causal LM output holding per-token loss values (reduction='none'), matching the shape of the labels tensor.

reset_parameters()

Resets module parameters.

TaskHead

Bases: Module, ModuleLateInit, ABC, Generic[THeadShared, THeadOutput]

Abstract base class for a task head that turns backbone hidden states into a typed output.

A head is confined to compute: it reads the shared hidden_states plus its own shared input PyTree, and returns its own output PyTree. Its parallelization and checkpoint mapping are kept out of the module as separate concerns.

Class Type Parameters:

Name Bound or Constraints Description Default
THeadShared

The head's own shared input, routed to it by name at composition time.

 required
THeadOutput

The output PyTree this head produces.

 required

forward(hidden_states, shared) abstractmethod

Computes the head output from hidden states and this head's shared input.

Parameters:

Name Type Description Default
hidden_states Tensor

Backbone hidden states of shape (B, S, H).

 required
shared THeadShared

The head's own shared input (e.g. labels, a pooling mask).

 required

Returns:

Type Description
THeadOutput

The head's output PyTree.

reset_parameters()

Resets the module parameters (i.e. performs random initialization).