M2.2 — Embeddings and Initial Model State

0x00 From Token IDs to Learned Vectors

In M2.1, we saw how a tokenizer divides text into tokens and assigns each token an integer ID. The next step is to convert these IDs into the vectors used in the Transformer’s calculations. This is the purpose of the input embedding layer, and the second stage of the text-to-vector conversion introduced at the beginning of Module 2.

Consider again the sentence:

The cat sleeps.

The DeepSeek-V3 tokenizer gives us four tokens and their IDs:

pieces: ["The", " cat", " sleeps", "."]
IDs:    [671, 7539, 114896, 16]

The ID 7539 identifies the token " cat" in the vocabulary. Its numerical value comes from the vocabulary’s numbering; it tells us nothing about how the word is used. Nearby IDs need not represent related words. We therefore associate each ID with a vector: an ordered list of numbers whose values can be adjusted as the model learns from text.

The model stores these vectors in an embedding table, with one row for each possible input ID. All rows have the same number of values, called the embedding dimension. The tokenizer’s ID tells us which row to select, an operation called an embedding lookup. If we call the table E, the four lookups for our sentence are:

671     selects E[671]       for "The"
7539    selects E[7539]      for " cat"
114896  selects E[114896]    for " sleeps"
16      selects E[16]        for "."

For example, E[7539] is the entire row for " cat". In DeepSeek-V3, the model configuration specifies 7,168 values per vector. Placing the four selected vectors in sentence order gives a matrix with four rows and 7,168 columns. This describes the sentence alone; the special tokens in M2.1’s formatted chat input select additional rows in exactly the same way.

The values in this table are learned along with the other model parameters. For a language model trained to predict the next token, training measures the errors in its predictions and adjusts the parameters to reduce those errors. This includes adjusting the embedding vectors. When we use a pretrained model, we load the resulting table from its saved weights.

At lookup, every occurrence of " cat" with ID 7539 selects the same row. Later layers use information from other tokens to change that representation. The input embedding is therefore the starting representation of a token; its representation in the context of a sentence is computed in the layers that follow.

Much of this article is devoted to how the embedding table behaves during training and how it is stored in a saved model. This accounts for its length: we will derive the updates to its rows, follow them through code, and examine what happens when the vocabulary or the way the table is stored changes.

We will first describe the table’s dimensions and explain how selecting a row relates to one-hot encoding. We will then implement lookup with PyTorch’s nn.Embedding and study initialization, the convention of multiplying the selected rows by sqrt(d_model), and training, including repeated tokens, padding, and whether the table’s gradient is stored densely or sparsely and how optimizers handle each form. Later sections examine ways to reduce the table’s memory use, add position and segment rows to the token rows, add rows for new tokens, share its values with the layer that predicts the next token, and divide the work across several GPUs. We will also compare lookup with the transformations used for image and audio features, and finish with tests and examples of implementation errors.

On a first reading, follow the lookup examples and the explanation of how training changes the selected rows. The longer derivations and checks for saving, sharing, or distributing the table can be revisited when working on those parts of a model.

0x01 Embedding Tables and Input Shapes

Our example has four tokens, each represented by a vector of 7,168 values. To describe the same operation for other sentences and models, we need notation for the number of rows and columns.

Let V be the number of rows in the embedding table, so valid IDs range from 0 to V - 1. We will begin with a table whose rows have the size used by the Transformer, written d_model. The table E then contains V × d_model learned values:

$$E \in \mathbb{R}^{V \times d_{\text{model}}}$$

Some models instead give the table a smaller row width, written d_embed, and map each selected row to d_model values with a learned matrix P of shape [d_embed, d_model], which this article calls the bridge. 0x05A explains this factorization and what it costs. The main examples take d_embed = d_model; the small trace in 0x02A and 0x02B uses a bridge with d_embed = 3 and d_model = 5.

Models commonly process several sequences together in a batch. Let B be the number of sequences and L their common length, after padding where necessary. Their token IDs form a two-dimensional array, or tensor, with shape [B,L]:

$$\text{input\_ids} \in \mathbb{N}^{B \times L}$$

For every entry in this tensor, lookup selects the corresponding row of E. We write the resulting collection of vectors as X_text:

$$X_{\text{text}} = E[\text{input\_ids}]$$

The batch still contains B sequences of L tokens. Each token now has d_model values, giving the output three dimensions:

$$X_{\text{text}} \in \mathbb{R}^{B \times L \times d_{\text{model}}}$$

For our sentence alone, B = 1 and L = 4: the input shape is [1,4], and the output shape is [1,4,7168]. The leading 1 records the batch size; inside it is the four-row matrix described in the introduction.

In an encoder-decoder model used for translation, this lookup is performed for both the source text and the target text. If their lengths are S and T, the corresponding output shapes are [B,S,d_model] and [B,T,d_model]. Each uses the embedding table associated with its vocabulary; whether those tables can be shared is discussed later.

Inputs that are already vectors

An image or audio encoder produces feature vectors before the input reaches the language model. These vectors may have a different number of values from the model’s text embeddings. To use them as input, we need to transform them to the required size.

Let H_media contain n feature vectors per example, each with d_media values. Its shape is:

$$H_{\text{media}} \in \mathbb{R}^{B \times n \times d_{\text{media}}}$$

One way to change the size is a learned linear layer, often called a projector. Its matrix W_p maps the d_media input values to d_model output values:

$$W_p \in \mathbb{R}^{d_{\text{media}} \times d_{\text{model}}}$$

We multiply each feature vector by this matrix and add a learned bias b, which has d_model values:

$$Z_{\text{media}} = H_{\text{media}} W_p + b$$

This operation preserves the number of feature vectors while changing their size:

$$Z_{\text{media}} \in \mathbb{R}^{B \times n \times d_{\text{model}}}$$

The same reasoning applies to document features and to event or action inputs: a discrete ID selects a row, while a continuous feature vector is transformed using its values. Casting those values to integers would not turn them into meaningful vocabulary IDs.

Some models use a multilayer network instead, which still produces one vector per feature vector, or a resampler or query adapter, which changes the number of vectors; the later sections on projectors explain these alternatives. Matching the vector size lets the language model accept the result. Learning to use the image or audio information also requires appropriate training.

The following diagram compares the two paths into the model:

Both paths end with vectors of d_model values. A linear projector keeps one output vector per feature vector; resamplers and query adapters change the count. “Position / layout / time” refers to the positions the language model supplies afterwards; a vision encoder may already have used patch positions to compute its features.

---

0x01A From symbols to addresses to learned rows

The historical problem is not that computers cannot store text. It is that a neural layer needs differentiable numeric state with a fixed channel width. A vocabulary first gives each discrete symbol an address:

token string -> integer ID in [0,V)

An integer ID by itself is a poor numeric feature. If dog -> 4812 and cat -> 2314, arithmetic on those integers invents an ordering that the vocabulary never intended. One-hot encoding removes that false magnitude:

$$\mathbf{o}_i \in \{0,1\}^{V}, \qquad (\mathbf{o}_i)_j=\mathbb{1}[j=i].$$

Every category is now equally distant from every other category, but the vector is sparse and V-dimensional. A learned matrix converts it to a dense row:

$$\mathbf{x}_i = \mathbf{o}_i E, \qquad E\in\mathbb{R}^{V\times d_{embed}}.$$

Because only component i of o_i is one, the product reduces to:

$$\mathbf{x}_i = E_{i,:}.$$

That is why embedding lookup has its final form: it is an optimized row gather for multiplication by a one-hot selector. The implementation does not need to materialize V-1 zeros.

The source figure uses Chinese annotations, but its object flow is exact: text is segmented, pieces receive integer indices, one-hot addresses are sparse and high-dimensional, and lookup produces dense lower-dimensional rows.

The terms sparse and dense here describe the representation, not whether the embedding parameter’s gradient is stored as a sparse tensor. Those are separate questions.

Why a learned row can carry useful structure

Initialization only assigns distinguishable coordinates. Training gives those coordinates task-dependent structure. If a loss L depends on selected row E_i, gradient descent applies:

$$E_i \leftarrow E_i - \eta\frac{\partial L}{\partial E_i}.$$

Rows used in contexts that demand similar predictions can receive updates with related directions. This does not guarantee a clean semantic geometry, and it does not mean every dimension receives a human-readable interpretation. It means the table is optimized jointly with the model objective.

The lookup table therefore has three identities at once:

address map:
    row i must still mean tokenizer ID i

parameter tensor:
    rows are learned under an optimizer and checkpoint

interface state:
    gathered rows establish scale and width for later computation

Breaking any one of them can preserve the tensor shape while invalidating the model.

From word rows to sentence vectors

Bengio et al. (2003) learned word rows together with a neural language model’s probability function. Similar words could then receive nearby vectors, so a word sequence never seen in training could still receive a high probability when its words resembled those of sentences that had been seen. Each row was still selected by a discrete word identity. Word2Vec (Mikolov et al., 2013) made such vectors cheap to train on billions of words with two simple architectures: CBOW predicts a word from its neighbours, and Skip-gram predicts the neighbours from the word. It exported a static table: each word, including a word with several meanings, kept exactly one row.

ELMo (Peters et al., 2018) made a token’s representation depend on its sentence by combining the internal layers of a bidirectional language model for each use of a word. Its input was not a word-table row: it computed each token’s context-independent representation from the token’s characters. In the Transformer (Vaswani et al., 2017) and BERT (Devlin et al., 2018), lookup rows became the initial state that attention layers transform. Sentence-BERT (Reimers and Gurevych, 2019) encodes each sentence once into a vector that can be compared by cosine similarity, instead of passing every pair of sentences through BERT. Finding the most similar pair among 10,000 sentences fell from about 65 hours to about 5 seconds.

A modern retrieval model contains several of these objects at once: token IDs, a lookup table, contextual token rows, a pooling rule, and a final indexed vector. The next section separates them.

0x01B “Embedding” names several different objects

The word embedding is overloaded. In model code, papers, and APIs it can refer to at least four objects with different construction paths and different valid comparisons.

Object	Construction	Typical shape	Context-dependent?	What makes it meaningful?
input lookup row	E[input_ids]	[B,L,d_model]	no, before later layers	language-model or task training of the lookup table
contextual token state	Transformer stack over input rows	[B,L,d_model]	yes	surrounding rows, masks, positions, and layer weights
pooled / retrieval vector	pool one or more contextual rows, often followed by a head and normalization	[B,d_retrieval]	yes	the pooling rule and contrastive/retrieval objective
projected modality state	projector maps encoder features into model width	[B,N,d_model]	depends on upstream encoder	paired multimodal training and the exact encoder/projector interface

Here d_retrieval is the width of the final retrieval vector: d_model when no head follows the pooling, otherwise the head’s output width.

Input lookup rows

M2.2’s main text object is:

$$X_{\text{lookup}} = E[\text{input\_ids}]$$

Before attention layers run, the same token ID selects the same table row regardless of sentence context, apart from any position/type additions made by the architecture:

"bank" in "river bank"   -> E[bank]
"bank" in "central bank" -> E[bank]

The context has not yet changed the token row. This differs from a classic Word2Vec-style standalone word embedding in training objective and vocabulary granularity, even though both are lookup tables.

Contextual token states

After the stack begins mixing rows, a token state depends on the sequence:

$$H^{(\ell+1)} = \operatorname{Block}_{\ell}(H^{(\ell)})$$

H_l["bank" | "river bank"]
    need not equal
H_l["bank" | "central bank"]

These outputs are often called contextual embeddings. They have the same [B,L,d_model] shape as the input lookup result, so shape alone cannot tell whether a tensor came from the embedding table or layer 20. Layer identity, masking, and position semantics are part of the object.

Pooled and retrieval vectors

A sentence/document/image embedding API usually adds another contract:

hidden = encoder(input_ids)                  # [B,L,d_model]
pooled = pool(hidden, attention_mask)        # [B,d_model]
z = projection(pooled)                       # [B,d_retrieval], optional
z = F.normalize(z, dim=-1)                   # only when the objective expects it

The pooling operation may use a designated row, masked mean pooling, attention pooling, or another learned head. A contrastive objective can then train nearby vectors to represent matched pairs and distant vectors to represent mismatches.

Without that objective, cosine similarity between arbitrary hidden states is a diagnostic number, not automatically a calibrated semantic-search score. Sentence-BERT is a useful concrete boundary: it modifies the BERT usage/training setup to produce sentence vectors that can be compared efficiently with cosine similarity. That property should not be inferred for every BERT hidden state or every mean-pooled language-model activation.

This is the end of M2.2’s treatment of pooling. M3.11 derives valid-row selection, pooling forward and backward rules, projection, normalization, contrastive objectives, task heads, and retrieval interaction contracts.

Projected modality states

The media branch in this article produces model-consumable state rows:

$$Z_{\text{media}} = \operatorname{Projector}(H_{\text{media}}) \in \mathbb{R}^{B\times N\times d_{\text{model}}}$$

Calling Z_media an image embedding does not make it a single retrieval vector. It may be a sequence of patch/query rows intended to participate in language attention. Conversely, a CLIP-style image embedding is commonly a pooled vector trained for cross-modal similarity. Both can originate from a vision encoder; their row count, objective, and consumer are different.

Use object-qualified names in debugging:

token lookup embeddings
layer-l contextual states
pooled retrieval embeddings
projected media state rows

The unqualified word embedding is not enough to establish shape, objective, normalization, or valid similarity metric.

The following diagram traces two paths through three of these objects:

The upper path produces the model’s initial input state; the lower path produces one retrieval vector per record. Before reusing a tensor from one path in the other, compare where it came from, which checkpoint produced it, what it could see, how it was pooled or normalized, what objective trained it, what consumes it, and which similarity metric it expects.

0x01C Ownership Handoff: Initial Rows Are Not Task Readout

M2.2’s implementation boundary ends with model-width initial state:

input_ids [B,L]
    -> input table and optional lexical bridge
        -> composed/scaled X0 [B,L,d_model]

The Transformer stack then creates contextual rows, and a later readout may select, pool, project, or normalize them:

X0 [B,L,d_model]
    -> M3.1-M3.7 contextual computation
        -> H [B,L,d_model]
            -> M3.11 row selection / pooling / task head
                -> task logits, values, spans, or retrieval representations

Equal widths do not collapse those stages:

Object	Producer	Immediate consumer	Canonical owner
lookup row	ID-addressed parameter table	input-state constructor	M2.2
initial model state	lookup, optional bridge/composition, and scale	Transformer blocks	M2.2
contextual row	Transformer block stack	later blocks or readout	M3.1-M3.7
pooled/task representation	selected contextual rows plus declared readout	task head or similarity function	M3.11
searchable index artifact	representation encoder plus corpus/index build	retrieval runtime	retrieval-system owner, with M8.10 deployment handoff

What crosses the boundary

M2.2 exports enough identity for downstream code to reconstruct the initial state:

tokenizer and ordered-vocabulary identity
input-table checkpoint and row map
padding and composed-address-space policy
factorization / projection / scale path
dtype, shard placement, and alias topology
X0 row order and validity metadata

It does not export a default sentence vector, pooling rule, contrastive objective, similarity threshold, or ANN index. Those objects cannot be inferred from X0.shape or from the fact that an API uses the word embedding.

M3.11 is the canonical owner for the next semantic boundary. It starts from contextual H, derives mask-correct pooling and its backward paths, types sequence/token/span heads, defines normalized retrieval representations and their objectives, and compares cross-, bi-, and late-interaction readouts. Physical ANN construction and fleet rollout remain downstream retrieval-system and M8.10 responsibilities.

This short handoff deliberately replaces a second full representation article inside M2.2. The terminology remains here because it prevents object confusion; the derivations live with the component that actually owns them.

---

0x02 Text path: IDs are addresses

From tokenization, the text side starts with integer IDs:

"cat" -> 2314
"dog" -> 4812
"."   -> 13

The integers are not meanings.

The fact that:

2314 < 4812

does not mean:

cat < dog

The ID is an address into a vocabulary-sized table. The learned row at that address is the first continuous representation the model sees.

For a small vocabulary:

ID	token
0	<pad>
1	<bos>
2	<eos>
13	.
2314	cat
4812	dog

the tokenized sequence:

cat dog .

might become:

input_ids = tensor([[2314, 4812, 13]])  # [B=1, L=3]

The Transformer block should not receive this as a numeric feature vector:

[2314, 4812, 13]

It should receive a sequence of dense rows:

[
  E[2314],
  E[4812],
  E[13],
]

The address/vector split is the local boundary:

token ID = which row?
embedding vector = learned row contents

This is why tokenizer/model mismatch is so destructive. If one tokenizer says ID 2314 means cat, but the loaded model trained row 2314 for another token, the shape is still valid. The meaning of the row address is wrong.

---

0x02A One complete tensor and code trace

Take:

B = 2
L = 4
V = 8
d_embed = 3
d_model = 5
pad_id = 0

with IDs:

input_ids = torch.tensor([
    [1, 4, 6, 2],
    [1, 5, 2, 0],
])  # [2,4]

attention_mask = input_ids.ne(0)  # [2,4]

The table is:

E: [V=8,d_embed=3]

Gather produces:

E[input_ids]: [2,4,3]

and a factorized bridge produces:

P: [3,5]
X = E[input_ids] @ P: [2,4,5]

The row trace is:

X[0,0] = E[1] @ P
X[0,1] = E[4] @ P
X[0,2] = E[6] @ P
X[0,3] = E[2] @ P

X[1,3] = E[0] @ P  # physical PAD row; later visibility must mask it

Code with assertions:

lookup = nn.Embedding(8, 3, padding_idx=0)
bridge = nn.Linear(3, 5, bias=False)

lexical = lookup(input_ids)
initial = bridge(lexical)

assert input_ids.shape == (2, 4)
assert input_ids.dtype == torch.int64
assert lexical.shape == (2, 4, 3)
assert initial.shape == (2, 4, 5)
assert attention_mask.shape == (2, 4)
assert torch.equal(initial[0, 0], bridge(lookup.weight[1]))

The mask is not applied by nn.Embedding; PAD lookup still materializes a row. Attention and pooling code must use the visibility metadata. padding_idx controls embedding-row gradient behavior, not the whole model’s attention.

Dtype and device trace

Typical objects use different dtypes:

input_ids:       integer index type
attention_mask:  bool or kernel/API-supported integer type
embedding rows:  bf16/fp16/fp32 according to weights/autocast
position IDs:    integer, owned by M3.3

Casting IDs to floating point is not required for lookup. Casting continuous rows back to integer destroys them. Device placement must also align:

assert input_ids.device == lookup.weight.device or runtime_can_route_ids()

An embedding on CPU with model blocks on GPU can be a deliberate offload design, but then the [B,L,d] transfer is part of latency. An accidental device copy is not free because lookup arithmetic is otherwise small.

Batch duplicates and gradients

In the example, ID 1 appears twice and ID 2 appears twice. Their input-side row gradients combine contributions from both batch rows. PAD ID 0 is selected once but receives no ordinary embedding gradient when padding_idx=0. Inspecting this trace catches three bugs at once: wrong address map, wrong PAD policy, and a mistaken assumption that each table row is updated at most once per batch.

0x02B Shape Is Not A Physical Tensor Contract

The previous trace established values and shapes. An implementation still has to answer a second set of questions:

where is each tensor stored?
what dtype interprets its bytes?
what stride maps a logical coordinate to storage?
which operation materializes a new tensor?
which artifact gives an integer its row meaning?

Those are not optional optimization details. A shape such as [2,4] cannot tell us whether its entries are integer addresses or floating features. A shape such as [8,3] cannot tell us which vocabulary ordering owns its rows. The immutable compatibility graph froze those semantic identities. This trace follows one admitted bundle through its forward tensor state.

For the synthetic forward fixture:

input_ids:
    shape       = [2,4]
    dtype       = int64
    device      = cpu
    strides     = [4,1] elements
    owner       = tokenizer_batch
    identity    = categorical_addresses:toy-token-batch-v1

embedding.weight:
    shape       = [8,3]
    dtype       = float32
    device      = cpu
    strides     = [3,1] elements
    owner       = embedding.weight
    identity    = toy-v8-d3-table-v1@toy-v8-ordered-address-space-v1

gathered_rows:
    shape       = [2,4,3]
    dtype       = float32
    device      = cpu
    strides     = [12,3,1] elements
    owner       = embedding_forward
    identity    = selected table + exact ID batch

The phrase row-major contiguous means that the last coordinate has stride one. If table row i has width d_e, coordinate k begins at logical element offset:

$$\operatorname{offset}_E(i,k)=i d_e+k.$$

For d_e=3, the selected rows in the fixture begin at:

ID	table row start, in elements	selected positions
0	0	(1,3), the PAD position
1	3	(0,0) and (1,0)
2	6	(0,3) and (1,2)
4	12	(0,1)
5	15	(1,1)
6	18	(0,2)

The contiguous output uses a different address equation:

$$\operatorname{offset}_X(b,\ell,k) =((bL+\ell)d_e)+k.$$

Therefore the two occurrences of ID 1 read the same logical source row but occupy distinct output ranges beginning at offsets 0 and 12. They have equal values; they are not aliases of the parameter row or of each other. This distinction matters later: backward must combine contributions from both positions into one owned parameter row, but forward has materialized two sequence positions.

PyTorch describes a tensor through storage, dtype, shape, stride, and storage offset, and warns that views and contiguous copies have different storage behavior. Its indexing note also distinguishes basic views from advanced indexing copies. The exact accelerator kernel may stage, cache, vectorize, or fuse row loads differently, so the equations above are a logical strided trace, not a claim about cache transactions.

Index dtype and weight dtype have different jobs

The official nn.Embedding contract accepts integer index tensors and appends embedding_dim to their shape. The parameter is floating point and the output follows the parameter/compute path. The address is not converted into a floating feature:

assert input_ids.dtype in (torch.int32, torch.int64)
assert embedding.weight.dtype.is_floating_point

rows = embedding(input_ids)

assert rows.shape == (*input_ids.shape, embedding.embedding_dim)
assert rows.dtype == embedding.weight.dtype  # absent another runtime policy

Framework autocast, quantized tables, CPU offload, and distributed routing can refine that last assertion. They must name the refinement. A primitive local lookup cannot silently use IDs on one device and a table on another. A runtime that permits that surface is performing an explicit copy or route whose owner, bytes, and synchronization belong in the runtime contract.

The PAD row is also physically selected during forward. padding_idx=0 does not remove position (1,3) from rows; it controls the table row’s gradient policy, and a newly constructed PyTorch embedding initializes that row to zeros. Attention visibility and pooled-representation masking remain separate operations.

Optional factorized bridge adds arithmetic, not addresses

For the same gathered tensor, the fixture applies:

P: [d_embed=3,d_model=5], float32, cpu, strides [5,1]
X_model = gathered_rows @ P

which gives:

[2,4,3] @ [3,5] -> [2,4,5]

The first two selected rows become:

E[1]       = [1, 0,  1]
E[1] @ P   = [1.5, 0, 1.5, -0.5, -1]

E[4]       = [2, 0, -1]
E[4] @ P   = [1.5, 0,   0,  0.5, -2]

The lookup still answers which lexical row. The bridge answers how that row enters model width. It costs:

$$\operatorname{MAC}_{bridge}=BLd_{embed}d_{model}.$$

For this fixture the bridge performs 2*4*3*5 = 120 MACs. Section 0x0A owns the full direct-versus-factorized architecture and parameter comparison; this section freezes only the forward tensor contract.

---

0x03 nn.Embedding and Harvard Embeddings.forward

The Harvard-style Transformer code makes the text bridge explicit:

class Embeddings(nn.Module):
    def __init__(self, d_model, vocab):
        super(Embeddings, self).__init__()
        self.lut = nn.Embedding(vocab, d_model)
        self.d_model = d_model

    def forward(self, x):
        return self.lut(x) * math.sqrt(self.d_model)

Each object has a concrete role.

Code object	Meaning
x	integer token IDs, not strings and not float features
nn.Embedding(vocab, d_model)	trainable lookup table with vocab rows and d_model columns
self.lut(x)	indexed row lookup
math.sqrt(self.d_model)	original Transformer / Harvard scaling convention
return value	content embeddings with shape [B,L,d_model]

Official PyTorch terminology is direct: nn.Embedding is a lookup table. Its trainable parameter has shape:

[num_embeddings, embedding_dim]

For this article:

num_embeddings = V
embedding_dim  = d_model

The input is an integer tensor of indices:

input_ids: LongTensor[B, L]

and the output appends the embedding dimension:

embedding(input_ids): FloatTensor[B, L, d_model]

So the source-compatible text bridge is:

embedding = nn.Embedding(
    num_embeddings=V,
    embedding_dim=d_model,
    padding_idx=pad_id,   # optional, but often important
)

x = embedding(input_ids)  # [B, L, d_model]

Then the Harvard wrapper multiplies by:

math.sqrt(d_model)

That scale is not a shape operation. It preserves [B,L,d_model]. But if the goal is to reproduce the original Transformer implementation closely, missing or double-applying that factor changes the numerical interface before positional encoding and attention.

The lookup table is also a training interface. During backpropagation, rows used by the batch receive gradient updates. If token 2314 appears, row E[2314] participates in the computation graph. Rows that do not appear in the batch usually do not receive a dense update from that example.

That is the first place token IDs become learned continuous states.

0x03A Initialization and the sqrt(d_model) convention

Before training, an embedding row is not a semantic dictionary entry. It is an initialized parameter vector. The scaling question therefore starts with the initializer, not with the word embedding.

Suppose the coordinates of row $E_i$ are independent with mean $\mu$ and variance $\sigma^2$. Linearity of expectation gives:

$$\mathbb{E}\left[\lVert E_i\rVert_2^2\right] =\sum_{k=1}^{d_{model}}\mathbb{E}[E_{i,k}^2] =d_{model}(\sigma^2+\mu^2).$$

For a zero-mean initializer and a scalar multiplier $s$:

$$\operatorname{Var}(sE_{i,k})=s^2\sigma^2, \qquad \mathbb{E}\left[\lVert sE_i\rVert_2^2\right] =s^2d_{model}\sigma^2.$$

This is the reusable derivation. If the desired coordinate variance is $\tau^2$, the matching scale is $s=\tau/\sigma$. The special choice $s=\sqrt{d_{model}}$ produces unit coordinate variance only when $\sigma^2=1/d_{model}$. It is not a property of every embedding initializer.

Reconstructing the Harvard source chain

The original Transformer paper states that embedding weights are multiplied by $\sqrt{d_{model}}$. It does not derive that choice from an initializer. The Harvard implementation exposes the multiplier in Embeddings.forward:

return self.lut(x) * math.sqrt(self.d_model)

But make_model later applies Xavier uniform initialization to every parameter with more than one dimension. That includes each rectangular embedding table:

for p in model.parameters():
    if p.dim() > 1:
        nn.init.xavier_uniform_(p)

For gain $g=1$, PyTorch defines Xavier uniform on a two-dimensional tensor as:

$$E_{i,k}\sim\mathcal{U}(-a,a), \qquad a=g\sqrt{\frac{6}{\operatorname{fan}_{in}+\operatorname{fan}_{out}}}.$$

Since a zero-mean uniform variable on $[-a,a]$ has variance $a^2/3$:

$$\sigma^2_{xavier} =\frac{2g^2}{\operatorname{fan}_{in}+\operatorname{fan}_{out}}.$$

For a PyTorch table stored as [V,d_model], the two-dimensional fan convention gives $\operatorname{fan}_{in}=d_{model}$ and $\operatorname{fan}_{out}=V$. After the Harvard multiplier:

$$\operatorname{Var}(\sqrt{d_{model}}E_{i,k}) =\frac{2g^2d_{model}}{V+d_{model}},$$ $$\mathbb{E}\left[\lVert\sqrt{d_{model}}E_i\rVert_2^2\right] =\frac{2g^2d_{model}^2}{V+d_{model}}.$$

At V=32,000, d_model=512, and g=1:

Xavier coordinate variance before scale:  2 / 32,512 ~= 0.00006151
coordinate variance after sqrt(512):       0.03150
expected squared row norm after scale:     16.13
approximate RMS per scaled coordinate:     0.1775

So the common sentence “Xavier has variance 1/d_model, then sqrt(d_model) restores variance one” is a useful square-matrix intuition, but it is not the exact Harvard/PyTorch rectangular-table calculation.

The positional comparison is also conditional. If the content coordinates have RMS $s\sigma$ and the positional coordinates at a given position have RMS $\rho_p$, then their pre-dropout RMS ratio is:

$$R_{content:position}=\frac{s\sigma}{\rho_p}.$$

Sinusoidal coordinates are bounded by one, but their exact RMS depends on the position and frequencies. Multiplication by $\sqrt{d_{model}}$ increases the content side of this ratio; it does not prove equal magnitude and does not by itself prove a training-stability theorem.

This convention is not universal. Modern architectures may use different initializers, learned positions, RoPE, embedding normalization, residual scaling, or no explicit embedding multiplier. Reproduce the checkpoint’s architecture; do not add sqrt(d_model) merely because the tensor is called an embedding.

Source-qualified input contract	Fresh initializer	Composition before the stack	Explicit token scale
original Transformer / Harvard reconstruction	Harvard’s global Xavier pass reaches [V,d_model]	scaled token row + fixed sinusoidal position, then dropout	sqrt(d_model)
original BERT code	truncated normal with configured initializer_range (default 0.02)	token + token-type + learned-position rows, then LayerNorm and dropout	none in the embedding path
GPT-2 code lineage	model-specific normal initialization	token row + learned-position row, then dropout	none in the embedding path
LLaMA-style RoPE decoder	model/checkpoint-specific	token row enters the residual stream; position rotates attention queries and keys	no additive-position balancing rule to infer

This table compares source contracts, not model quality. M3.3 owns the geometry and extrapolation behavior of sinusoidal, learned, relative, and rotary position mechanisms.

The common explanation “the scale directly prevents attention dot products from exploding” is incomplete. Attention separately scales QK^T by 1/sqrt(d_k). The input multiplier establishes the content-state scale before projection and composition; its effect on later dot products is mediated by the learned projection and normalization stack.

Fresh initialization, exact load, and growth are different states

A constructor, a checkpoint loader, and a vocabulary-resize procedure may all produce a tensor shaped [V,d_model]. They do not have the same permission to write its rows.

Materialization state	Rows allowed to change	Required evidence	Forbidden shortcut
fresh initialization	all rows, according to a named initializer and post-init policy	seed/generator, initializer parameters, post-init mutators, scale owner	describing framework defaults as the checkpoint’s trained distribution
exact checkpoint load	none, except an admitted dtype/device materialization	source key, shape, address hash, source/target dtype, finite check, tensor hash	silently retaining random constructor rows for a missing key
vocabulary growth	appended or explicitly remapped rows only	old-prefix equality, new-row initializer, tokenizer/config update, optimizer-state expansion, tie restoration	resizing the input table without its output and optimizer dependents

def materialize_embedding(spec, checkpoint=None):
    table = allocate(spec.shape, dtype=spec.parameter_dtype)

    if spec.state == "fresh_initialization":
        initialize_(table, spec.initializer, generator=spec.generator)
        apply_post_init_constraints_(table, spec.padding_policy)

    elif spec.state == "exact_checkpoint_load":
        saved = require_tensor(checkpoint, spec.source_key)
        require_shape_address_and_finite(saved, spec)
        table.copy_(admitted_cast(saved, spec.parameter_dtype))

    elif spec.state == "vocabulary_growth":
        old = require_tensor(checkpoint, spec.source_key)
        require_exact_old_prefix(old, spec)
        table[: old.shape[0]].copy_(old)
        initialize_(table[old.shape[0] :], spec.new_row_initializer)
        expand_optimizer_rows_and_restore_ties(spec)

    else:
        raise ValueError("unknown embedding materialization state")

    return table

An exact load checks more than shape:

def assert_embedding_checkpoint(module, state_dict, key, address_hash):
    saved = state_dict[key]
    assert saved.shape == module.weight.shape
    assert state_dict["tokenizer_address_hash"] == address_hash
    assert torch.isfinite(saved).all()
    assert saved.dtype == module.weight.dtype or can_safely_cast(saved.dtype)

Current Hugging Face resize_token_embeddings documentation exposes both a model-default initializer and a mean/covariance-based growth option, and says the model re-ties weights when its class implements tie_weights(). That is a version-scoped helper contract, not a universal rule. A custom implementation still has to prove that old rows were preserved, new output rows and biases were handled, optimizer state was expanded, aliases were restored, and the tokenizer hash now names the larger address space. M2.1 owns tokenizer-side allocation; Section 0x05D owns the full release lifecycle.

0x03B nn.Embedding options change training behavior

The compact constructor hides several independent policies:

nn.Embedding(
    num_embeddings=V,
    embedding_dim=d_model,
    padding_idx=pad_id,
    max_norm=None,
    norm_type=2.0,
    scale_grad_by_freq=False,
    sparse=False,
    _weight=None,
)

Their contracts are:

Option	Mechanism	Integration consequence
padding_idx	prevents that row from receiving ordinary embedding gradients	masks and tokenizer PAD ID must agree
max_norm	renormalizes selected rows whose norm exceeds the threshold	forward performs an in-place weight modification
norm_type	chooses the norm used by max_norm	usually L2, but not fixed
scale_grad_by_freq	divides row gradient by batch frequency	changes repeated-token weighting
sparse	emits sparse gradient for table weights	restricts compatible optimizers and distributed paths
_weight	initializes from supplied table	supplied shape and padding semantics must match

padding_idx is a gradient rule, not a universal zero-row invariant

In the current PyTorch source, a freshly constructed nn.Embedding samples its weight with normal_() and then zeros weight[padding_idx]. Supplying _weight skips reset_parameters, so the supplied padding row is not automatically zeroed. A later global initializer, checkpoint load, direct assignment, or tied output path can also produce a nonzero row.

Forward lookup still returns the row currently stored at padding_idx. padding_idx suppresses the ordinary lookup-gradient contribution for that row; it does not make the row invisible to attention, remove it from a loss, or protect it from every other mutation path.

For example, with dense AdamW and a present gradient tensor, decoupled weight decay applies to the parameter before the gradient update:

$$E_{pad}^{+}=(1-\eta\lambda)E_{pad} \qquad\text{when}\qquad \nabla_{E_{pad}}\mathcal{L}=0.$$

A zero row remains zero under that multiplication. A nonzero custom PAD row does not. Optimizers may skip parameters whose gradient is absent, and sparse optimizer paths have different support, so record the actual optimizer and gradient representation rather than universalizing this equation.

max_norm makes lookup forward stateful

Let $S$ be the set of distinct row IDs selected by the input. Ignoring the kernel’s small numerical guard, max-norm applies:

$$E_i\leftarrow \begin{cases} E_i\dfrac{m}{\lVert E_i\rVert_p}, & i\in S\ \text{and}\ \lVert E_i\rVert_p>m,\\ E_i, & \text{otherwise}. \end{cases}$$

Only selected rows are candidates. Repeated IDs do not imply repeated semantic renormalization. With p=2, m=2, and selected IDs [1,2,1,0]:

row 1 before: [3.0, 4.0, 0.0, 0.0]   norm 5.0
row 1 after:  [1.2, 1.6, 0.0, 0.0]   norm 2.0

row 2 before: [0.0, 0.0, 0.6, 0.8]   norm 1.0
row 2 after:  unchanged

row 3 before: [6.0, 8.0, 0.0, 0.0]   norm 10.0
row 3 after:  unchanged because row 3 was not selected

PyTorch performs this selected-row renormalization in-place during forward, before returning gathered rows. A checkpoint or table hash can therefore change after an inference-only call even when backward and the optimizer never run. If a differentiable computation uses embedding.weight before the embedding call, clone the weight or reorder operations according to the official API warning; otherwise autograd can see a versioned tensor modified in place.

embedding = nn.Embedding(V, d, max_norm=1.0)

# Clone if this pre-forward use participates in autograd.
penalty = regularizer(embedding.weight.clone())
x = embedding(input_ids)
loss = task_loss(x) + penalty

The implementation trace must name the mutation explicitly:

selected = torch.unique(input_ids)
before = embedding.weight.detach()[selected].clone()
output = embedding(input_ids)  # may mutate selected rows when max_norm is set
after = embedding.weight.detach()[selected]

assert torch.linalg.vector_norm(after, ord=2, dim=-1).max() <= 1.0 + tolerance
record_rows_that_changed(selected, before, after)

With sparse=True, inspect the optimizer rather than assuming any optimizer is valid. Sparse gradients also interact with gradient clipping, weight decay, and distributed collectives differently from dense parameters.

0x03C Numerical state is a ledger, not one dtype

“The embedding is bf16” is underspecified. At least six numerical states can exist around one logical table:

State	Example dtype	Owner	What must be recorded
serialized checkpoint tensor	fp32, bf16, fp16, or encoded/quantized	checkpoint format	logical tensor key, encoding, source dtype, hashes
live parameter storage	fp32, bf16, fp16	module/runtime	storage dtype, device, sharding, alias identity
gathered forward output	usually follows the live table for direct lookup	operator dispatch	observed output dtype, subsequent scale/add cast policy
autocast compute	op-specific fp32, fp16, or bf16	autocast dispatcher	enabled region, target dtype, actual operator behavior
gradient	dense or sparse; dtype/runtime dependent	autograd	layout, dtype, scaling state, reduction owner
optimizer/master state	often fp32 moments; master copy optional	optimizer/runtime	moment dtype, master-copy existence, update dtype

PyTorch autocast chooses dtypes per operator; entering an autocast region does not prove that the parameter itself was converted. Gradient scaling multiplies the loss to protect small fp16 gradients and unscales before the optimizer step; it does not define checkpoint dtype or guarantee a separate fp32 master table. Some systems keep fp32 parameters and use lower-precision compute, some store lower-precision parameters with fp32 optimizer state, and some maintain an explicit master copy. Those are different manifests.

embedding_numerical_state:
  checkpoint_dtype: float32
  parameter_storage_dtype: bfloat16
  lookup_output_dtype: bfloat16
  input_scale: sqrt_d_model
  scale_application_count: 1
  autocast:
    enabled: true
    target_dtype: bfloat16
    observed_lookup_policy: follows_parameter_storage
  gradient:
    layout: dense
    dtype: bfloat16
    grad_scaler: false
  optimizer:
    name: AdamW
    first_moment_dtype: float32
    second_moment_dtype: float32
    master_weight: explicit_float32_copy
  max_norm:
    enabled: false

This is an illustrative manifest, not PyTorch’s universal AMP layout. A reproducible run records observed tensor dtypes and optimizer state. It also checks finite values after admitted casts: lower precision has finite range and rounding limits, and mathematically equivalent operation orders need not be bitwise identical across devices or releases.

The visual has two independent axes. The horizontal path controls which rows may be written. The vertical ledger controls which representation of those rows exists at each runtime stage. A valid shape on one axis does not repair a missing contract on the other.

Executable Numerical-State Gate

The local NumPy oracle freezes this contract with a synthetic [5,4] table. It checks four initializer ledgers, loads a checkpoint whose PAD row is intentionally nonzero, applies max_norm=2 to IDs [1,2,1,0], and proves:

changed by forward:       row 1, [3,4,0,0] -> [1.2,1.6,0,0]
selected but unchanged:   rows 0 and 2
unselected and unchanged: rows 3 and 4, even though row 3 norm is 10
growth transition:        old five-row prefix preserved exactly
new optimizer rows:       explicit zero state
weight tie:               restored after growth

Eighteen invalid fixtures reject mixed state provenance, undeclared initializers, address/shape/cast mismatches, non-finite values, invalid norm policies, out-of-range IDs, mutated old rows, missing new-row policy, malformed optimizer state, a lost tied alias, double scaling, and an implicit master-weight assumption. Passing the fixture proves those declared contracts. It does not prove a framework kernel, random-sample distribution, convergence result, or production AMP layout.

---

0x04 One-hot as conceptual row selection

One-hot vectors are useful for understanding embedding lookup.

For a vocabulary of size V, a token ID can be represented as a vector with one 1 and all other entries 0:

id = 3
one_hot(id) = [0, 0, 0, 1, 0, ...]

If:

$$E \in \mathbb{R}^{V \times d_{\text{model}}}$$

then:

$$\text{one\_hot}(id) E = E[id]$$

Write the scalar selector as (e_i\in{0,1}^{V}), where:

$$(e_i)_v=\mathbb{1}[v=i].$$

Then coordinate (k) of the product is:

$$(e_i^\top E)_k =\sum_{v=0}^{V-1}(e_i)_vE_{v,k} =\sum_{v=0}^{V-1}\mathbb{1}[v=i]E_{v,k} =E_{i,k}.$$

The final form is a gather because every term except row i is multiplied by zero. This is exact algebra, not the statement that a production lookup first builds a sparse vector.

For a sequence (I=(i_0,\ldots,i_{L-1})), stack the selectors into:

$$S\in\{0,1\}^{L\times V}, \qquad S_{\ell,v}=\mathbb{1}[v=i_\ell].$$

The sequence result is:

$$X=SE\in\mathbb{R}^{L\times d_e}, \qquad X_{\ell,k}=E_{i_\ell,k}.$$

For a whole batch:

$$\text{one\_hot}(\text{input\_ids}) \in \{0,1\}^{B \times L \times V}$$ $$\text{one\_hot}(\text{input\_ids}) E \in \mathbb{R}^{B \times L \times d_{\text{model}}}$$

More explicitly, let:

$$S_{b,\ell,v}=\mathbb{1}[v=I_{b,\ell}].$$

Contracting the vocabulary axis gives:

$$X_{b,\ell,k} =\sum_{v=0}^{V-1}S_{b,\ell,v}E_{v,k} =E_{I_{b,\ell},k}.$$

This is the conceptual bridge.

It is not the practical implementation.

For real vocabularies, the one-hot tensor would be mostly zeros and very large. nn.Embedding does indexed lookup instead. It returns the same selected rows without materializing the one-hot tensor.

That distinction matters for bugs:

correct:
    nn.Embedding(...)(input_ids)

conceptual:
    one_hot(input_ids) @ E

wrong:
    nn.Embedding(...)(one_hot(input_ids))

nn.Embedding expects integer indices. It does not expect strings, one-hot vectors, or already embedded floats.

The declared vocabulary width is essential. Inferring num_classes from the largest ID in one batch produces a selector whose last dimension changes with the sample. The table has a release-wide row order, so the conceptual selector must always use that declared V, even when the current batch touches only a small subset.

The executable equivalent is:

def one_hot_reference(ids, weight):
    V = weight.shape[0]  # declared table width, not ids.max()+1
    selectors = F.one_hot(ids, num_classes=V)
    selectors = selectors.to(dtype=weight.dtype, device=weight.device)
    return selectors @ weight

def gather_path(ids, weight):
    return F.embedding(ids, weight)

dense = one_hot_reference(input_ids, embedding.weight)
gathered = gather_path(input_ids, embedding.weight)
torch.testing.assert_close(dense, gathered)

torch.nn.functional.one_hot returns a LongTensor, so the conceptual matrix product needs an explicit cast to the table’s compute dtype. The gather path avoids both the casted selector and its vocabulary-sized intermediate.

Why materializing one-hot is the wrong execution plan

For B=8, L=2048, and V=128,000:

ID entries:              8 * 2048             = 16,384
one-hot entries:         8 * 2048 * 128,000   = 2,097,152,000

int64 IDs:                                      128 KiB
int64 one-hot result:                            15.625 GiB
float16 selector after cast:                      3.906 GiB
entry-count ratio:                              128,000x

Those numbers count logical payload only. They do not include allocator, kernel, or temporary-workspace overhead. Indexed lookup reduces the operation to the information actually needed: read BL addresses and produce BLd_e floating values. There is no mathematical need to execute the zero multiplications in a dense [BL,V] @ [V,d_e] product.

---

0x04A Backward pass: gather becomes scatter-add

Owner question. Given repeated IDs, optional frequency scaling, sparse or dense gradient storage, and a table spread across ranks, which contributions belong to each row and which process is allowed to apply its update?

The answer has two layers that must not be collapsed:

mathematical owner:
    every occurrence of row j contributes to one logical dE[j]

execution owner:
    a declared rank stores, reduces, and updates all or part of row j

The first layer is invariant. The second can use replication, routed rows, masked local lookup, column shards, fused optimizers, or other schedules.

Deriving scatter-add from the forward selector

Flatten the B x L position axes into N=BL. Let i_t be the global token ID at position t, let E in R^{V x d} be the table, and define the conceptual selector:

$$S_{t,j}=\mathbb{1}[i_t=j], \qquad S\in\{0,1\}^{N\times V}.$$

The gather proof from 0x04 can be written as a matrix product without materializing S:

$$X=SE, \qquad X\in\mathbb{R}^{N\times d}.$$

For upstream gradient $G=\partial\mathcal{L}/\partial X$, take a differential:

$$dX=S\,dE,$$ $$d\mathcal{L} =\langle G,dX\rangle_F =\operatorname{tr}(G^TS\,dE) =\operatorname{tr}((S^TG)^T dE).$$

Therefore:

$$\boxed{\frac{\partial\mathcal{L}}{\partial E}=S^TG}$$

and row j receives:

$$\boxed{ \frac{\partial\mathcal{L}}{\partial E_j} =\sum_{t=1}^{N}\mathbb{1}[i_t=j]G_t }$$

The transpose changes a forward gather into a backward scatter-add. It also explains why duplicate IDs cannot use overwrite semantics: several rows of G map to the same parameter row.

If padding_idx=p, the embedding operation removes that row’s ordinary lookup gradient. Define:

$$S^{(p)}_{t,j}=\mathbb{1}[i_t=j]\mathbb{1}[j\ne p].$$

Then the input-path gradient is $(S^{(p)})^TG$. This does not erase the PAD row from other parameter uses, regularizers, or a tied output path.

One duplicate-ID tensor trace

Use V=8, d=3, padding_idx=0, and two local batches:

origin rank 0
IDs:       [1,       4,       1,       0]
G rows:    [[1,0,1], [0,2,0], [3,1,0], [9,9,9]]

origin rank 1
IDs:       [4,       6,       1,       4]
G rows:    [[2,0,1], [0,1,3], [1,2,1], [4,1,0]]

The PAD occurrence contributes nothing through this operator. Summing all other occurrences gives:

dE[1] = [1,0,1] + [3,1,0] + [1,2,1] = [5,3,2]
dE[4] = [0,2,0] + [2,0,1] + [4,1,0] = [6,3,1]
dE[6] = [0,1,3]
dE[0] = [0,0,0] through the lookup path
all other rows = [0,0,0]

A dense reference keeps the rule executable:

def embedding_backward_dense(ids, grad_output, vocab_size, padding_idx=None):
    ids = ids.reshape(-1)
    rows = grad_output.reshape(ids.numel(), -1)
    keep = torch.ones_like(ids, dtype=torch.bool)
    if padding_idx is not None:
        keep &= ids != padding_idx

    grad_weight = rows.new_zeros((vocab_size, rows.shape[-1]))
    grad_weight.index_add_(0, ids[keep], rows[keep])
    return grad_weight

The implementation may use atomics, sorting, segmented reduction, or a fused kernel. The article’s correctness contract is the resulting sum, not one ordering of floating-point additions. Deterministic mode and device choice can therefore affect bitwise results without changing the algebraic owner.

Frequency scaling has a count scope

Let the non-PAD count of row j in the operator’s local mini-batch be:

$$c_j=\sum_{t=1}^{N}\mathbb{1}[i_t=j]\mathbb{1}[j\ne p].$$

With scale_grad_by_freq=True, the bounded PyTorch contract is:

$$\frac{\partial\mathcal{L}}{\partial E_j} = \begin{cases} \dfrac{1}{c_j}\sum_{t:i_t=j}G_t,& c_j>0,\\ 0,& c_j=0. \end{cases}$$

The count belongs to the embedding invocation. In ordinary data parallelism, each rank executes that invocation before gradient synchronization. Rank-local frequency scaling followed by a sum is generally not global-frequency scaling:

$$\sum_{r=1}^{R}\frac{1}{c_j^{(r)}} \sum_{t:i_t^{(r)}=j}G_t^{(r)} \ne \frac{1}{\sum_r c_j^{(r)}} \sum_r\sum_{t:i_t^{(r)}=j}G_t^{(r)}.$$

For the trace above:

rank-local frequency scaling, then SUM
row 1: [4,1,1]/2 + [1,2,1]/1 = [3,2.5,1.5]
row 4: [0,2,0]/1 + [6,1,1]/2 = [3,2.5,0.5]

global effective-batch frequency scaling
row 1: [5,3,2]/3 = [1.6667,1,0.6667]
row 4: [6,3,1]/3 = [2,1,0.3333]

Neither expression is automatically the intended training objective. Record the count scope, data-parallel reduction (SUM or AVG), gradient accumulation window, and loss normalization. Changing world size or microbatch partitioning can otherwise change the effective row update while every tensor shape remains valid.

The options are not freely composable in current PyTorch. In the pinned Embedding.cpp source, embedding_sparse_backward rejects scale_grad_by_freq=True with a not-supported error. The equations above describe the dense frequency-scaled derivative and the distributed count-scope problem; a PyTorch sparse-layout execution must leave frequency scaling off or use a separately implemented and verified reduction path.

Dense and sparse are storage contracts, not different derivatives

The dense representation materializes [V,d], including zero rows. A sparse COO-style representation can initially store one row value per valid occurrence:

uncoalesced row indices: [1,4,1,4,6,1,4]
uncoalesced values:      seven [d] upstream rows

After coalescing duplicate indices:

coalesced row indices: [1,4,6]
coalesced values:      [[5,3,2], [6,3,1], [0,1,3]]
logical dense shape:   [8,3]

Coalescing does not change the derivative. It establishes one stored value per touched row and makes duplicate summation explicit:

def coalesced_sparse_rows(ids, grad_output, padding_idx=None):
    ids = ids.reshape(-1)
    values = grad_output.reshape(ids.numel(), -1)
    if padding_idx is not None:
        keep = ids != padding_idx
        ids, values = ids[keep], values[keep]

    unique_ids, inverse = torch.unique(ids, sorted=True, return_inverse=True)
    coalesced = values.new_zeros((unique_ids.numel(), values.shape[-1]))
    coalesced.index_add_(0, inverse, values)
    return unique_ids, coalesced

PyTorch’s current nn.Embedding documentation says sparse=True makes the weight gradient a sparse-layout tensor. That is narrower than saying the parameter, optimizer state, checkpoint, or every gradient path is sparse. SparseAdam explicitly targets dense-layout parameters receiving sparse-layout gradients and updates only materialized entries and their moments.

For V=128,000, d=4096, fp32 gradient values, and U=4096 unique touched rows:

dense gradient values:       V*d*4 = 1.953 GiB
coalesced sparse values:      U*d*4 = 64 MiB
int64 row indices:            U*8   = 32 KiB

This is a logical payload ledger. Sparse tensor objects, allocator metadata, sort/coalesce workspace, optimizer state, and collective buffers add overhead. If almost every row is touched, or a later path creates a dense contribution, the sparse advantage shrinks or disappears.

Optimizer, clipping, and regularization contracts

Current PyTorch nn.Embedding documentation lists only these built-in optimizer families for sparse embedding gradients:

Optimizer	CPU sparse gradient	CUDA sparse gradient	Important qualification
SGD	documented	documented	supported optimizer name does not prove every option, weight-decay, momentum, foreach, or fused combination
SparseAdam	documented	documented	dense parameter plus sparse-layout gradient; masked Adam approximation, no generic dense AdamW semantics
Adagrad	documented	not listed	CPU-only sparse support in the current embedding note
Adam, AdamW, and unlisted optimizers	do not assume	do not assume	require an explicit compatible path or dense gradient

An admission check should inspect the actual gradient after backward:

loss.backward()
grad = embedding.weight.grad

assert grad is not None
if embedding.sparse:
    assert grad.layout == torch.sparse_coo
    grad = grad.coalesce()
    require_sparse_optimizer_contract(optimizer, grad.device)
else:
    assert grad.layout == torch.strided

Gradient clipping also needs a representation and ownership policy. For a coalesced sparse gradient sharded by row, the global L2 norm can be computed without densifying:

$$q_r=\sum_{u\in\mathcal{U}_r}\lVert g_u\rVert_2^2, \qquad q=\operatorname{AllReduceSUM}(q_r),$$ $$\gamma=\min\left(1,\frac{m}{\sqrt{q}+\epsilon}\right), \qquad g_u\leftarrow\gamma g_u.$$

Calling a generic clipping helper without checking sparse-layout and sharded support is not a contract. PyTorch FSDP, for example, exposes its own global-norm method when gradients are sharded because local clip_grad_norm_ is not enough.

Regularization can deliberately destroy touched-row-only semantics. Dense decoupled weight decay conceptually updates every stepped row, including rows absent from the batch. A sparse system must choose among no decay, a supported masked approximation, a lazy timestamped decay, periodic dense maintenance, or densification. Those choices are not numerically interchangeable.

A tied output head can make the total gradient dense

Section 0x0A owns the full weight-tying contract, but this backward derivation must protect the gradient boundary. If logits use the same table:

$$Z=HE^T, \qquad H\in\mathbb{R}^{N\times d}, \qquad Z\in\mathbb{R}^{N\times V},$$

then for $D=\partial\mathcal{L}/\partial Z$:

$$\left.\frac{\partial\mathcal{L}}{\partial E}\right|_{output}=D^TH.$$

D normally has a vocabulary-wide softmax contribution, so $D^TH$ is generally dense even when the input lookup touched only a few rows. A tied parameter gets:

$$\frac{\partial\mathcal{L}}{\partial E} = \underbrace{(S^{(p)})^TG}_{\text{sparse-capable input path}} + \underbrace{D^TH}_{\text{generally dense output path}}.$$

sparse=True on the input module therefore does not prove that the shared parameter’s total gradient remains sparse. The runtime must admit a common layout and optimizer path; silently dropping either contribution is a correctness failure.

Distributed ownership starts with the objective convention

Let origin rank r own a local batch and produce $G^{(r)}_{local}$. For a replicated table with SUM reduction:

$$G_E=\sum_{r=0}^{R-1}(S_r^{(p)})^TG_r.$$

With AVG, either the collective or the loss convention introduces 1/R. Applying both divides twice; applying neither changes the effective learning rate. The manifest must name:

loss reduction within each rank
microbatch accumulation rule
data-parallel gradient collective: SUM or AVG
model-parallel ownership axis
optimizer-state owner

Data parallelism and model parallelism are independent axes. A row can be partitioned across tensor-parallel ranks and each resulting shard can still be replicated across a data-parallel group.

Schedule A: replicated table plus data-parallel reduction

Each rank holds all V*d parameters and optimizer state:

local_grad = scatter_add(local_ids, local_upstream, shape=(V, d))
global_grad = all_reduce_sum(local_grad, data_parallel_group)
embedding.weight.grad.copy_(global_grad)
optimizer.step()  # identical state and update on every replica

This has simple address semantics but a dense gradient collective when the runtime uses ordinary dense DDP buckets. Sparse communication needs a separate supported design; sparse=True alone does not rewrite a data-parallel reducer.

Schedule B: row owners with routed IDs and gradients

Partition rows into disjoint intervals [s_r,e_r) that cover [0,V) exactly once. The owner function is:

$$o(j)=r\quad\text{iff}\quad s_r\le j<e_r.$$

Forward preserves origin-position metadata:

requests = pack_by_owner(input_ids, origin_rank, origin_position)
owner_requests = all_to_all(requests, model_parallel_group)

owner_rows = local_table[owner_requests.global_id - shard_start]
responses = pack_by_origin(owner_rows, owner_requests.origin_position)
X = all_to_all(responses, model_parallel_group).restore_position_order()

Backward reverses the route:

owner_grads = all_to_all(
    pack_by_saved_owner(grad_output, origin_position),
    model_parallel_group,
)

local_grad = scatter_add(
    owner_grads.global_id - shard_start,
    owner_grads.value,
    shape=local_table.shape,
)
local_optimizer.step(local_table, local_grad)

The origin rank owns sequence order. The row-owner rank owns parameter and optimizer state. Losing origin_position, applying global IDs directly to a local shard, or stepping a non-owner copy can all produce plausible shapes and wrong rows.

TorchRec’s official sharding material uses the analogous input_dist -> lookup -> output_dist decomposition and supports table-wise, row-wise, column-wise, table-row-wise, and data-parallel plans. Its examples use embedding-bag systems, so this article uses TorchRec as a general sharding alternative rather than claiming that its pooled operator is the LM token-table implementation.

Schedule C: Megatron-Core masked local vocabulary lookup

Megatron-Core’s VocabParallelEmbedding is a pinned language-model example of row ownership without routing each ID to one owner. At commit 542ff5368611682f3f8d5487595647c13e02801b, each tensor-parallel rank:

1. derives its contiguous [vocab_start_index,vocab_end_index);
2. marks non-local IDs with input_mask;
3. subtracts vocab_start_index from local IDs and maps non-local IDs to safe local index zero;
4. performs local F.embedding;
5. zeros outputs at non-local positions;
6. reduces the partial outputs across the tensor-parallel group, or optionally reduce-scatters embedding outputs along the sequence-parallel path.

The core trace is:

input_mask = (ids < vocab_start) | (ids >= vocab_end)
masked_input = ids.clone() - vocab_start
masked_input[input_mask] = 0

output_parallel = F.embedding(masked_input, local_weight)
output_parallel[input_mask, :] = 0.0
output = reduce_from_tensor_model_parallel_region(output_parallel)

Because exactly one row shard contributes a nonzero value for each valid ID, the reduction reconstructs the full embedding output. In backward, each local table receives scatter-add contributions only for IDs in its vocabulary range. The tensor-parallel row owner can update its shard without all-reducing that shard’s weight gradient across the same ownership axis. A separate data-parallel group may still reduce replicas of the shard.

This is an official-version-scoped source trace, not the only vocabulary- parallel schedule. The optional reduce-scatter branch changes output placement and the autograd collective pairing; it must be recorded in the manifest.

Schedule D: column-wise table ownership

Partition the feature columns into disjoint intervals [a_r,b_r) that cover [0,d) exactly once. Every column owner needs the IDs, but rank r gathers only:

$$X^{(r)}=E[:,a_r:b_r][i], \qquad X^{(r)}\in\mathbb{R}^{N\times d_r}.$$

If the next operator needs a full hidden row, concatenate the partial outputs:

$$X=\operatorname{Concat}_{feature}(X^{(0)},\ldots,X^{(R-1)}).$$

Backward splits G over the same column intervals and scatter-adds each slice into its owner:

for column_owner, (start, end) in column_shards:
    local_grad = scatter_add(ids, grad_output[:, start:end], shape=(V, end-start))
    column_owner.optimizer.step(local_grad)

Materializing X on one origin can pair forward concatenation/all-gather with a backward split. Replicating a full X and computing losses on several ranks can require reductions of partial-input gradients. Keeping X sharded can avoid a collective only when the next operator accepts that layout. “Column-wise” names weight placement, not a complete activation schedule.

Communication ledger for the toy trace

Let values use b bytes, IDs use b_i, route metadata use b_m, and let N_remote count routed occurrences whose owner differs from their origin. A simple owner-routed logical payload is:

$$M_{row-route} \approx N_{remote}(b_i+b_m)+2N_{remote}db.$$

The two value terms are the forward row response and backward row gradient. For the two-rank trace, N_remote=2, d=3, b=4, and b_i=b_m=8:

ID + metadata requests: 2*(8+8) = 32 bytes
forward row responses:   2*3*4   = 24 bytes
backward row gradients:  2*3*4   = 24 bytes
logical routed payload:             80 bytes

For comparison:

replicated dense DP gradient buffer per rank: V*d*b = 8*3*4 = 96 bytes
Megatron masked-local output reduction buffer: N*d*b = 8*3*4 = 96 bytes
column-wise full-output concatenation payload: N*d*b = 96 bytes

These are tensor payloads, not wire-time models. Collective algorithms, world-size factors, local routes, deduplication, padding, topology, bucket fusion, overlap, and retained route metadata change actual traffic. M4 owns the general communication/performance accounting; this section owns which tensor must move for the embedding contract.

Executable Distributed-Reconciliation Gate

The local NumPy oracle executes the [V=8,d=3] trace under five contracts:

single-table dense scatter-add
replicated local scatter-add followed by SUM
row-owner routed backward
Megatron-style masked local row ownership
column-wise feature ownership followed by concatenation

Every distributed reconstruction must equal the single-table reference exactly for this integer-valued fixture. The oracle separately proves that:

rank-local and global frequency scaling differ
coalesced sparse rows densify to the same input gradient
tied output contribution makes all rows materialized in the toy total gradient
optimizer admission checks gradient layout, device, clipping, and tying

Nineteen negative fixtures reject PAD leakage, out-of-range IDs, malformed upstream rows, row/column shard gaps and overlaps, wrong update owners, unsupported sparse optimizers or devices, dense gradients sent to a sparse-only optimizer, undeclared clipping, a tied dense path under a sparse-only contract, the unsupported sparse-plus-frequency option pair, wrong frequency scope, wrong collective reduction, and missing coalescing.

Passing this gate proves the declared arithmetic and ownership reconstruction. It does not prove NCCL ordering, topology performance, a fused optimizer kernel, or bitwise equivalence under arbitrary floating-point reduction schedules.

0x04B Lookup, Bag Reduction, And Projection Are Different Operators

The word embedding is used for APIs with different shape and semantic contracts. Equal output width does not make these operators interchangeable.

Dense one-hot linear and indexed lookup

A bias-free linear layer over a materialized categorical selector can reproduce lookup. PyTorch stores a linear weight as [d_out,d_in], so set:

$$W_{linear}=E^\top\in\mathbb{R}^{d_e\times V}.$$

Then:

$$\operatorname{Linear}(S)=SW_{linear}^\top=SE.$$

The values match, but the execution interfaces do not:

indexed lookup:
    input  = integer addresses [B,L]
    weight = table [V,d_e]
    output = selected rows [B,L,d_e]

dense one-hot linear:
    input  = floating selectors [B,L,V]
    weight = transposed table [d_e,V]
    output = dense contraction [B,L,d_e]

Adding a bias would break exact lookup equivalence. Feeding raw integer IDs to the linear layer would treat each ID as a scalar feature, which is a different and invalid model.

EmbeddingBag gathers and removes the position axis

nn.EmbeddingBag computes a sum, mean, or coordinate-wise maximum over bags without exposing the full intermediate row tensor. Let offsets delimit bag (j) as the half-open range ([o_j,o_{j+1})). Excluding padding_idx, sum and mean are:

$$z_j^{sum} =\sum_{t=o_j}^{o_{j+1}-1} \mathbb{1}[i_t\ne i_{pad}]E_{i_t},$$ $$z_j^{mean} =\frac{z_j^{sum}} {\sum_{t=o_j}^{o_{j+1}-1}\mathbb{1}[i_t\ne i_{pad}]}.$$

For the two fixture bags:

flat IDs = [1,4,6,2 | 1,5,2,0]
offsets  = [0,4,8]
pad ID   = 0, excluded from reduction

sum:
    bag 0 -> [2.0, 3.0, 1.5]
    bag 1 -> [1.5, 0.5, 3.0]

mean:
    bag 0 -> [0.5000, 0.7500, 0.3750]  # denominator 4
    bag 1 -> [0.5000, 0.1667, 1.0000]  # denominator 3, PAD excluded

The output shape is [num_bags,d_e], not [B,L,d_e]. This may be useful for unordered feature bags, but it cannot replace token lookup before a Transformer when token order and per-position attention are required.

Factorized lookup preserves addresses and positions

Factorized lookup composes two operators:

$$I:[B,L] \xrightarrow{\operatorname{gather}(E)} X_e:[B,L,d_e] \xrightarrow{P} X_0:[B,L,d_{model}].$$

Unlike EmbeddingBag, it preserves every input position. Unlike direct lookup, it adds a dense bridge and allows d_e != d_model. Section 0x0A owns its architecture tradeoffs; this loop establishes that it is not a pooling operation.

A continuous projector has no vocabulary row address

For encoder features:

$$F\in\mathbb{R}^{B\times N\times d_{in}}, \qquad W_p\in\mathbb{R}^{d_{in}\times d_{out}},$$

a count-preserving projector computes:

$$Z_{b,n,:}=F_{b,n,:}W_p+b_p.$$

Every floating input coordinate contributes according to learned weights. No integer selects one row from a vocabulary table. The fixture maps [2,2,2] @ [2,5] -> [2,2,5]; it does not consume input_ids, tokenizer row order, or padding_idx. X1.2 owns how real non-text encoders and projectors are trained and integrated. M2.2 keeps this comparison because the two outputs can both be [B,N,d_model] while naming different construction paths.

Operator	Input object	Persistent weight	Output	Position/count effect	Invalid substitution
indexed lookup	categorical addresses [*]	[V,d_e] table	[* ,d_e]	preserves every address position	pass float features as IDs
one-hot linear	categorical selectors [...,V]	[d_e,V] linear weight	[...,d_e]	preserves positions, materializes V	infer selector width from batch maximum
EmbeddingBag	IDs plus bag boundaries	[V,d_e] table	[num_bags,d_e]	reduces positions inside each bag	use before order-sensitive attention
factorized lookup	categorical addresses [B,L]	[V,d_e] plus [d_e,d_model]	[B,L,d_model]	preserves positions, changes width	call the bridge a second lookup
continuous projector	float feature rows [...,d_in]	affine map to d_out	[...,d_out]	usually preserves row count	feed token IDs as scalar features

Executable Operator-Reconstruction Gate

The local oracle checks all of the following:

[ ] scalar one-hot multiplication equals scalar row selection
[ ] sequence one-hot multiplication equals sequence gather
[ ] batched one-hot contraction equals batched gather exactly
[ ] num_classes comes from the declared address space, not batch max
[ ] IDs use an admitted integer dtype and lie in [0,V)
[ ] table, IDs, bridge, and projector declare shape and element strides
[ ] primitive lookup tensors are co-located or a runtime route is named
[ ] gathered output is materialized separately from persistent table storage
[ ] duplicate IDs create distinct output positions from one logical row
[ ] PAD is physically selected while gradient/visibility policies remain separate
[ ] factorized bridge input width and dtype/device state reconcile
[ ] bag offsets form monotonic half-open ranges and exclude PAD as declared
[ ] continuous projector consumes floating features, never categorical addresses

The frozen report exposes the decisive outcomes directly:

batch_equal:                              true
gather_output_shares_table_storage:       false
duplicate ID 1 output positions:          [[0,0],[1,0]]
PAD output positions:                     [[1,3]]
operator contracts reconstructed:         5
invalid fixtures rejected as declared:    16 / 16

Representative rejection classes are:

Invalid state	Stable oracle error
float IDs or another non-address dtype	INDEX_DTYPE_MISMATCH
ID outside [0,V)	ID_OUT_OF_RANGE
false contiguous-stride declaration	STRIDE_MISMATCH
local lookup split across undeclared devices	DEVICE_MISMATCH
bridge does not consume the table width	BRIDGE_INPUT_WIDTH_MISMATCH
bag offsets do not delimit monotonic half-open ranges	INVALID_BAG_OFFSETS
continuous projector receives categorical addresses	OPERATOR_CONTRACT_MISMATCH

Sixteen invalid fixtures target those contracts. Passing proves this forward construction for the synthetic values. It does not prove a production kernel’s speed, a distributed communication schedule, backward correctness, or learned representation quality.

---

0x05 Table shape, training, PAD row, source/target tables

The embedding table has one row per vocabulary entry:

E: [V, d_model]

A batch of IDs selects rows:

input_ids: [B, L]
E[input_ids]: [B, L, d_model]

The table is usually trainable. It is not a fixed dictionary of meanings. The row for cat becomes useful because the model is trained in contexts where that row helps predict, classify, align, or condition later computations.

There are four table details that cause real implementation bugs.

Vocabulary size

The tokenizer’s largest emitted ID must fit inside the table:

0 <= input_ids < V

If the tokenizer emits an ID outside the embedding table, the lookup fails. If the ID is in range but came from the wrong tokenizer, the lookup succeeds and retrieves the wrong learned row.

Source and target tables

Encoder-decoder code may use separate source and target vocabularies:

src_embed = Embeddings(d_model, src_vocab)
tgt_embed = Embeddings(d_model, tgt_vocab)

The shapes can look identical:

src: [B,S]
tgt: [B,T]

but the address spaces can differ:

E_src: [V_src,d_model]
E_tgt: [V_tgt,d_model]

Swapping source and target embedding tables is a semantic bug, not a rank bug. The integer tensor is still [B,L]; the rows come from the wrong address space.

PAD row

Padding adds rectangular shape:

input_ids:      [B,L_max]
attention_mask: [B,L_max]

The PAD token has an ID, and therefore a potential embedding row:

E[pad_id]: [d_model]

PyTorch’s padding_idx lets a specific row behave as padding in the embedding layer: the padding entry does not contribute gradient through the embedding module, and new embeddings initialize that row as zeros by default.

That is only one part of the contract. The PAD ID also has to agree with attention masks, loss masks, label masks, tokenizer conventions, and any batching code. If one component thinks PAD is 0 and another thinks it is -1, the tensor shape can survive while lookup, masking, or supervision breaks.

A common LLaMA-style tokenizer convention is:

pad_id = -1

That is metadata, not a valid row address for ordinary nn.Embedding lookup. Code must replace, mask, or otherwise handle it before passing IDs into the table.

Output vocabulary

The vocabulary dimension appears again at the output head:

hidden: [B,L,d_model]
logits: [B,L,V]

The input embedding table and output head are related but not automatically the same object. Some implementations tie weights. Some do not. Tying source embeddings, target embeddings, and output projection only makes sense when the vocabulary assumptions are compatible.

---

0x05A Factorized embeddings: vocabulary width need not equal hidden width

The simplest Transformer chooses d_embed = d_model, producing one table:

$$E\in\mathbb{R}^{V\times H}, \qquad N_{direct}=VH.$$

ALBERT makes a different design choice. It first gathers a narrower row of width E_d, then projects it into Transformer hidden width H:

$$E_{vocab}\in\mathbb{R}^{V\times E_d}, \qquad P\in\mathbb{R}^{E_d\times H},$$ $$X_{hidden}=E_{vocab}[I]P,$$ $$N_{factorized}=VE_d+E_dH.$$

The factorization saves parameters when:

$$VE_d+E_dH < VH,$$

or equivalently:

$$E_d < \frac{VH}{V+H}.$$

For V=128,000, H=4096, and E_d=256:

direct table:
    128,000 * 4,096 = 524,288,000 parameters

factorized:
    128,000 * 256 + 256 * 4,096
    = 32,768,000 + 1,048,576
    = 33,816,576 parameters

parameter ratio:
    0.0645, about 15.5x smaller for this bridge

The saving is not free. Every token row passes through P, adding:

$$\operatorname{MAC}_{bridge}=BL\,E_dH.$$

More importantly, all vocabulary rows enter hidden space through an E_d-dimensional bottleneck. ALBERT’s argument is that context-independent lexical identity may not require the full contextual hidden width. This is a modeling hypothesis and architecture choice, not a theorem that every model should factorize.

class FactorizedEmbedding(nn.Module):
    def __init__(self, vocab_size, embedding_dim, hidden_dim, padding_idx=None):
        super().__init__()
        self.lookup = nn.Embedding(
            vocab_size,
            embedding_dim,
            padding_idx=padding_idx,
        )
        self.bridge = nn.Linear(embedding_dim, hidden_dim, bias=False)

    def forward(self, input_ids):
        lexical = self.lookup(input_ids)  # [B,L,E_d]
        hidden = self.bridge(lexical)     # [B,L,H]
        return hidden

Factorization complicates naive weight tying. The input table is [V,E_d], while an ordinary output head expects [V,H]. A model can add an output bridge, tie only compatible factors, or keep the output head independent. “Both are vocabulary matrices” is not enough to alias their storage.

0x05B Initial state may sum several address spaces

BERT-style inputs demonstrate another design: token identity, segment/type, and learned position each select a row of the same width, then are added:

$$X_0 = E_{token}[I] + E_{type}[T] + E_{position}[P].$$

The sum is valid because every term has shape [B,L,H]. Addition preserves row count and hidden width. It also superposes the three sources, so downstream layers do not receive separate token/type/position channels.

The rigorous explanation is linear. Concatenate three one-hot selectors and a block-stacked parameter matrix:

$$\begin{aligned} \mathbf{c} &=[\mathbf{o}_{token};\mathbf{o}_{type};\mathbf{o}_{position}],\\ W &=\begin{bmatrix}E_{token}\\E_{type}\\E_{position}\end{bmatrix}. \end{aligned}$$

Then:

$$\mathbf{c}^{T}W = E_{token}[i]+E_{type}[t]+E_{position}[p].$$

No “frequency band” argument is required. The model learns three row tables in a shared output space because their sum is what the next layer optimizes.

class BertStyleInitialState(nn.Module):
    def __init__(self, vocab_size, type_count, max_positions, hidden_dim):
        super().__init__()
        self.token = nn.Embedding(vocab_size, hidden_dim)
        self.token_type = nn.Embedding(type_count, hidden_dim)
        self.position = nn.Embedding(max_positions, hidden_dim)
        self.norm = nn.LayerNorm(hidden_dim)
        self.dropout = nn.Dropout(0.1)

    def forward(self, input_ids, token_type_ids, position_ids):
        x = (
            self.token(input_ids)
            + self.token_type(token_type_ids)
            + self.position(position_ids)
        )
        return self.dropout(self.norm(x))

This is a historical architecture example, not M2.2 taking ownership of position policy. M3.3 owns how position_ids are constructed and how modern alternatives inject geometry.

0x05C Training lifecycle of the input table

The table passes through distinct states:

1. allocate rows according to tokenizer vocabulary
2. initialize or import row values
3. train jointly with the model objective
4. optionally freeze, factorize, tie, shard, or quantize
5. load with exact tokenizer/address compatibility
6. serve gathers under runtime dtype/placement

Random versus pretrained initialization

Random initialization is natural when the tokenizer and model are trained together. Importing Word2Vec/GloVe/FastText rows is possible only when token and dimension contracts can be reconciled:

source embedding vocabulary -> model tokenizer vocabulary
source word-level rows       -> model word/subword rows
source width d_source        -> model width d_embed/d_model
source normalization         -> model tokenizer normalization

Exact vocabulary matches can copy rows. Partial matches need an initialization policy for missing rows. Word-level vectors do not directly define subword pieces; averaging characters or constituent rows is a heuristic.

def initialize_from_external(model_embedding, tokenizer, external):
    copied = 0
    with torch.no_grad():
        for token, token_id in tokenizer.vocab.items():
            if token in external and external[token].shape[-1] == model_embedding.embedding_dim:
                model_embedding.weight[token_id].copy_(external[token])
                copied += 1
    return copied

Report coverage and whether copied rows remain trainable. “Initialized from pretrained embeddings” without those details is not reproducible.

Row frequency and update opportunity

For input-side lookup, row v participates when ID v occurs. If corpus count is c_v, then the number of direct lookup-gradient contributions scales with c_v. High-frequency rows can dominate optimizer steps; rare rows may remain near initialization.

0x04A derives the exact local-count rule and shows why rank-local frequency scaling followed by data-parallel reduction is not the same as one global effective-batch count. This lifecycle section adds the longer horizon: data sampling, tokenization, packing, gradient accumulation, world size, and output tying all affect how often a row receives an update.

Subsampling frequent words in classic Word2Vec and frequency-scaled embedding gradients are related responses to frequency imbalance, but they intervene at different layers: one changes training examples; the other changes gradient aggregation.

Freeze versus joint training

Freezing the input table can preserve imported geometry or reduce trainable state. Joint training lets rows specialize to the model objective. The tradeoff is not simply quality versus speed:

freeze:
    stable address vectors, no table optimizer state, but cannot adapt rare/new
    pieces or align scale to a new architecture

train:
    task/model adaptation and vocabulary specialization, but full optimizer and
    checkpoint state plus potential catastrophic drift from imported geometry

In most modern language-model pretraining, the table is trained jointly. In transfer setups, freezing schedules and learning-rate multipliers are explicit hyperparameters.

Optimizer groups and weight decay

Embedding parameters can use a separate group:

optimizer = AdamW([
    {"params": dense_blocks, "lr": lr, "weight_decay": wd},
    {"params": [model.token_embedding.weight],
     "lr": embedding_lr,
     "weight_decay": embedding_wd},
])

Reasons include sparse-gradient compatibility, PAD-row invariants, pretrained initialization, and different effective update frequency. There is no universal rule that embeddings should or should not receive weight decay; document the checkpoint’s optimizer contract.

0x05D Vocabulary Growth Is A Model-State Migration

Owner question. When a tokenizer gains, loses, or reorders addresses, which model tensors, optimizer states, distributed shards, release artifacts, and runtime routes must change together so that preserved rows remain preserved and new rows begin from a declared state?

“Add tokens and resize the embedding” describes only one local API call. A release migration crosses at least four state spaces:

symbolic state:
    ordered vocabulary, normalizer, pre-tokenizer, special IDs, template

model state:
    input rows, output rows, output bias, aliases, factorization bridges

training state:
    gradients, master weights, Adam moments, step counters, scheduler state

runtime state:
    tensor shards, quantized packs, tokenizer/model routing, caches, rollback tuple

M2.1 owns how a candidate vocabulary is fitted, evaluated, and assigned addresses. This section consumes two immutable ordered vocabularies and owns the row-level conversion between their model states. M3.8 owns the complete output distribution and decoding policy. M8 owns quantization algorithms and the generic canary machinery. Here, output rows, quantized packs, and canaries appear only where vocabulary identity makes them part of the embedding migration contract.

Start by classifying the address change

Equal old and new vocabulary sizes do not imply compatibility. A migration manifest must classify the transformation before any tensor is allocated:

Change	Address effect	Rows mechanically reusable?	Required semantic proof
append-only	old ordered vocabulary is an exact prefix of the new one	every old row at the same ID	old normalization, pieces, special roles, and template are unchanged
permutation/reindex	pieces may be identical but IDs move	yes, through an explicit permutation	every old piece and control role maps to exactly one new address
removal/pruning	some old addresses disappear	surviving rows through an injective map	no retained dataset/session requires removed IDs
split	one old unit becomes several new units	no exact semantic row exists for the new pieces	initializer and adaptation policy are declared
merge	several old units become one new unit	no unique exact source row exists	composition initializer and retraining evidence are declared
replacement/joint vocabulary	many pieces and IDs change	overlap rows only	complete old-to-new overlap map plus initialization for the complement

Append-only is therefore a property to prove, not a synonym for “the new file contains the old strings.” The proof is:

V_new[0:V_old] == V_old
special_role_to_id_new == special_role_to_id_old
tokenizer_semantics_hash_new extends, rather than rewrites, the old contract

If one old ID moves, raw encoded datasets and persisted ID sequences no longer have the same meaning even though all tensor dimensions are valid.

The migration map separates exact copies from initialization

Let the old and new ordered vocabularies have sizes $V_0$ and $V_1$. Define an exact-row map

$$R\in\{0,1\}^{V_1\times V_0},$$

where $R_{v,j}=1$ means new address $v$ has exactly the same symbolic meaning as old address $j$. For an exact migration, each new row has at most one exact source and each preserved old row has at most one destination:

$$\sum_j R_{v,j}\le 1, \qquad \sum_v R_{v,j}\le 1.$$

Define the preserved-row mask

$$m_v=\sum_j R_{v,j}\in\{0,1\}.$$

For an old lexical table $E_0\in\mathbb{R}^{V_0\times D}$ and explicit initializer rows $E_{init}\in\mathbb{R}^{V_1\times D}$, the new table is

$$E_1 = R E_0 + \operatorname{diag}(\mathbf{1}-m)E_{init}.$$

Row-wise, this says exactly what the compact matrix notation can hide:

for new_id in range(V1):
    if exact_source[new_id] is not None:
        E1[new_id] = E0[exact_source[new_id]]
        row_status[new_id] = "exact_copy"
    else:
        E1[new_id] = initialize_new_row(new_id, policy, evidence)
        row_status[new_id] = "initialized"

A split, merge, decomposition average, FOCUS combination, or semantic-neighbor copy belongs to initialize_new_row. It may be a good warm start, but it is not an exact preservation edge in $R$. Keeping those ledgers separate prevents an approximate semantic guess from passing an old-row equality test.

The same exact map applies to every vocabulary-indexed state:

$$W_{out,1}=R W_{out,0}+\operatorname{diag}(\mathbf1-m)W_{out,init},$$ $$b_1=R b_0+(\mathbf1-m)\odot b_{init}.$$

For tied weights, $E_1$ and $W_{out,1}$ are not two independently converted copies. Convert the one canonical lexical parameter, then restore the alias or distributed replica contract from 0x0A.

Preserved logits do not preserve normalized probabilities

Suppose all old output rows and the context vector $h$ are unchanged. Their raw logits $z_i=h^T e_i+b_i$ are then exactly preserved. Before expansion, the softmax denominator is

$$Z=\sum_{i=1}^{V_0}\exp(z_i).$$

If $M=V_1-V_0$ new rows produce logits $a_1,\ldots,a_M$, their added mass is

$$A=\sum_{j=1}^{M}\exp(a_j).$$

Every old-token probability changes from $p_i=\exp(z_i)/Z$ to

$$p_i' = \frac{\exp(z_i)}{Z+A} = p_i\frac{Z}{Z+A}.$$

The total probability assigned to new tokens is

$$q_{new}=\frac{A}{Z+A},$$

and the KL divergence from the old distribution to the old coordinates of the expanded distribution is

$$D_{KL}(P_0\Vert P_1) = \sum_i p_i\log\frac{p_i}{p_i'} = \log\left(1+\frac{A}{Z}\right).$$

This derivation explains why shape-preserving row copies are not a behavioral no-op. The new rows enter the denominator immediately, before they have appeared in a prompt or target.

Hewitt’s vocabulary-expansion analysis derives the single-new-row form and motivates mean initialization. If every new output row is initialized to the mean old row

$$\mu=\frac{1}{V_0}\sum_{i=1}^{V_0}e_i,$$

and the bias convention is compatible, then its logit is the mean old logit. By Jensen’s inequality,

$$\exp\left(\frac{1}{V_0}\sum_i z_i\right) \le \frac{1}{V_0}\sum_i\exp(z_i) = \frac{Z}{V_0}.$$

For $M$ identical mean-initialized rows,

$$A\le\frac{MZ}{V_0}, \qquad D_{KL}(P_0\Vert P_1) \le \log\left(1+\frac{M}{V_0}\right).$$

The bound does not automatically survive covariance noise, different output biases, untied input/output initialization, a changed hidden state, or a full tokenizer replacement. It is a local initialization argument, not a promise of downstream quality. A zero row is also not neutral: its dot-product logit is zero, so it contributes exp(0)=1; if the trained old logits are all strongly negative, that mass can dominate.

Initialization is a design choice, not checkpoint preservation

Policy for a non-preserved row	Information used	Strength	Failure to test
constructor-matched random	original model initializer	reproduces fresh-allocation convention	new output mass can be too large or row norm mismatched
zero	no lexical prior	deterministic and easy to audit	zero logit can dominate a negative-logit model; symmetry among new rows
old mean or mean/covariance	global old-table geometry	bounded denominator disturbance under the stated mean assumptions	covariance noise weakens the bound; global mean may ignore token semantics
decomposition average	old tokenizer expansion of the new string	reuses model-native pieces	several positions are collapsed into one static row; order is usually lost
overlap/semantic combination	shared pieces plus an auxiliary embedding space	useful for tokenizer/language replacement	auxiliary geometry and overlap coverage become additional artifacts
output-bias prior	desired initial token prior	directly controls denominator mass	an overly negative bias can make new rows unlearnably rare
exact row remap	identity proved by old/new address manifests	preserves one trained row exactly	invalid when normalization, boundary, or control semantics changed

FOCUS is one externally verified alternative for vocabulary specialization: it represents target tokens using combinations of overlapping source/target tokens selected with semantic similarity in an auxiliary static embedding space. Its reported comparisons are specific to the paper’s multilingual-specialization setup; they do not make FOCUS a universal initializer. WECHSEL is another cross-lingual alternative. Both illustrate why initializer_id must bind code, auxiliary model, overlap set, and hyperparameters rather than say “semantic initialization.”

Current Hugging Face Transformers documentation exposes resize_token_embeddings(..., mean_resizing=True). The pinned implementation at commit 85c307a184264643392ef37b6e902209812a6e23 allocates new input rows, initializes added rows from old mean/covariance when requested, separately resizes an untied language-model head, and calls tie_weights() after resizing. This is a useful implementation boundary, not a complete optimizer, dataset, shard, or release migration.

Architecture determines which vocabulary axes move

The converter begins from the architecture manifest, not key-name guesses:

Architecture	Vocabulary-indexed tensors to migrate	State that normally remains unchanged
direct tied decoder	one canonical [V,H] table, output bias [V]	Transformer trunk
direct untied decoder	input [V,H], output [V,H], bias [V]	Transformer trunk
factorized lexical tie	lexical [V,D], bias [V]	input bridge [D,H], output bridge [H,D] unless architecture changes
factorized input, direct output	lexical [V,D], direct output [V,H], bias [V]	input bridge [D,H]
seq2seq target-only tie	target lexical/output table and target bias	source table when source vocabulary is unchanged
compatible three-way tie	one joint source/target/output table and target bias	trunk; only if all three roles share the same ordered map

Resizing the bridge in a factorized model because V changed is a category error: the bridge does not have a vocabulary axis. Conversely, updating input rows while leaving an untied output head at $V_0$ produces a model that can read new IDs but cannot score them. Special IDs require their own semantic check. If PAD, BOS, EOS, or control IDs move, masks, templates, ignore-index rules, and runtime configuration migrate with them.

def migrate_vocabulary_state(old, spec):
    # Convert canonical logical tensors before creating physical aliases/shards.
    lexical = migrate_rows(old.lexical, spec.exact_map, spec.lexical_initializer)

    if spec.output_is_tied:
        output = lexical
    else:
        output = migrate_rows(old.output, spec.exact_map, spec.output_initializer)

    bias = migrate_vector(old.output_bias, spec.exact_map, spec.bias_initializer)

    model = construct_architecture(spec.new_config)
    model.load_canonical(lexical=lexical, output=output, output_bias=bias)
    model.restore_declared_aliases()
    assert_model_address_contract(model, spec)
    return model

Adam state migrates by row; its age usually does not

For Adam or AdamW, a vocabulary table commonly has first and second moments with the same logical shape as the parameter:

$$M_0,V_0\in\mathbb{R}^{V_0\times D}.$$

Under a zero-moment policy for new rows, migrate them as

$$M_1=R M_0, \qquad V_1^{(2)}=R V_0^{(2)}.$$

The notation $V_1^{(2)}$ denotes the second moment, not vocabulary size. For an output bias, apply the same rule to vectors of length $V$. Exact old rows retain their exact moment rows. New rows begin with zero moments unless the migration manifest declares and validates a different warm start.

At optimizer step $t$, ignoring weight decay for the moment,

$$\begin{aligned} m_t &= \beta_1m_{t-1}+(1-\beta_1)g_t,\\ v_t &= \beta_2v_{t-1}+(1-\beta_2)g_t^2,\\ \widehat m_t &= \frac{m_t}{1-\beta_1^t},\\ \widehat v_t &= \frac{v_t}{1-\beta_2^t},\\ \theta_t &= \theta_{t-1}-\eta\frac{\widehat m_t}{\sqrt{\widehat v_t}+\epsilon}. \end{aligned}$$

Many stock optimizers store one scalar step for the whole parameter, not one age per row. Preserving the old scalar step preserves old-row bias correction, but new rows inherit that large $t$ despite having zero moments. Resetting the whole parameter step makes new rows young by damaging the update convention for every old row. Per-row age requires a custom optimizer/state representation or separating new rows into a distinct parameter and reconciling that choice with lookup, tying, sharding, and checkpoint code.

def migrate_adam_table_state(old_state, exact_map, new_shape):
    exp_avg = zeros(new_shape, dtype=old_state.exp_avg.dtype)
    exp_avg_sq = zeros(new_shape, dtype=old_state.exp_avg_sq.dtype)
    for new_id, old_id in exact_map.items():
        exp_avg[new_id] = old_state.exp_avg[old_id]
        exp_avg_sq[new_id] = old_state.exp_avg_sq[old_id]
    return AdamState(
        exp_avg=exp_avg,
        exp_avg_sq=exp_avg_sq,
        step=old_state.step,  # parameter-level age, not a per-row claim
    )

PyTorch documents that optimizer state is keyed by parameter IDs and that load matches saved IDs to live parameters by zipping parameter-group order, without additional verification. Parameter names do not change this default behavior. If resized parameters differ from those in the saved optimizer, a custom register_load_state_dict_pre_hook must adapt the state dictionary before load. The official Optimizer.load_state_dict documentation also requires scheduler construction before optimizer-state load so loaded learning rates are not overwritten.

The safe order is therefore:

1. construct the new model architecture
2. load/migrate canonical model tensors
3. restore tied aliases or declared replicas
4. construct optimizer over the final live Parameters
5. migrate optimizer state by canonical parameter identity and row map
6. construct/load scheduler in the framework-documented order
7. execute one deterministic reconciliation step

Loading optimizer state before re-tying can leave state attached to a stale parameter object. Registering both endpoints of a tie creates two moment sets for one logical parameter.

”Freeze old rows” must close every update path

A continued-training schedule may choose:

full continued pretraining
new rows only
new rows plus output head
new rows plus adapters
row-specific warmup followed by joint training

For a tied decoder, zeroing only the sparse input-lookup gradient does not freeze old rows. The output softmax generally contributes a dense gradient over the full vocabulary. Adam moments can also carry an old update forward, and AdamW’s decoupled weight decay can move a row even when its current gradient is zero. A genuine old-row freeze must mask the combined gradient, prevent moment-driven updates and decay for those rows, and remain valid under shard and alias ownership. A separate new-row parameter can simplify optimizer age but complicates the classifier and checkpoint topology; there is no free schedule.

Convert the logical state, then prove physical placement

Vocabulary migration and resharding are two transformations:

logical conversion:
    old ordered rows + old optimizer state -> new ordered rows + new state

physical materialization:
    new logical tensors -> TP/PP/DP shards, replicas, optimizer fragments, packs

The simplest auditable schedule converts a canonical logical tensor and then partitions it. A streaming converter is acceptable only if concatenating its outputs reproduces the same canonical result.

For tensor-parallel vocabulary shards, new intervals must form an exact cover:

$$[0,V_1)=\biguplus_{r=0}^{P-1}[s_r,e_r).$$

Tied input/output roles must use the same row intervals. Output bias follows the same vocabulary partition. Pipeline replicas must receive the same migrated row shard, logical parameter ID, moments, step convention, and gradient/update-ownership contract described in 0x0A.

Distributed optimizer formats add another compatibility axis. PyTorch’s distributed checkpoint state-dictionary APIs canonicalize optimizer parameter IDs to fully qualified names and can return state that is reshardable across supported parallelisms, but the API is documented as experimental. Megatron Core’s current distributed-checkpoint documentation distinguishes a faster dp_reshardable optimizer format, which does not permit arbitrary model-parallel changes, from a slower fully_reshardable format. It also documents version-specific legacy optimizer compatibility limits. Record the framework version and checkpoint format; do not infer reshardability from the word “distributed.”

Quantized or packed embedding/output tensors are derived runtime artifacts. The conservative boundary is:

migrate canonical high-precision logical state
    -> restore aliases/replicas
        -> recalibrate or regenerate quantized packs and group metadata

Directly appending packed bytes is valid only when group boundaries, scale and zero-point state, padding, alignment, row order, and runtime loader contract are all explicitly migrated. M8 owns those derivations.

Tokenizer and model revisions form a version-pair matrix

Let T0/M0 be the active tokenizer/model pair and T1/M1 the candidate:

Request tokenizer	Loaded model	Admission	Reason
T0	M0	control	original ordered addresses and rows agree
T0	M1	compatibility test only	append-only can preserve old IDs, but continued training may still change behavior
T1	M0	reject	T1 may emit IDs outside M0 or attach changed meaning to an old ID
T1	M1	candidate	only valid after manifest, tensor, optimizer/runtime, and behavior gates pass

An API receiving text can route the text to the tokenizer paired with the model. An API receiving pre-tokenized IDs must require the tokenizer semantic revision and reject an incompatible pair. Persisted IDs without tokenizer provenance are not safely replayable after migration.

A releasable bundle binds:

old and new ordered-vocabulary digests
old-to-new exact row map and initializer manifest
tokenizer semantics and raw release-file digests
model config and canonical tensor digests
input/output/bias role map and alias graph
optimizer format, parameter identity map, moments, step, scheduler state
TP/PP/DP shard and replica map
quantized-pack/calibration revision, when deployed
encoded-dataset compatibility and retokenization decision
runtime/model/tokenizer route identity
golden cases, offline report, canary record, rollback target

The tokenizer and model must advance atomically as one routing tuple. This is the representation-specific constraint that M8.10’s generic release machinery must preserve.

Lifecycle: conversion evidence precedes traffic

PROPOSED_DIFF
    -> LOGICAL_STATE_CONVERTED
        -> SHARDS_AND_PACKS_MATERIALIZED
            -> OFFLINE_MECHANICAL_AUDIT
                -> ONE_STEP_TRAINING_RECONCILIATION
                    -> SHADOW_VERSION_PAIR
                        -> CANARY_VERSION_PAIR
                            -> RAMP
                                -> ACTIVE

Every live stage retains a rollback edge to the complete T0/M0 tuple. Keep the old tokenizer, model, runtime, configuration, and required capacity available. Do not send a session from T1/M1 back to T0/M0 while reusing foreign token IDs, prefix-cache entries, or K/V rows. Reconstruct from source text or another explicitly compatible representation. M5 owns cache mechanics; M8.10 owns generic traffic control and recovery-time objectives.

Mechanical and behavioral gates are different

The migration cannot be approved by one perplexity number. It needs a ladder:

address gate:
    old/new vocabulary digests, exact-map uniqueness, special-ID semantics

tensor gate:
    preserved rows bitwise equal; new rows finite; expected shapes/dtypes

classifier gate:
    old raw logits preserved when required; new probability mass measured

optimizer gate:
    old moments/step preserved; new-state policy explicit; one-step oracle passes

ownership gate:
    aliases restored; TP cover exact; PP replicas/state/update contract valid

artifact gate:
    shard and quantized-pack digests bind the candidate tuple

behavior gate:
    old-text regressions, new-domain gains, special/control behavior, generation

release gate:
    T1/M0 rejected; T1/M1 admitted only with a warm, loadable T0/M0 rollback

The local vocabulary-migration oracle executes a synthetic append-only migration, derives the softmax mass and KL terms, migrates direct/factorized/tied tensor roles and Adam states, reconciles one deterministic update, reconstructs tensor-parallel shards, and evaluates the four tokenizer/model pairs. Its negative fixtures reject shape-correct but semantically stale maps, moved special IDs, corrupted old rows or moments, reset optimizer age, duplicate tied state, factorized bridge resize, shard gaps/overlaps, stale packs, unsafe version pairs, missing rollback artifacts, and unversioned persisted IDs.

Passing it proves the declared synthetic transformation. It does not prove tokenizer quality, production framework loading, collective ordering, quantized-kernel parity, downstream model quality, or canary safety.

Checkpoint serialization contract

Record at least:

tokenizer semantic/release hash and ordered vocabulary digest
embedding shape, dtype, and factorization dimensions
PAD/control rows and scale convention
weight-tying aliases or distributed replica topology
source/target sharing policy and exact migration map
optimizer parameter identities, moment shapes, step and scheduler convention
sharding layout and distributed-checkpoint format/version
quantization representation and calibration revision, if deployed
candidate version pair and complete rollback version pair

Some serialization formats do not preserve Python object aliases directly; load code may need to re-tie parameters after restoring values. Test pointer identity for colocated ties and logical replica/update identity for distributed ties, not only value equality.

---

0x06 Positional handoff: src_embed and tgt_embed are sequential

In the Harvard construction, the embedding module is wrapped with positional encoding:

src_embed = nn.Sequential(
    Embeddings(d_model, src_vocab),
    PositionalEncoding(d_model, dropout),
)

tgt_embed = nn.Sequential(
    Embeddings(d_model, tgt_vocab),
    PositionalEncoding(d_model, dropout),
)

So these are different operations:

Embeddings(d_model, vocab)(x)

and:

src_embed(x)

The first performs content lookup and scaling. The second performs content lookup, scaling, positional addition, and positional dropout in the Harvard-style code.

M2.2 owns the embedding half:

input_ids -> learned rows -> [B,L,d_model]

M3.3 owns the first positional-geometry handoff:

[B,L,d_model] + position/layout/time information

The boundary matters because the shape is unchanged:

after embedding:           [B,L,d_model]
after positional encoding: [B,L,d_model]

A model can receive the correct rank while losing position information, receiving position information twice, or applying it at the wrong interface.

That is why “the embedding output has shape [B,L,d_model]” is necessary but not sufficient. It says the width is correct. It does not prove the row has the right content row, position row, mask behavior, or source/target table.

---

0x07 Published X1 handoff: H_media before a projector

The following three sections are retained because the existing article already contains a useful implementation boundary and figures. They are not canonical M2.2 depth. X1.1 owns the choice of non-text units and geometry; X1.2 owns encoders, projectors, resamplers, and fusion; X1.3 owns multimodal packing, training, and runtime state.

For non-text inputs, this preview starts after a processor or encoder has produced modality-specific feature states.

Examples:

H_img:   [B, N_img,   d_vision]
H_doc:   [B, N_doc,   d_doc]
H_audio: [B, N_audio, d_audio]
H_video: [B, N_video, d_video]
H_event: [B, N_event, d_event]

These are not token IDs. They are already continuous features.

That means this is usually wrong:

bad_ids = image_features.long()
text_embedding(bad_ids)

It may produce integer tensors. It destroys the meaning of the visual features.

The local problem is instead a width and interface problem:

vision encoder output width: d_vision
language model state width:  d_model

For example:

H_img: [B,N_img,1024]
LLM expects: [B,N,d_model=4096]

The projector bridge must convert:

[B,N_img,1024] -> [B,N_img,4096]

The projector is not just “making the shape fit.” It is the learned alignment layer between the feature space produced by one component and the state space consumed by another.

---

0x08 Published X1 handoff: projectors and adapters

The simplest count-preserving projector is a linear layer:

class MediaProjector(nn.Module):
    def __init__(self, d_media, d_model):
        super().__init__()
        self.proj = nn.Linear(d_media, d_model)

    def forward(self, h):
        return self.proj(h)

The shape trace is:

h:          [B,n_media,d_media]
proj.weight implementation storage: [d_model,d_media]
output:     [B,n_media,d_model]

Conceptually:

$$Z_{\text{media}} = H_{\text{media}} W_p + b$$

where:

$$W_p \in \mathbb{R}^{d_{\text{media}} \times d_{\text{model}}}$$

and:

$$Z_{\text{media}} \in \mathbb{R}^{B \times n_{\text{media}} \times d_{\text{model}}}$$

The count is preserved:

n_media in
n_media out

This covers the common “project visual features into language embedding width” interface. LLaVA-style systems use this kind of connector idea: visual features from a vision encoder are mapped into the language model’s embedding space so they can be consumed as language-compatible rows.

Count-changing adapters

Some adapters change both width and row count:

H_media: [B,n_media,d_media]
    -> adapter / resampler / query bottleneck
Z_media: [B,r,d_model]

Here:

r != n_media

That is an interface choice, not a small implementation detail. It changes how many rows later attention sees. A query bottleneck can compress hundreds or thousands of visual patch features into a fixed number of learned-query outputs. BLIP-2’s Q-Former is a source-backed example of this kind of bridge between a frozen image encoder and a frozen language model.

The retained published handoff only needs:

Adapter	X1 treatment
linear projector	preserves row count, changes width
MLP projector	preserves row count, changes width with more capacity
pooling/resampler	changes row count and width
query bottleneck	extracts r rows from a larger feature set

The downstream risk is immediate. If an adapter changes:

[B,n,d] -> [B,r,d_model]

then masks, row metadata, labels, and geometry tensors must follow r, not the old n.

---

0x09 Published X1 handoff: mixed-origin X0

Once all branches have width d_model, they can be assembled along the sequence or unit axis:

X0 = torch.cat(
    [
        image_states,   # [B, N_img, d_model]
        text_states,    # [B, L_text, d_model]
        event_states,   # [B, N_event, d_model]
    ],
    dim=1,
)

The result is:

X0: [B, N_total, d_model]
N_total = N_img + L_text + N_event

The concatenation axis is the row axis. Concatenating along the feature axis would produce a different object:

wrong:
    [B,n,d_model] + [B,n,d_model] -> [B,n,2*d_model]

That may be useful in some architectures, but it is not the usual “append media rows before text rows” context construction.

Mixed-origin state requires more than equal width.

Invariant	Why it matters
row order	prompt templates and training data expect particular row order
row type	image, text, event, PAD, and tool rows may need different conventions
masks	ragged media and text batches need correct visibility/padding masks
geometry metadata	X1 aligns layout/time metadata; M3.3-M3.5 own language position geometry
labels	M3.9 decides which rows are supervised and which are ignored

This bounded handoff stops at width alignment and state assembly. X1.1 owns non-text geometry, X1.2 owns mixed-origin fusion, and X1.3 owns multimodal packing and runtime state. M3.9 owns general training labels and masks; M3.10 owns single-request inference behavior.

The debug question is:

Does row 17 still correspond to the intended token, patch, span, event, or PAD row?

If the answer is no, [B,N_total,d_model] is only a convincing shape.

---

0x0A Weight tying is an algebra, storage, and ownership contract

Embedding lookup and the vocabulary classifier happen at opposite boundaries of the model:

input:
    token ID v -> row E[v,:] -> hidden-width state

output:
    hidden-width state -> one score per vocabulary row

Their framework layouts can match:

token_embedding.weight: [V,H]
lm_head.weight:         [V,H]

but equal shape is only the first admission check. The full contract has four independent questions:

semantic identity:
    does row v name the same token in every role?

algebraic compatibility:
    does each role consume the same [V,H] parameter orientation?

physical realization:
    is there one Parameter, or declared synchronized replicas?

update ownership:
    where are all gradient contributions combined, optimizer state stored,
    and the authoritative update applied?

Owner question. When input lookup and output classification claim to use one vocabulary table, what exact algebra, object identity, shard layout, checkpoint graph, and optimizer schedule make that statement true before and after serialization?

M2.2 owns those identities. M3.8 owns the complete output-head score interface, cross-entropy boundary, and decoding policy. M4 owns general collective cost models. M8 owns quantization derivations and packed kernels. This section keeps only the consequences needed to prove that the shared table remains one logical parameter.

The papers do not make one interchangeable claim

Weight tying entered language modeling through related but distinct arguments. Collapsing them into “sharing saves parameters” loses the actual design space.

Source	Question and intervention	Important scope
Press and Wolf	Treat the topmost classifier matrix as an output embedding, analyze its row updates, and set input and output embeddings equal. They also study decoder input/output tying and three-way translation tying.	Their update analysis distinguishes neural language models from skip-gram; the translation variants require compatible source/target subword identity. Results are source-reported, not a universal tying theorem.
Inan, Khosravi, and Socher	Start from an augmented target distribution informed by word-vector similarity, show why the classifier and input-embedding subspaces approach one another under stated assumptions, then evaluate explicit embedding reuse.	The theoretical section assumes equal input/hidden widths and removes output bias for the derivation. Augmented loss and hard parameter reuse are separable interventions.
Attention Is All You Need	Share compatible encoder embedding, decoder embedding, and pre-softmax weights; multiply embedding-role outputs by $\sqrt{d_{model}}$.	This is three-role reuse under the paper’s vocabulary setup. The role-specific input scale does not create another parameter copy.
ALBERT	Factor vocabulary width away from hidden width, then use a hidden-to-embedding decoder path before vocabulary classification.	Direct [V,H] tying is impossible when lexical rows are [V,D] with $D\ne H$; only the compatible lexical classifier factor can be shared.

Press and Wolf’s neural-language-model observation is especially useful for understanding the gradient: an input-only table updates rows selected by the current batch, whereas the ordinary output softmax gives every vocabulary row an update. Their tied row therefore evolves under both roles. The paper also warns, through its skip-gram analysis, that the same tying argument should not be transported into every embedding objective without re-derivation.

The original Transformer-era diagram is a compact intuition: one stored matrix is read by lookup and transposed by the output linear operation.

The diagram does not establish vocabulary identity, pointer identity, factorization compatibility, or distributed synchronization. Those are the rest of this section.

Deriving the two-path gradient from one shared parameter

Flatten all supervised output positions into $N$ rows. Let:

S: [M,V]       one-hot selector for M input positions
E: [V,H]       shared token table
G_X: [M,H]     gradient arriving at gathered input states
C: [N,H]       final contextual rows presented to the classifier
Y: [N,V]       target distributions, usually one-hot
b: [V]         optional untied output bias

The input and output uses are:

$$X=SE, \qquad Z=CE^T+\mathbf{1}b^T, \qquad P=\operatorname{softmax}(Z).$$

For row-wise cross entropy,

$$\mathcal L=-\sum_{n=1}^{N}\sum_{v=1}^{V}Y_{n,v}\log P_{n,v}.$$

The log-softmax derivative gives the logit residual:

$$D = \frac{\partial\mathcal L}{\partial Z} =P-Y, \qquad D\in\mathbb{R}^{N\times V}.$$

Now differentiate the classifier $Z=CE^T$. Using $d\mathcal L=\operatorname{tr}(D^T dZ)$:

$$\begin{aligned} dZ &=dC\,E^T+C\,dE^T,\\ d\mathcal L &=\operatorname{tr}(D^T dC\,E^T) +\operatorname{tr}(D^T C\,dE^T). \end{aligned}$$

Reordering each trace to place the differential last yields:

$$\boxed{ \left.\frac{\partial\mathcal L}{\partial E}\right|_{output}=D^TC }, \qquad \boxed{ \frac{\partial\mathcal L}{\partial C}=DE }, \qquad \boxed{ \frac{\partial\mathcal L}{\partial b}=D^T\mathbf1 }.$$

The backward section derived the transpose of lookup:

$$\left.\frac{\partial\mathcal L}{\partial E}\right|_{input}=S^TG_X.$$

Because both graph edges terminate at the same parameter object, autograd must accumulate:

$$\boxed{ \frac{\partial\mathcal L}{\partial E} =S^TG_X+D^TC }.$$

The row-level form shows exactly why tying destroys input-only sparsity:

$$\frac{\partial\mathcal L}{\partial E_v} = \sum_{m:i_m=v}G_{X,m} + \sum_{n=1}^{N}(P_{n,v}-Y_{n,v})C_n.$$

The first sum is empty for vocabulary rows absent from the input. The second sum is generally nonzero for every row because softmax assigns probability to every finite candidate. An output bias remains a separate [V] parameter and does not become tied merely because the classifier weight does.

If the input role multiplies gathered rows by a scalar $s$, as in the original Transformer’s $s=\sqrt{H}$ convention, the input contribution becomes $sS^TG_X$. The output contribution remains $D^TC$ unless the output operation declares its own scale. One parameter can participate in differently scaled operations without becoming two parameters.

A two-position tensor trace

Use V=4, H=3, input IDs [1,3,1], and the following upstream input rows:

G_X = [[1,0,1],
       [0,2,1],
       [2,1,0]]

The input contribution is:

row 0: [0,0,0]
row 1: [3,1,1]
row 2: [0,0,0]
row 3: [0,2,1]

For two classifier rows, choose:

C = [[1,2,-1],
     [0,1, 2]]

D = P-Y
  = [[ 0.1,-0.4, 0.2, 0.1],
     [-0.2, 0.1,-0.3, 0.4]]

Each row of D sums to zero, as a softmax-cross-entropy residual should. The output contribution D.T @ C is:

row 0: [ 0.1, 0.0,-0.5]
row 1: [-0.4,-0.7, 0.6]
row 2: [ 0.2, 0.1,-0.8]
row 3: [ 0.1, 0.6, 0.7]

The tied total is therefore:

row 0: [0.1,0.0,-0.5]   # absent from input, still updated by output
row 1: [2.6,0.3, 1.6]
row 2: [0.2,0.1,-0.8]   # absent from input, still updated by output
row 3: [0.1,2.6, 1.7]

This trace is also an optimizer admission test: a sparse-only update path cannot consume the complete tied gradient without a declared densification or another mathematically equivalent classifier implementation.

Direct, factorized, and incompatible tying are different architectures

For a direct table $E\in\mathbb{R}^{V\times H}$:

untied input plus output weights:  2*V*H
tied input/output weight:             V*H

The familiar saving is one matrix:

$$\Delta N=VH,$$

excluding optional output bias and optimizer state.

Now let lexical width be $D\ne H$:

$$A\in\mathbb{R}^{V\times D}, \qquad P_{in}\in\mathbb{R}^{D\times H}, \qquad X=SA P_{in}.$$

An ordinary direct output matrix is [V,H] and cannot alias A. A compatible factorized classifier first maps contextual rows back to lexical width:

$$U=C R_{out}, \qquad R_{out}\in\mathbb{R}^{H\times D}, \qquad Z=UA^T+\mathbf1b^T.$$

The lexical table then receives:

$$\left.\frac{\partial\mathcal L}{\partial A}\right|_{input} =S^TG_XP_{in}^T,$$ $$\left.\frac{\partial\mathcal L}{\partial A}\right|_{output} =D^TU,$$

and the output bridge receives:

$$\frac{\partial\mathcal L}{\partial R_{out}} =C^TDA.$$

This is the shape pattern used by ALBERT’s masked-LM head: hidden states are mapped from hidden width to embedding width before the decoder whose [V,D] weight is tied to the lexical table. The input bridge and output bridge are separate learned matrices in that implementation. Requiring $R_{out}=P_{in}^T$ is another, stronger architecture constraint; it is not implied by lexical-table tying.

Ignoring biases, normalization, and activations, the parameter ledger is:

Architecture	Weight parameters	Direct alias legal?
direct input + direct output, untied	$2VH$	no alias declared
direct input/output tied	$VH$	yes, same [V,H] table
factorized input + direct [V,H] output	$VD+DH+VH$	lexical/output shapes differ
factorized lexical decoder tied, separate input/output bridges	$VD+2DH$	tie [V,D] lexical classifier only
factorized lexical decoder plus transposed shared bridge	$VD+DH$	only under the additional bridge-sharing contract

For V=128,000, H=4096, and D=256:

direct untied:                         1,048,576,000
direct tied:                             524,288,000
factorized input + direct output:         558,104,576
factorized lexical-tied, two bridges:      34,865,152
factorized lexical + shared bridge:        33,816,576

The compute ledger also changes. For M gathered input rows and N classifier rows:

direct output classifier:       N*V*H MACs
factorized input bridge:        M*D*H MACs
factorized output bridge:       N*H*D MACs
factorized output classifier:   N*V*D MACs

These are logical dense MACs. Fused losses, vocabulary parallelism, sampled or adaptive classifiers, and kernel packing change execution without changing the declared parameter graph.

Eligibility matrix: row meaning comes before shape

Candidate sharing relation	Admission rule
decoder input <-> decoder output	Same ordered target vocabulary and compatible stored width.
encoder input <-> decoder input	Same ordered joint vocabulary, not merely equal row count.
encoder input <-> decoder input <-> output	All three roles must satisfy the same row identity; Press-Wolf three-way tying used a union source/target vocabulary.
token table <-> type or position table	Invalid: rows belong to different address spaces even when widths match.
hashed bucket table <-> vocabulary output	Invalid as a direct alias when many tokens share a bucket and the classifier still needs one logit row per token. A decomposed classifier would be another architecture.
imported frozen table <-> trainable output	Tying would also freeze or jointly train the classifier. That training-policy change must be explicit.
factorized [V,D] input <-> direct [V,H] output	Invalid for $D\ne H$; insert and own a compatible output bridge or keep the head untied.

Equal V_src and V_tgt do not prove row compatibility. A permutation of the target vocabulary preserves every shape and makes three-way tying semantically wrong.

Colocated implementation: one object, one gradient, one optimizer state

A minimal direct tie assigns the Parameter reference, preferably before constructing the optimizer:

class DirectTiedModel(nn.Module):
    def __init__(self, vocab_size, hidden_size):
        super().__init__()
        self.token_embedding = nn.Embedding(vocab_size, hidden_size)
        self.lm_head = nn.Linear(hidden_size, vocab_size, bias=False)
        self.lm_head.weight = self.token_embedding.weight

    def embed(self, input_ids):
        return self.token_embedding(input_ids)

    def project(self, hidden):
        return F.linear(hidden, self.token_embedding.weight)

The post-construction gate should test object and storage identity, not values:

left = model.token_embedding.weight
right = model.lm_head.weight

assert left is right
assert left.untyped_storage().data_ptr() == \
       right.untyped_storage().data_ptr()
assert left.storage_offset() == right.storage_offset()
assert left.shape == right.shape
assert left.stride() == right.stride()

data.copy_(...), equal tensor hashes, and two Parameters that happen to share initial values all fail the colocated tying contract. They can acquire separate gradients, separate optimizer states, and separate updates.

Optimizer construction is part of the graph freeze:

params = list(model.parameters())
assert len({id(p) for p in params}) == len(params)

optimizer = AdamW(params, ...)
assert optimizer_registers_once(model.token_embedding.weight, optimizer)

Replacing lm_head.weight after optimizer construction can leave the optimizer holding the old Parameter. A valid restore order is:

construct architecture and alias map
load or assign checkpoint values
re-establish requested aliases
verify object/storage identity
construct optimizer and restore one logical optimizer state
run one two-path gradient/update probe

A state dictionary is not the model’s alias graph

PyTorch documents that Module.state_dict() is a shallow mapping containing references to module parameters and buffers. torch.save() preserves tensor storage sharing in the serialized Python object. Those facts do not mean that load_state_dict() will redesign an independently constructed module into a tied one: the usual load path copies checkpoint values into the Parameter objects already registered by the model.

Therefore the checkpoint contract needs both data and topology:

canonical logical tensor:
    token_table.weight

alias map:
    lm_head.weight -> token_table.weight

semantic prerequisites:
    ordered vocabulary hash, shape, dtype, factorization, shard map

post-load observation:
    alias/storage identity or declared synchronized-replica relation

If a checkpoint exposes both names, validate their values before choosing one canonical tensor. Two present alias keys with different values are not a minor serialization detail; they describe an untied or corrupted state and need an explicit conversion policy.

The current Transformers source makes the topology model-aware through a {target: source} tied-key mapping gated by tie_word_embeddings. Its load path can use the present side when the other key is absent, checks equality when both are present, and assigns the target attribute to the source Parameter. The exact behavior is version-scoped; the invariant here is the explicit mapping plus post-load identity verification.

Safetensors deliberately does not encode arbitrary PyTorch shared-storage graphs. Its official save_model/load_model helpers are model-aware, can drop duplicate names, and document that saved metadata is not sufficient to recover an arbitrary shared structure by itself. A plain tensor-name loader therefore cannot be the sole owner of alias restoration.

def restore_tied_checkpoint(model, tensors, alias_map):
    validate_vocab_shapes_and_shards(model, tensors)
    reject_divergent_present_alias_values(tensors, alias_map)
    model.load_state_dict(tensors, strict=False)
    apply_alias_map(model, alias_map)
    verify_alias_graph(model, alias_map)
    return model

PyTorch’s load_state_dict(assign=True) changes which tensor properties are preserved and requires optimizer construction after loading under the official contract. Treat it as a separate restore mode and rerun the same alias graph gate; do not infer that assign=True repairs an undeclared topology.

Conversion is an architecture migration

Migration	Required policy
tied -> untied	Clone the canonical tensor into a new output Parameter, create separate optimizer state, and record that future updates may diverge.
untied -> tied	Choose a canonical value or merge rule deliberately; validate the quality change. Equality at one instant is not a justification for silently discarding one trained matrix.
direct tied -> factorized tied	Define lexical width, input/output bridges, and a conversion or retraining procedure. A reshape cannot preserve the original classifier exactly in general.
one vocabulary -> joint three-way vocabulary	Build an explicit row mapping, initialize new rows, migrate output bias and optimizer state, and update tokenizer/checkpoint identity together.
native tied -> separately quantized input/output packs	Preserve one logical source tensor and record role-specific derived packs, or declare runtime untying. Different independent scales are not one mutable training parameter.

Section 0x05D owns full vocabulary migration and release mechanics. This section establishes that a tying-mode change cannot be hidden inside an ordinary checkpoint load.

Tensor-parallel output ownership

With vocabulary-row tensor parallelism, rank r owns:

$$E_r=E[s_r:e_r,:] \in\mathbb{R}^{V_r\times H}.$$

The masked-local input lookup from section 0x04A reconstructs hidden rows from the owner of each token. The same local shard can produce local classifier scores:

$$Z_r=CE_r^T \in\mathbb{R}^{N\times V_r}.$$

The system may all-gather the Z_r blocks for full logits, or use a vocabulary-parallel cross entropy that computes global normalization and target terms without materializing [N,V] on every rank. M3.8 owns the score/loss interface. The tying requirement is narrower:

input row interval on rank r == output row interval on rank r
local input/output roles reference one local shard when colocated
output bias uses the same vocabulary partition
checkpoint row order reconstructs the same global table

If input and output partitions differ, equal global values are insufficient; the runtime needs a declared reshard/remapping step and the checkpoint cannot label the local storages as direct aliases.

Pipeline-parallel tying uses synchronized replicas, not pointer identity

The first pipeline stage performs input lookup while the last stage performs output projection. Different processes cannot share one Python Parameter pointer. A valid logical tie maintains two physical replicas per tensor-parallel row shard:

first stage: E_r^(in)
last stage:  E_r^(out)

initialization:
    E_r^(in) == E_r^(out)

gradient reconciliation:
    g_r = g_r^(input) + g_r^(output)

update:
    both replicas apply the same update from equal optimizer state,
    or one authority steps and broadcasts the new value

The pinned Megatron-Core implementation makes this exception explicit. With pipeline parallelism it creates a last-stage copy, initializes first/last-stage values through an embedding-group all-reduce, marks the copies as shared, and requires their gradients to be reduced before the update. Its sharded state dictionary maps the last-stage output copy to the first-stage embedding’s logical checkpoint key with a replica identifier.

For each row shard, the reconciliation equation is:

$$g_r^{logical}=g_r^{input}+g_r^{output}.$$

If every replica begins with the same value and optimizer state and consumes the same reconciled gradient, deterministic identical updates preserve equality. A mere periodic value copy is weaker: independent optimizer moments can make replicas diverge again at the next step.

This yields two separate verification predicates:

colocated tie:
    SAME_PARAMETER_OBJECT && SAME_STORAGE

cross-stage logical tie:
    SAME_LOGICAL_PARAMETER_ID
    && SAME_TP_ROW_INTERVAL
    && REPLICA_GROUP_DECLARED
    && GRADIENT_RECONCILIATION_DECLARED
    && OPTIMIZER_AUTHORITY_DECLARED

Calling both realizations “equal tensors” erases the distinction this section is designed to protect.

Executable Tying And Checkpoint Gate

The local NumPy oracle freezes four independent proofs:

two-path gradient:
    dE_tied == dE_input + dE_output

architecture ledger:
    direct and factorized parameter/MAC formulas reconcile numerically

colocated restore:
    equal copies fail; one object/storage alias passes

distributed restore:
    tensor-parallel shards reconstruct the logical table and pipeline replicas
    remain equal only after gradient plus optimizer-state reconciliation

The fixture also compares direct untied, direct tied, factorized untied-output, factorized lexical-tied, target-only seq2seq tying, and compatible three-way tying. Negative cases reject row-order mismatches, shape-only ties, divergent duplicate keys, missing alias maps, duplicated optimizer registration, stale optimizer references, incompatible factorization, shard gaps/overlaps, input/output partition mismatch, unsynchronized pipeline gradients, divergent replica optimizer state, and tied claims over independently quantized mutable copies.

Passing this oracle proves the synthetic algebra and declared ownership graph. It does not prove framework autograd, NCCL ordering, production checkpoint conversion, or model-quality preservation after changing the architecture.

0x0A1 Training-state and runtime ledger

For a dense trainable table with V rows, width H, and b_w bytes per stored weight:

$$M_{weight}=VHb_w.$$

Mixed-precision Adam training may additionally hold a master weight and two optimizer moments. A rough ledger is:

forward weight, bf16:       2 * V * H bytes
master weight, fp32:        4 * V * H bytes
first moment, fp32:         4 * V * H bytes
second moment, fp32:        4 * V * H bytes
gradient, implementation:   dense or sparse policy-dependent

This can make a “simple lookup table” expensive during training. At inference, the logical arithmetic per selected row is tiny, but memory access is irregular and table residency matters. Repeated IDs may benefit from cache locality; large distributed vocabularies require ownership and communication.

Rare rows create an optimization problem. If row v occurs c_v times in the training corpus, its input-side gradient evidence is proportional to those occurrences, while an untied output head still sees vocabulary-wide softmax gradients. Undertrained-token audits inspect frequency, row norms, output probability pathologies, and tokenization neighborhoods rather than assuming every allocated row learned equally well.

Padding and regularization interaction

padding_idx suppresses ordinary embedding-gradient updates for the PAD row, but global optimizer or regularization logic can still matter. Decoupled weight decay applied blindly to all dense parameters may move a supposedly fixed row. Models that require a zero PAD row can enforce it after optimizer step or exclude the row/parameter according to their implementation contract.

@torch.no_grad()
def enforce_padding_row(embedding, pad_id):
    if pad_id is not None:
        embedding.weight[pad_id].zero_()

Do not add this hook to a checkpoint that intentionally learned a nonzero PAD row. The invariant comes from the model definition.

Local parameter, arithmetic, and placement ledger

For a text table E: [V,d_model]:

$$N_{\text{embedding params}} = Vd_{model}.$$

Lookup returns $Bn$ rows, so its logical output contains $Bnd_{model}$ elements. It is a gather, not the dense multiplication by a materialized one-hot matrix used in the conceptual proof. At V=32,000 and d_model=4096, the table has 131,072,000 parameters and occupies 250 MiB at two bytes per parameter.

For a linear media projector from d_media to d_model:

$$N_{\text{projector params}} = d_{media}d_{model},$$ $$\operatorname{MAC}_{\text{projector}} = Bn_{media}d_{media}d_{model}.$$

With B=1, n_media=576, d_media=1024, and d_model=4096:

projector parameters: 4,194,304 = 8 MiB at 2 bytes/parameter
projector MACs:       2,415,919,104
projected states:     2,359,296 elements = 4.5 MiB at 2 bytes/element

A two-layer projector replaces the parameter term with d_media*d_hidden + d_hidden*d_model and materializes an intermediate [B,n_media,d_hidden]. Count-changing resamplers add a separate row-reduction contract; width matching alone does not explain their cost or information loss.

Placement matters because these are boundary tensors. A vocabulary-sharded embedding must map global token IDs to the owning shard and combine the lookup result according to its sharding scheme. One implementation may mask non-local IDs, perform local lookup, and all-reduce hidden outputs; another may route IDs to owners. It does not necessarily gather full table rows. A vision encoder and language model on different devices must transfer [B,n_media,d_media] or the projected [B,n_media,d_model], depending on where the projector is placed. Moving the projector changes transfer width and compute placement without changing the article’s final shape. A device/dtype conversion that silently moves the tensor through host memory can dominate a small projector even though every shape is correct.

---

0x0B Shape ledger

Use this as the local debugging table.

Object	Shape	Owner / note
text token IDs	[B,L_text]	M2.1 output / M2.2 input
source IDs	[B,S]	data path / encoder input
target IDs	[B,T]	data path / decoder input
embedding table E	[V,d_model]	M2.2
source embedding table	[V_src,d_model]	may differ from target
target embedding table	[V_tgt,d_model]	decoder input table
PAD row	[d_model]	row at pad_id if valid
text embeddings	[B,L_text,d_model]	E[input_ids]
source content embeddings	[B,S,d_model]	before PE if only Embeddings
target content embeddings	[B,T,d_model]	before PE if only Embeddings
full src_embed(src)	[B,S,d_model]	embedding plus PE in Harvard code
full tgt_embed(tgt)	[B,T,d_model]	embedding plus PE in Harvard code
PE buffer	[1,max_len,d_model]	M3.3 handoff
media features H_media	[B,n_media,d_media]	X1 checkpoint boundary
image features H_img	[B,N_img,d_vision]	vision encoder output
audio features H_audio	[B,N_audio,d_audio]	audio encoder output
document features H_doc	[B,N_doc,d_doc]	document processor output
projector conceptual weight W_p	[d_media,d_model]	framework may store transpose
projected media states Z_media	[B,n_media,d_model]	published X1 boundary
count-changing adapter output	[B,r,d_model]	row count changed
row-type metadata	[B,N_total] or side metadata	must align to rows
mixed state sequence X0	[B,N_total,d_model]	published X1 output
output head weight	[V,d_model]	boundary note
logits / log-probs	[B,L,V]	M3.8/M3.10/A15 boundary

The main pattern:

correct last dimension is necessary.
correct row identity is also necessary.

Many bugs preserve the first and break the second.

---

0x0C Conformance tests for the initial-state boundary

Small tests can prove the algebra independently of a full training run.

Lookup equals one-hot multiplication

def test_lookup_matches_one_hot():
    torch.manual_seed(0)
    V, H = 11, 5
    ids = torch.tensor([[2, 7, 2], [0, 10, 4]])
    weight = torch.randn(V, H)

    gathered = F.embedding(ids, weight)
    selected = F.one_hot(ids, num_classes=V).to(weight.dtype) @ weight

    torch.testing.assert_close(gathered, selected)

This proves forward equivalence for the test case; it does not justify materializing one-hot tensors in production.

Repeated IDs accumulate gradients

def test_repeated_ids_scatter_add():
    embedding = nn.Embedding(8, 3)
    ids = torch.tensor([4, 2, 4])
    weights = torch.tensor([1.0, 1.0, 2.0]).unsqueeze(-1)

    loss = (embedding(ids) * weights).sum()
    loss.backward()

    torch.testing.assert_close(
        embedding.weight.grad[4],
        torch.full((3,), 3.0),
    )
    torch.testing.assert_close(
        embedding.weight.grad[2],
        torch.full((3,), 1.0),
    )

Padding, bounds, dtype, and tying

def assert_initial_state_contract(model, tokenizer, batch):
    ids = batch["input_ids"]
    assert ids.dtype in (torch.int32, torch.int64)
    assert ids.min().item() >= 0
    assert ids.max().item() < model.token_embedding.num_embeddings
    assert tokenizer.fingerprint == model.config.tokenizer_hash

    x = model.token_embedding(ids)
    assert x.shape == (*ids.shape, model.config.embedding_dim)
    assert torch.isfinite(x).all()

    if model.config.padding_idx is not None:
        pad = model.config.padding_idx
        pad_positions = ids == pad
        if pad_positions.any() and model.config.zero_pad_row:
            torch.testing.assert_close(
                x[pad_positions],
                torch.zeros_like(x[pad_positions]),
            )

    if model.config.tie_word_embeddings:
        assert model.token_embedding.weight.data_ptr() == \
               model.lm_head.weight.data_ptr()

Scale and factorization checks

def test_harvard_scaling(wrapper):
    ids = torch.tensor([[1, 3]])
    raw = wrapper.lut(ids)
    scaled = wrapper(ids)
    torch.testing.assert_close(scaled, raw * math.sqrt(wrapper.d_model))


def test_factorized_shape(module):
    ids = torch.tensor([[1, 2, 3], [4, 5, 6]])
    lexical = module.lookup(ids)
    hidden = module(ids)
    assert lexical.shape == (2, 3, module.lookup.embedding_dim)
    assert hidden.shape == (2, 3, module.bridge.out_features)

Checkpoint tests should additionally verify vocabulary row count, tokenizer hash, table dtype, sharding metadata, PAD policy, factorization dimensions, and actual aliasing of tied parameters. These are release invariants, not comments.

0x0D Failure lab

These are practical bugs where the tensor rank survives but the bridge contract breaks.

The list is intentionally broad, but breadth alone is not closure. Each bold failure name is a stable review marker tied to an input-state invariant or a neighboring owner. That makes the lab an ownership map rather than a collection of warnings.

Artifact and table-identity bugs

One release hash used as every kind of identity.

Resharding changes serialized artifacts without necessarily changing the logical tensor; changing row order can preserve shape while changing semantics. Keep address, logical-table, checkpoint, and serialized-artifact identities separate.

Equal vocabulary size admitted as address compatibility.

Two tokenizers can have the same V and different ID-to-piece assignments. Bind the table to the tokenizer semantic hash and ordered vocabulary digest.

Tied weights accepted from equal values.

Two independent tensors can compare equal at load time and diverge on the next update. Verify shared storage after materialization.

Shard map accepted from local shapes alone.

Overlapping intervals duplicate owners; gaps make rows unreachable. Require an exact, non-overlapping cover of [0,V).

ID / embedding boundary

Token IDs treated as numeric magnitudes.
ID 9000 is not semantically larger than ID 42. Both are table addresses.

Strings passed into nn.Embedding.
The layer expects integer indices, not token strings.

One-hot vectors passed into nn.Embedding.
One-hot is a conceptual explanation for one_hot @ E. The implementation expects IDs.

Already embedded floats passed into nn.Embedding.
The shape may look sequence-like, but embedding lookup is row indexing. It is not a generic projection layer.

Media features cast to IDs.
Converting [B,n,d_media] floats to integers and using a text embedding table destroys continuous feature meaning.

Vocabulary and table bugs

Tokenizer vocabulary and embedding row count mismatch.
The tokenizer emits IDs outside [0,V-1], or the model has rows the tokenizer never emits.

Tokenizer/model mismatch.
An ID is in range but addresses a row trained for a different token.

Source and target embedding tables swapped.
Both tensors may be [B,L]; the address space is wrong.

Shared embedding assumed when the implementation uses separate tables.
Parameter tying and checkpoint loading can silently become invalid assumptions.

Target Generator initialized with V_src.
This works only when source and target vocabularies are actually compatible.

A reordered vocabulary is called append-only.

Every tensor shape is valid, but old ID v now names a different token. Prove that the complete old ordered vocabulary is an exact prefix and that special roles retain their IDs.

An approximate initializer is recorded as an exact row copy.

A decomposition average or semantic neighbor can be useful, but it does not preserve one trained row’s meaning. Keep exact-map edges and initialization provenance in separate ledgers.

Old raw logits are preserved, so old probabilities are assumed preserved.

New output rows add A to the softmax denominator and scale every old probability by Z/(Z+A). Measure new-token mass and the resulting KL term.

The embedding grows but the untied output head or bias does not.

The model can read the new ID but cannot score it, or the classifier has a misaligned vocabulary axis. Migrate every declared vocabulary-indexed role.

Adam moments are resized but the step is reset.

The new shape loads, but old rows use a different bias-correction convention. Preserve parameter age or declare a custom row-age design; do not reset the whole table to make only the new rows young.

Only the sparse lookup gradient is masked to freeze old tied rows.

Dense output gradients, existing moments, or decoupled weight decay still move the rows. Freeze the combined update path, not one gradient producer.

Physical shards are edited independently without a canonical reconstruction.

Local files can all pass shape checks while row order, moments, or tied intervals diverge. Convert logical state first, then prove the shard cover and reconstruction, or prove an equivalent streaming transformation.

Quantized bytes are appended without regenerating group metadata.

Group boundaries, scales, padding, packing, or alignment can become stale. Migrate canonical state and regenerate the derived deployment artifact unless the packed-state conversion is separately proved.

T1/M0 is admitted because most IDs are still in range.

One new or semantically changed ID is enough to corrupt lookup and output meaning. Route tokenizer and model as an admitted version pair.

Rollback restores M0 while retaining T1 IDs or candidate cache state.

The process is live but conditions on foreign addresses or K/V rows. Restore the complete old tuple and reconstruct state from compatible source data.

PAD bugs

PAD ID, padding_idx, attention mask, and label mask disagree.
Lookup runs, but PAD rows may be updated, attended to, or supervised incorrectly.

pad_id=-1 passed into embedding lookup.
-1 may be tokenizer metadata, but ordinary embedding lookup expects valid nonnegative row IDs.

PAD row updates unexpectedly.
If padding_idx is not set or the wrong ID is used, PAD can train like content.

PAD row frozen when intended as content.
The reverse bug: padding_idx can freeze a row that another component treats as a real token.

Backward and distributed ownership bugs

Duplicate IDs overwrite instead of accumulate.

Assignment to dE[id] retains only one occurrence. The backward operation is scatter-add, or an equivalent sorted/segmented reduction.

Frequency scaling uses an undeclared count scope.

Rank-local counts, global effective-batch counts, and accumulation-window counts produce different updates. Record the scope and reduction convention.

Sparse layout is confused with sparse values.

A mostly-zero dense tensor is not a sparse-layout gradient, and an uncoalesced sparse tensor can contain duplicate row indices. Optimizer admission must inspect layout and coalescing state.

Optimizer name is treated as a complete sparse contract.

Device, weight decay, momentum/fused options, clipping, and gradient layout can invalidate a superficially supported optimizer family.

Local clipping is called global clipping.

Each owner clips with its local norm, so the global sharded gradient exceeds the declared threshold. Reduce the norm statistic across owners before scaling.

Tied output gradient is dropped to preserve input sparsity.

The lookup path can be sparse while the vocabulary-wide output path is dense. The shared parameter must receive both contributions under one admitted layout.

Data-parallel AVG and loss averaging divide twice.

The gradient is scaled by an unintended extra world-size factor. Freeze loss, microbatch, accumulation, and collective conventions together.

Row-shard intervals contain a gap or overlap.

A valid global ID has no owner or two owners. Validate exact coverage of [0,V) before forward.

Global IDs index a local shard directly.

The lookup is in bounds for some IDs but addresses the wrong row. Subtract the owner’s shard start only after proving ownership.

Origin-position metadata is lost during routing.

Correct rows return in expert/owner order rather than sequence order. Preserve origin rank and flattened position through both collectives.

A non-owner rank steps a row shard.

Forward reconstruction can be correct while optimizer state diverges. The manifest must name parameter and optimizer-state owners.

Column-wise weight placement is mistaken for a full activation schedule.

The next operator requires a full hidden row, but no concatenate/all-gather or compatible sharded layout is declared.

Tensor-parallel and data-parallel groups are conflated.

The row shard is correctly unique across tensor-parallel ranks but its replicas are never reduced across the data-parallel group, or the reverse.

Weight-tying, checkpoint, and output-ownership bugs

Equal values are mistaken for one parameter.

embed.weight and lm_head.weight have equal hashes after initialization but are distinct Parameter objects. They receive separate gradients and optimizer state and can diverge on the first update. Verify object and storage identity for a colocated tie.

A checkpoint stores values but loses the alias graph.

The loader copies two keys into two independently constructed Parameters and never reapplies lm_head.weight -> token_embedding.weight. Serialize a canonical logical tensor plus an alias map, restore the topology explicitly, and verify it after loading.

Two declared alias keys contain different values.

Silently selecting one tensor discards trained state and hides an untied or corrupt checkpoint. Reject the load unless an explicit untied-to-tied migration policy chooses or merges the canonical value and validates the quality change.

The optimizer is built before an alias is restored.

The optimizer retains the replaced output Parameter while forward uses the embedding Parameter. Re-establish and verify aliases before constructing the optimizer, or rebuild optimizer parameter groups and restore one logical state.

One tied Parameter is registered twice with the optimizer.

The same logical weight can be stepped twice or carry duplicate state. Deduplicate by object identity and assert exactly one optimizer registration per colocated logical parameter.

A factorized lexical table is tied directly to a hidden-width head.

[V,D] and [V,H] have different coordinate spaces when $D\ne H$. A reshape or copied prefix is not a classifier derivation. Insert an owned hidden-to-lexical bridge and tie only the compatible [V,D] decoder factor, or keep the head untied.

Source and target tables are tied from equal row counts.

Equal V and shape do not prove equal row meaning. Require the same ordered vocabulary identity, including special-token assignments, or an explicit joint vocabulary migration.

Tensor-parallel input and output shards use different row intervals.

Each local shape can be valid while lookup row v and classifier logit v refer to different global rows. Require identical input/output partitions or a declared reshard and remapping operation.

Pipeline replicas reduce values but not the two gradient paths.

The first stage steps only the input contribution and the last stage steps only the output contribution. Reconcile the logical sum before update, then apply one authoritative update or identical updates from equal optimizer state.

Pipeline replica values match while optimizer moments differ.

A periodic parameter broadcast makes one checkpoint look correct, but the next local optimizer step diverges again. Synchronize or centralize optimizer state as part of the logical tie.

The reconciled pipeline gradient is reduced twice.

Both a tied-embedding group and a wider data-parallel collective sum the same contribution without a declared normalization convention. Freeze group membership and reduction order, then reconcile the oracle against one logical gradient.

Input and output roles create independent quantized packs and call them tied.

Different scales, calibration state, or mutable packed values are two physical representations. Keep one logical source tensor with role-specific derived packs and a regeneration policy, or explicitly declare runtime untying.

Harvard reproduction bugs

Missing sqrt(d_model).
Shapes match, but the source-compatible scale is wrong.

Applying sqrt(d_model) twice.
This happens when both the embedding wrapper and caller scale.

Treating sqrt(d_model) as an initializer-independent law.

The output variance is d_model * sigma^2, not automatically one. Record the actual initializer, post-initialization pass, and scale owner.

Calling rectangular Xavier variance 1/d_model.

For a PyTorch [V,d_model] tensor with gain one, the Xavier variance is 2/(V+d_model). The square/fan approximation is not the exact Harvard table ledger.

Fresh initialization confused with checkpoint loading.

A missing checkpoint key leaves random constructor rows behind. An exact load must fail rather than silently convert missing trained state into fresh state.

Vocabulary growth overwrites old rows.

Resize code reinitializes the full table rather than only appended/remapped rows. Compare the old prefix exactly and update optimizer state, output rows, biases, configuration, tokenizer identity, and tied aliases.

Assuming padding_idx forces a permanent zero row.

The constructor zeros a fresh PAD row, but supplied weights bypass that reset; global initialization, checkpoint load, tying, and regularization can define or change a nonzero row. The lookup-gradient rule is narrower than a zero-row invariant.

Ignoring max_norm as a forward mutation.

Selected rows can change during lookup even under inference and without an optimizer step. Hashing, autograd regularizers, checkpoint comparisons, and concurrent readers must account for that state change.

Equating autocast with converted parameters or fp32 master weights.

Autocast is operator-specific. Parameter storage, gathered output, gradient, optimizer moments, loss scaling, and an optional master copy need separate observed fields.

Calling only Embeddings when code expects Embeddings + PositionalEncoding.
The output is [B,L,d_model], but rows lack the positional component expected by the Harvard stack.

Calling src_embed when only raw content embeddings were intended.
Position encoding can be applied too early or twice.

Projector bugs

Mistaking d_media for d_model.
A vision encoder may output [B,N,1024]; the language model may expect [B,N,4096].

Projector maps to the correct width but wrong semantic space.
Z_media: [B,N,d_model] is shape-compatible but not aligned if the projector was not trained for this encoder and language model.

Projector trained for one encoder is used with another encoder.
The same d_media can hide incompatible feature distributions.

Projector frozen or trained contrary to the setup.
Both choices can be correct in different systems. The bug is making the wrong choice for the interface being reproduced.

Embedding-object boundary bugs

Treating lookup rows as contextual states.

The same token ID selects the same input row before the Transformer stack; context-dependent meaning is produced later.

Using arbitrary hidden-state cosine as a calibrated retrieval score.

Vector width and cosine computability do not prove that a downstream readout, normalization rule, or representation objective trained the space for retrieval. M3.11 owns that proof.

Calling projected patch rows a single image embedding.

[B,N,d_model] model-consumable rows and a pooled [B,d_retrieval] task vector have different row counts, consumers, and objectives.

Treating equal-width downstream representations as interchangeable.

Vectors from different layers, pooling heads, checkpoints, or objectives can share a dimension without sharing a coordinate system. M2.2 records only the handoff; M3.11 owns their comparison and admission.

Count-changing adapter bugs

Adapter changes n but masks still use old n.
[B,n,d] -> [B,r,d_model] succeeds while masks, row metadata, or labels still assume n.

Adapter reorders rows without metadata update.
The output rows have the correct width but no longer correspond to the original patches, spans, or events.

Query bottleneck collapses needed evidence.
The adapter produces valid [B,r,d_model] rows, but small OCR, visual, or audio evidence is compressed away.

Mixed-origin X0 bugs

Text and media concatenated along the feature axis.
The usual mixed-context path concatenates rows along dim=1, not features along dim=-1.

Rows have correct shape but wrong order.
Training may expect image rows before the question. Inference may append them elsewhere.

Row-type metadata lost.
The model or wrapper may no longer distinguish image rows, text rows, event rows, and PAD rows in the way training expected.

Geometry metadata dropped at concat.
Rows have width d_model, but layout, time, or order metadata no longer aligns for X1 and M3.3.

Ragged media batches padded without media masks.
Padded media rows become evidence rows.

Output-head and weight-tying bugs

Output head and input embedding vocabulary dimensions disagree.
Input lookup works; logits have the wrong vocabulary dimension.

Assuming tied weights when the implementation does not tie.
Checkpoint loading, parameter counting, and output behavior can be wrong.

Tying incompatible source, target, and output tables.
Weight sharing is not a free operation across different vocabularies.

Using transposed weight layout incorrectly.
Conceptually hidden @ W_out.T; framework storage for nn.Linear(d_model,V) has its own layout. Custom implementations must respect it.

Article-Wide Failure Closure

The lab contains 84 named scenarios across twelve ownership groups:

$$4+5+15+4+13+12+11+4+4+3+5+4=84.$$

Each scenario keeps its exact article marker and names the first contract that can reject it. A guard does not need to reproduce a full training run: tensor and gradient errors close through executable oracles, while semantic ownership, provenance, and cross-module errors close through immutable manifests.

The count describes this article’s edited input-state surface. Contextual pooling, contrastive-objective, retrieval-index, and representation-runtime failures are intentionally counted and closed by M3.11 or their downstream retrieval/deployment owners instead of being duplicated here.

---

0x0E Handoff to M3 and X1

M2.2 produces text model-width rows:

X_text: [B,L_text,d_model]

That solves the width problem.

It does not solve the geometry problem.

After this bridge, the model still needs geometry and computation:

text order
attention visibility and score algebra
position IDs / attention bias
RoPE and context extension
residual/normalization/block execution

M3.1 starts attention algebra. M3.3 through M3.5 own positional geometry. The parallel route now enters X1.1 for non-text units and geometry, then X1.2 for encoders, projectors, resampling, and fusion.

After M3.1-M3.7 have produced contextual rows, M3.11 owns the separate route from those rows to pooled representations, task heads, typed scores, and retrieval readout. M2.2 does not reopen that route.

---

0xFF Reference

Primary Papers And Official Technical Sources

• A Neural Probabilistic Language Model
• Using the Output Embedding to Improve Language Models
• Tying Word Vectors and Word Classifiers: A Loss Framework for Language Modeling
• Attention Is All You Need
• The Annotated Transformer
• torch.nn.Embedding — PyTorch documentation
• torch.nn.Embedding implementation — PyTorch source
• Embedding backward implementation — pinned PyTorch source
• SparseAdam — PyTorch documentation
• clip_grad_norm_ — PyTorch documentation
• FullyShardedDataParallel gradient clipping — PyTorch documentation
• Exploring TorchRec sharding — PyTorch tutorial
• TorchRec distributed embedding execution — PyTorch tutorial
• VocabParallelEmbedding — Megatron-Core documentation
• VocabParallelEmbedding — pinned Megatron-Core source
• torch.nn.init — PyTorch documentation
• Automatic Mixed Precision — PyTorch documentation
• Numerical accuracy — PyTorch documentation
• AdamW — PyTorch documentation
• Adam: A Method for Stochastic Optimization
• torch.nn.functional.one_hot — PyTorch documentation
• torch.nn.EmbeddingBag — PyTorch documentation
• torch.nn.Linear — PyTorch documentation
• Tensor attributes and strided layout — PyTorch documentation
• Tensor views and materialization — PyTorch documentation
• Module state dictionaries — PyTorch documentation
• Module.load_state_dict — PyTorch documentation
• Serialization semantics — PyTorch documentation
• PreTrainedModel.resize_token_embeddings — Hugging Face Transformers documentation
• FOCUS: Effective Embedding Initialization for Monolingual Specialization of Multilingual Models
• WECHSEL: Effective Initialization of Subword Embeddings for Cross-Lingual Transfer
• Optimizer.state_dict — PyTorch documentation
• Optimizer.load_state_dict — PyTorch documentation
• Optimizer.register_load_state_dict_pre_hook — PyTorch documentation
• Distributed Checkpoint state dictionaries — PyTorch documentation
• Megatron-Core distributed optimizer checkpoint formats
• Megatron-Core optimizer-state sharding helpers
• Transformers tied-weight restoration — pinned source
• ALBERT factorized masked-LM head — pinned source
• Torch shared tensors — Safetensors documentation
• Megatron-Core shared embedding/output ownership — pinned source
• Megatron-Core GPT output-layer tying — pinned source
• BLIP-2: Bootstrapping Language-Image Pre-training with Frozen Image Encoders and Large Language Models
• Visual Instruction Tuning
• Learning Transferable Visual Models From Natural Language Supervision
• Efficient Estimation of Word Representations in Vector Space
• Distributed Representations of Words and Phrases and their Compositionality
• Deep contextualized word representations
• BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding
• BERT reference implementation
• GPT-2 reference implementation
• LLaMA: Open and Efficient Foundation Language Models
• Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks
• C-Pack: Packed Resources For General Chinese Embeddings
• ALBERT: A Lite BERT for Self-supervised Learning of Language Representations
• Word Embeddings: Encoding Lexical Semantics — PyTorch tutorial
• DeepSeek-V3 model configuration — pinned revision

Contextual Surveys And Engineering Explanations

These sources preserve useful explanatory context. The article does not rely on them alone for a mechanism, implementation, or measurement claim.

• Initializing New Word Embeddings for Pretrained Language Models — John Hewitt

Local Corpus And Downstream Owners

• Transformer Series (2) — Overall Architecture
• Transformer Series (3) — Data Processing
• Transformer Series (6) — Token
• Transformer Series (7) — Embedding
• Transformer Series (8) — Position Encoding
• Transformer Series (15) — Sampling and Output
• Transformer System M2.1 — Tokenization and Sequence Contracts
• Transformer System M3.1 — Scaled Dot-Product Attention and Masks
• Transformer System M3.3 — Positional Geometry and Attention Bias
• Transformer System X1.1 — Unitization and Geometry Metadata
• Transformer System X1.2 — Encoders, Projectors, Resamplers, and Fusion
• Transformer System X1.3 — Masks, Training, Runtime, and State
• Transformer System M3.8 — Output Head and Decoding Policy
• Transformer System M3.9 — Training Execution
• Transformer System M3.10 — Single-Request Inference Execution
• Transformer System M3.11 — Representation Readout and Task-Head Inference
• M4.1 - Transformer Cost Model: Parameters, Activations, FLOPs, and Memory Traffic
• M5.1 - Autoregressive Serving State: KV Cache Lifecycle and Scheduler Pressure
• M6.2 - Expert Parallelism and DeepSeek MoE: Load Balancing, Shared Experts, and Fine-Grained Routing
• M8.10 - Deployment Evaluation, Release, and Rollback