Skip to content

← Flax step 44 of 100

Hard Primitives

Mini T5: Shared Embeddings Across Two Stacks

Jump to the task

Connect an encoder to a decoder

Build a small encoder-decoder Transformer inspired by T5. The encoder reads the source tokens. The decoder reads its own input tokens and attends to the encoder's states. Return raw vocabulary scores for every decoder position, using one embedding table for both inputs and the output projection.

This exercise evaluates a randomly initialized model. It does not train, tokenize text, sample tokens or load a T5 checkpoint. Treat the layers below as the exact exercise architecture, not a complete T5 implementation.

source IDs (S,) -> shared embedding -> encoder blocks -> RMSNorm -> context (S, D)
                                                                     |
decoder IDs (T,) -> shared embedding -> decoder blocks <--------------+
                                          |
                                       RMSNorm
                                          |
                                multiply by embedding.T
                                          |
                                  raw logits (T, V)

Normalize each token's features

Define T5RMSNorm with a learned vector gamma of shape (D,), initialized to ones, and eps=1e-6. For each token row, compute:

rms_norm(x) = gamma * x / sqrt(mean(x ** 2, axis=-1, keepdims=True) + eps)

There is no mean subtraction or learned bias. Different norm instances have independent scales. This formula does not guarantee faster training or the same results as LayerNorm.

Encode with two residual updates

Each T5EncoderBlock contains unmasked self-attention and a ReLU FFN:

a   = x + self_attention(RMS0(x))
out = a + Dense(D)(relu(Dense(F)(RMS1(a))))

Every source position can attend to every source position. Create a new block for each encoder layer, then apply one final T5RMSNorm to produce the context. The FFN expands from D to F features and projects back.

Decode with three residual updates

Each T5DecoderBlock uses four norm instances, since cross-attention normalizes its query and context separately:

a   = x + causal_self_attention(RMS0(x))
b   = a + cross_attention(query=RMS1(a), key_value=RMS2(context))
out = b + Dense(D)(relu(Dense(F)(RMS3(b))))

The causal mask is a lower-triangular (T, T) array, including the diagonal. Decoder position i can read decoder positions 0 through i. Cross-attention has no mask: each decoder position can read all S context rows. Its scores have shape (H, T, S), so source and decoder lengths need not match. All supplied positions are valid; this exercise does not add padding masks.

Keep the context fixed across decoder blocks. After the last block, apply one final T5RMSNorm to the decoder states.

Reuse one embedding table

Construct one unnamed nn.Embed(V, D) with its default initializer. Call that same instance on both ID arrays. If its matrix is W, use hidden @ W.T for the output head, without a bias or softmax.

This shares V * D parameters across three uses. Compared with three independent matrices of that size, it saves 2 * V * D parameters. Gradients from the uses accumulate into the same matrix during training; sharing does not by itself guarantee better predictions.

Here is the lookup and output projection with known weights, leaving the Transformer blocks out so the shapes are easy to check:

import jax.numpy as jnp

W = jnp.array([[1., 0.], [0., 2.], [1., 1.]])
source_ids = jnp.array([0, 2], dtype=jnp.int32)
decoder_ids = jnp.array([1], dtype=jnp.int32)
source = W[source_ids]
decoder = W[decoder_ids]
logits = decoder @ W.T

assert source.shape == (2, 2)
assert decoder.shape == (1, 2)
assert jnp.array_equal(logits, jnp.array([[0., 4., 2.]]))
assert logits.reshape(-1).shape == (3,)

What omitting positions changes

This model adds neither position embeddings nor relative-position bias. With no source positions or source mask, permuting the source IDs permutes encoder rows and leaves decoder logits unchanged. The model therefore cannot distinguish source sequences that only reorder the same tokens. The decoder's causal mask still supplies an ordering constraint, so arbitrary decoder permutations do not have that property.

Problem

Implement mini_t5_forward(seed, x_enc_ids, x_dec_ids, vocab_size, d_model, num_heads, d_ff, num_enc_layers, num_dec_layers):

  1. Cast configuration values to Python integers and both ID arrays to jnp.int32. Define T5RMSNorm, T5EncoderBlock, T5DecoderBlock and MiniT5 as compact Linen modules.
  2. Follow the computations above. Use unnamed submodules in the shown order. Root scopes are Embed_0, T5EncoderBlock_0 through T5EncoderBlock_{LE-1}, T5RMSNorm_0, T5DecoderBlock_0 through T5DecoderBlock_{LD-1} and T5RMSNorm_1.
  3. All attention layers use num_heads=H and qkv_features=D, with default projection biases and zero dropout. Each block's FFN uses nn.Dense(F), nn.relu and nn.Dense(D), retaining default biases. Call cross-attention with the normalized query and normalized context as its two positional arguments.
  4. Initialize the whole model directly with jax.random.PRNGKey(int(seed)) and both ID arrays, apply those variables to the same arrays, and return the logits with .reshape(-1).

The source and decoder inputs are one-dimensional float arrays of lengths S and T containing integer-valued IDs in [0, V). S, T, V, D, H and F are positive, D is divisible by H, and both layer counts are nonnegative. Final norms remain even with zero blocks. With zero decoder blocks, the returned logits do not depend on the source.

Return length T * V, with vocabulary entries varying fastest. Keep seed and all six model configuration arguments static under JIT. Preserve module names, call order and initializer defaults for the seeded tests. This is a full-sequence forward pass, without an incremental cache.

The diagram's bidirectional, causal and cross-attention panels correspond to encoder self-attention, decoder self-attention and decoder cross-attention.

Loading visualization…