Tied Input/Output Embedding

Reuse one vocabulary matrix

An input embedding turns a token ID into a vector. An output projection turns a hidden vector into one score per vocabulary entry. Weight tying uses the same learned matrix for both operations.

Let W have shape (V, D), where V is vocabulary size and D is embedding width:

x = W[ids]            # (T, D): look up one row for each token ID
logits = x @ W.T      # (T, V): score every vocabulary row

This exercise feeds the lookup vectors directly into the output projection. There is no Transformer, context mixing or training step in between. It demonstrates parameter sharing, not a complete next-token predictor.

What each score means

logits[t, v] is the dot product of W[ids[t]] and W[v]. These are raw scores, not probabilities. Do not apply softmax.

The score for the input token itself is its squared vector norm: logits[t, ids[t]] = sum(W[ids[t]] ** 2). That score is nonnegative, but it is not necessarily the largest score. Another vector with a larger norm can have a larger dot product. Randomly initialized vectors also do not come with learned semantic meaning.

Work through a small example

Use known weights here so every score is easy to check:

import jax.numpy as jnp

W = jnp.array([[1., 0.], [2., 0.], [0., 1.]])
ids = jnp.array([0, 2, 0], dtype=jnp.int32)
x = W[ids]
logits = x @ W.T

assert jnp.array_equal(logits, jnp.array([
    [1., 2., 0.],
    [0., 0., 1.],
    [1., 2., 0.],
]))
assert logits[0, 1] > logits[0, 0]  # The input token need not score highest.
assert jnp.array_equal(logits[0], logits[2])  # Repeated IDs give repeated rows.

For this direct lookup-and-project computation, changing the token order only changes the output row order. A full language model normally transforms embeddings using their context before projecting to vocabulary scores.

What tying saves

Two independent vocabulary weight matrices contain 2 * V * D parameters. Sharing one leaves V * D, saving V * D parameters. For V=50000 and D=4096, that saves 204,800,000 parameters. This halves the weights in these two matrices, not necessarily the whole model. The exercise has no output bias or other trainable parameters.

Sharing also means gradients from lookup and output scoring accumulate into the same parameter. Even a vocabulary row not selected by ids can receive a gradient through its output score. Whether tying improves a trained model's quality depends on the model and training setup.

Use an explicit shared parameter here

Declare a single parameter named embedding with self.param("embedding", nn.initializers.normal(stddev=0.02), (V, D)) inside a compact Linen module. Both lookup and projection must use that array. Two independently trainable arrays would not be tied, even if initialized to equal values.

Linen also supports tying through an nn.Embed instance and its attend method, which projects using that instance's embedding matrix. This exercise uses self.param to make the shared array explicit and specifies its initializer. The default nn.Embed initializer is different.