Skip to content

Embeddings

About

The d9d.module.block.embedding package provides the SplitTokenEmbeddings token embedding layer. It is built from several named embedding tables, called splits. Each split owns a contiguous range of the global vocabulary, in the order given by split_order.

Use Cases

  • Regular token embedding: Use a single split with the full vocabulary size.
  • Prompt tuning: Add new tokens to your tokenizer and use two splits. The first split holds the original token embeddings, and the second one holds the new prompt tokens. Train only the nn.Embedding of the second split.

The SplitLanguageModellingHead uses the same splits for the output vocabulary.

Usage

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
import torch

from d9d.module.block.embedding import SplitTokenEmbeddings

embeddings = SplitTokenEmbeddings(
    split_vocab_size={"orig": 151936, "prompt": 16},
    split_order=["orig", "prompt"],
    hidden_size=2048,
)
embeddings.reset_parameters()

# Train only the new prompt tokens.
embeddings.token_embedding["orig"].requires_grad_(False)

# Indices 151936 and above belong to the "prompt" split.
input_ids = torch.tensor([[1, 2, 151936, 151937]])
hidden_states = embeddings(input_ids)  # (1, 4, 2048)

API Reference

d9d.module.block.embedding

Embedding layers.

SplitTokenEmbeddings

Bases: Module, ModuleLateInit

Token embedding layer composed of several named, independent embedding tables.

Each named split (e.g. "orig", "special", "prompt_prefix") owns a contiguous range of global vocabulary indices. This is useful for model adaptation, where different sets of tokens need different initialization or training behavior.

__init__(split_vocab_size, split_order, hidden_size)

Constructs the SplitTokenEmbeddings object.

Parameters:

Name Type Description Default
split_vocab_size dict[str, int]

Mapping from split names to their vocabulary sizes.

 required
split_order Sequence[str]

Order in which splits are concatenated to form the global vocabulary. Every name must be a key of split_vocab_size.

 required
hidden_size int

Dimensionality of the embedding vectors.

 required
forward(input_ids)

Looks up embeddings for global vocabulary indices.

Parameters:

Name Type Description Default
input_ids Tensor

Global vocabulary indices. Can have any shape.

 required

Returns:

Type Description
Tensor

Embeddings. Shape: (*input_ids.shape, hidden_size).

Raises:

Type Description
ValueError

If split_order is empty.
reset_parameters()

Resets the parameters of all embedding splits.