Model Heads

About

The d9d.module.block.head package provides task heads. A task head turns backbone hidden states into a typed output. Each head takes the hidden states and its own shared input, such as labels or a pooling mask.

Heads

Causal Language Modelling

SplitLanguageModellingHead computes the per-token cross-entropy loss for causal language modeling. Labels equal to LM_IGNORE_INDEX are ignored.

It uses the fused linear cross-entropy kernel from Cut Cross-Entropy (paper), so it does not materialize the full logits tensor. It requires the d9d[cce] extra.

The vocabulary can consist of several independent splits, like in SplitTokenEmbeddings.

Classification

ClassificationHead applies dropout and a linear projection, and returns fp32 logits.

Embedding

EmbeddingHead returns fp32 embeddings, with an optional linear projection and L2 normalization.

Both ClassificationHead and EmbeddingHead take an optional pooling_mask that selects the tokens to use. You can build it with d9d.dataset.token_pooling_mask_from_attention_mask.

Usage

import torch

from d9d.module.block.head import LM_IGNORE_INDEX, SequenceCausalLMHeadShared, SplitLanguageModellingHead

head = SplitLanguageModellingHead(
    split_vocab_size={"orig": 151936, "special": 8},
    split_order=["orig", "special"],
    hidden_size=2048,
).to(device="cuda", dtype=torch.bfloat16)

hidden_states = torch.randn(2, 16, 2048, device="cuda", dtype=torch.bfloat16)
labels = torch.randint(0, 151944, (2, 16), device="cuda")

output = head(hidden_states, SequenceCausalLMHeadShared(labels=labels))
loss = output.cross_entropy[labels != LM_IGNORE_INDEX].mean()

API Reference

`d9d.module.block.head`

Task heads that turn backbone hidden states into a typed output, and their I/O contracts.

`LM_IGNORE_INDEX = -100` `module-attribute`

Label index that the LM head ignores when computing the loss.

`ClassificationHead`

Bases: TaskHead[SequencePoolingHeadShared, SequenceClassificationOutput]

Classification head on top of model hidden states.

Applies dropout and a linear projection to produce logits for num_labels classes. An optional pooling mask selects specific tokens (e.g. [CLS] tokens) before the projection.

`init(hidden_size, num_labels, dropout)`

Constructs the ClassificationHead object.

Parameters:

Name	Type	Description	Default
`hidden_size`	`int`	Hidden size.	required
`num_labels`	`int`	Number of output classes.	required
`dropout`	`float`	Dropout probability.	required

`forward(hidden_states, shared)`

Computes class logits from hidden states.

Parameters:

Name	Type	Description	Default
`hidden_states`	`Tensor`	Hidden states. Shape: `(batch, seq_len, hidden_size)`.	required
`shared`	`SequencePoolingHeadShared`	The head shared input. Its optional `pooling_mask` selects hidden states as `hidden_states[pooling_mask == 1]`, which flattens the batch and sequence dimensions.	required

Returns:

Type	Description
`SequenceClassificationOutput`	The classification output holding the unnormalized fp32 logits.

`reset_parameters()`

Resets module parameters.

`EmbeddingHead`

Bases: TaskHead[SequencePoolingHeadShared, SequenceEmbeddingOutput]

Head that extracts dense embeddings from hidden states.

Optionally applies a linear projection and L2 normalization, e.g. for contrastive learning or retrieval. An optional pooling mask selects specific tokens (e.g. the last token) before the projection.

`init(hidden_size, embedding_dim, normalize)`

Constructs the EmbeddingHead object.

Parameters:

Name	Type	Description	Default
`hidden_size`	`int`	Hidden size.	required
`embedding_dim`	`int \| None`	Dimensionality of the output embedding. If `None`, no linear projection is 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`	Hidden states. Shape: `(batch, seq_len, hidden_size)`.	required
`shared`	`SequencePoolingHeadShared`	The head shared input. Its optional `pooling_mask` selects hidden states as `hidden_states[pooling_mask == 1]`, which flattens the batch and sequence dimensions.	required

Returns:

Type	Description
`SequenceEmbeddingOutput`	The embedding output holding the fp32 embeddings.

`reset_parameters()`

Resets module parameters.

`SequenceCausalLMHeadShared` `dataclass`

The shared input a causal language modeling head consumes.

Attributes:

Name	Type	Description
`labels`	`Tensor`	Target token indices for the loss computation. Shape: `(batch, seq_len)`.

`SequenceCausalLMOutput` `dataclass`

The output of a causal language modeling head.

Attributes:

Name	Type	Description
`cross_entropy`	`Tensor`	Per-token cross-entropy loss, the negative log-probability of each label. It is zero at ignored positions. Shape: `(batch, seq_len)`.

`SequenceClassificationOutput` `dataclass`

The output of a classification head.

Attributes:

Name	Type	Description
`scores`	`Tensor`	Classification logits. Shape: `(num_pooled_tokens, num_labels)` with a pooling mask, `(batch, seq_len, num_labels)` without one.

`SequenceEmbeddingOutput` `dataclass`

The output of an embedding head.

Attributes:

Name	Type	Description
`embeddings`	`Tensor`	Pooled embeddings. Shape: `(num_pooled_tokens, embedding_dim)` with a pooling mask, `(batch, seq_len, embedding_dim)` without one.

`SequencePoolingHeadShared` `dataclass`

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

Attributes:

Name	Type	Description
`pooling_mask`	`Tensor \| None`	Binary mask of the tokens to pool. You can build it from an attention mask with `d9d.dataset.token_pooling_mask_from_attention_mask`. Shape: `(batch, seq_len)`.

`SplitLanguageModellingHead`

Bases: TaskHead[SequenceCausalLMHeadShared, SequenceCausalLMOutput]

Language modeling head with a split vocabulary that computes per-token cross-entropy loss.

Each vocabulary split (e.g. regular and special tokens) has its own linear layer. The forward pass concatenates their weights in split_order and computes the loss with the fused Cut Cross-Entropy kernel, so full logits are not materialized. Requires the d9d[cce] extra.

`init(split_vocab_size, split_order, hidden_size)`

Constructs the SplitLanguageModellingHead object.

Parameters:

Name	Type	Description	Default
`split_vocab_size`	`Mapping[str, int]`	Mapping from split names to their vocabulary sizes.	required
`split_order`	`Sequence[str]`	Order in which vocabulary splits are concatenated. This defines which global indices belong to each split.	required
`hidden_size`	`int`	Hidden 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`	Hidden states. Shape: `(batch, seq_len, hidden_size)`.	required
`shared`	`SequenceCausalLMHeadShared`	The head shared input. Its `labels` index the global vocabulary formed by concatenating the splits in `split_order`. Labels equal to `LM_IGNORE_INDEX` are ignored.	required

Returns:

Type	Description
`SequenceCausalLMOutput`	The causal LM output holding the unreduced per-token loss, with the shape of `labels`.

`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. Shape: `(batch, seq_len, hidden_size)`.	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).

Model Heads

About

Heads

Causal Language Modelling

Classification

Embedding

Usage

API Reference

d9d.module.block.head

LM_IGNORE_INDEX = -100 module-attribute

ClassificationHead

__init__(hidden_size, num_labels, dropout)

forward(hidden_states, shared)

reset_parameters()

EmbeddingHead

__init__(hidden_size, embedding_dim, normalize)

forward(hidden_states, shared)

reset_parameters()

SequenceCausalLMHeadShared dataclass

SequenceCausalLMOutput dataclass

SequenceClassificationOutput dataclass

SequenceEmbeddingOutput dataclass

SequencePoolingHeadShared dataclass

SplitLanguageModellingHead

__init__(split_vocab_size, split_order, hidden_size)

forward(hidden_states, shared)

reset_parameters()

TaskHead

forward(hidden_states, shared) abstractmethod

reset_parameters()

`d9d.module.block.head`

`LM_IGNORE_INDEX = -100` `module-attribute`

`ClassificationHead`

`init(hidden_size, num_labels, dropout)`

`forward(hidden_states, shared)`

`reset_parameters()`

`EmbeddingHead`

`init(hidden_size, embedding_dim, normalize)`

`forward(hidden_states, shared)`

`reset_parameters()`

`SequenceCausalLMHeadShared` `dataclass`

`SequenceCausalLMOutput` `dataclass`

`SequenceClassificationOutput` `dataclass`

`SequenceEmbeddingOutput` `dataclass`

`SequencePoolingHeadShared` `dataclass`

`SplitLanguageModellingHead`

`init(split_vocab_size, split_order, hidden_size)`

`forward(hidden_states, shared)`

`reset_parameters()`

`TaskHead`

`forward(hidden_states, shared)` `abstractmethod`

`reset_parameters()`