Skip to content

← Flax NNX step 99 of 100

Medium Primitives

NNX Tied Input and Output Embeddings

Jump to the task

One table, two uses

An input embedding table E has shape (V, D): one D-dimensional vector for each of V vocabulary entries. Looking up token IDs produces input vectors. An output projection scores each vocabulary entry from a hidden vector. Weight tying uses the same table for both operations:

hidden = E[token_ids]       # (T, D)
logits = hidden @ E.T       # (T, V)

There is one learned parameter, not a second matrix initialized to the same values. A copied matrix could change independently during training. Reading E and its transpose on each forward pass keeps both uses tied.

In a language model, attention and other layers normally transform the input vectors before the output projection. This exercise leaves those layers out to isolate weight sharing. It is not a next-token predictor by itself, and no training or softmax is requested.

What each output means

For token position t and candidate vocabulary entry v:

logits[t, v] = dot(E[token_ids[t]], E[v])

The candidates cover the entire vocabulary, not just tokens present in the input. Logits are unnormalized scores and may be negative. Dot products depend on vector lengths as well as directions, so these are not cosine similarities. A token's own row is not guaranteed to have the largest score.

A complete example

import jax.numpy as jnp
from flax import nnx

class TiedTable(nnx.Module):
    def __init__(self):
        self.embedding = nnx.Param(jnp.array([[1.0, 0.0],
                                              [2.0, 0.0],
                                              [0.0, 3.0]]))

    def __call__(self, ids):
        table = self.embedding[...]
        return table[ids] @ table.T

model = TiedTable()
ids = jnp.array([0, 2, 0], dtype=jnp.int32)
expected = jnp.array([[1.0, 2.0, 0.0],
                      [0.0, 0.0, 9.0],
                      [1.0, 2.0, 0.0]])
assert jnp.allclose(model(ids), expected)
# Updating one table changes both the lookup and projection paths.
model.embedding[...] = model.embedding[...].at[0].set(jnp.array([2.0, 0.0]))
assert jnp.allclose(model(ids)[0], jnp.array([4.0, 4.0, 0.0]))

Token zero initially scores token one higher than itself because the second vector is longer in the same direction. Repeated input IDs give repeated output rows. Each position is processed independently here.

Initialization and parameter count

Initialize the exercise table with one normal draw using rngs.params(): jax.random.normal(key, (V, D), dtype=jnp.float32) / sqrt(D). Entries have variance 1/D, so a row's expected squared length is one. This is the initialization specified by this exercise, not a general claim that all embedding tables use this scale or a Xavier initializer.

An untied, bias-free embedding plus output projection has 2*V*D parameters. Sharing the table uses V*D, saving V*D. For V=50000 and D=768, that saves 38,400,000 parameters. The output multiplication still needs to compute all vocabulary scores; tying does not eliminate it. Model quality depends on the architecture and training setup.

Both uses contribute to gradients of the same parameter. The lookup contribution accumulates for repeated IDs. The output-projection path can also contribute to rows that were never looked up, depending on the loss. Weight sharing is therefore more than a forward-pass storage trick.

Problem

Write tied_io_embed(seed, token_ids, vocab_size, d_model) using a TiedEmbed(nnx.Module) with one nnx.Param named embedding.

  1. Convert seed, vocab_size and d_model to Python integers. Let V be vocab_size and D be d_model. Create nnx.Rngs(int(seed)) and draw one key with rngs.params().
  2. Initialize the float32 (V, D) table as described above. Do not add another projection matrix or an output bias.
  3. Cast token_ids to jnp.int32. Read the table with self.embedding[...], look up its rows, and multiply those vectors by the same table transposed.
  4. Return the logits flattened to (T * V,) in token-position order.

token_ids is a one-dimensional array of length T containing integer IDs passed as floats. Every ID is in [0, V). V and D are positive. The cast is a harness conversion, not validation or rounding of arbitrary fractional IDs. Scalar inputs also arrive as floats representing integers. Keep seed, V and D static under JIT; IDs can be dynamic array values.

Each entrypoint call initializes a fresh model. Forward calls on a retained model should not mutate its table. There are no position embeddings, contextual layers or probability normalization in this exercise.