Skip to content

← Flax step 33 of 100

Medium Primitives

Learned Position Embedding

Jump to the task

Give each position its own vector

Token embeddings identify vocabulary items. To represent where each token occurs, a model can add a separate vector for its position in the sequence. A learned position table stores one trainable row per absolute position: P has shape (max_T, D), and position t uses P[t].

For token features x with shape (B, T, D), the addition is x + P[:T][None, :, :]. Every sequence in the batch uses the same first T position vectors. Position zero means the start of the sequence here; cached decoding would need to select rows using the appropriate offset.

This exercise returns the position vectors alone. It initializes the table but does not add token embeddings or train any parameters.

Learned and fixed positions

A learned table adds max_T * D parameters. Training can adjust its entries to suit the task. A fixed sinusoidal encoding uses a formula instead, so it needs no trainable position table and can be evaluated at positions beyond those seen during training.

Being able to compute an encoding at a new position does not guarantee that the model will work well on longer sequences. A learned table has the additional constraint that its stored rows only cover positions 0 through max_T - 1. Extending that table requires an explicit choice about the new rows. Neither approach is universally more accurate.

Declare a parameter in Linen

self.param(name, initializer, shape) creates a parameter during init. During apply, it reads that parameter from the supplied variables. apply does not train the parameter. The @nn.compact decorator lets the module declare parameters inside its forward method.

This complete example creates a table and selects its first three rows:

import jax
import jax.numpy as jnp
import flax.linen as nn

class LearnedPos(nn.Module):
    max_T: int
    d_model: int

    @nn.compact
    def __call__(self, T):
        table = self.param(
            "pos_embed",
            nn.initializers.normal(stddev=0.02),
            (self.max_T, self.d_model),
        )
        return table[:T]

module = LearnedPos(max_T=32, d_model=4)
variables = module.init(jax.random.PRNGKey(0), 3)
positions = module.apply(variables, 3)
table = variables["params"]["pos_embed"]

assert table.shape == (32, 4)
assert positions.shape == (3, 4)
assert jnp.array_equal(positions, table[:3])

x = jnp.zeros((2, 3, 4))
combined = x + positions[None, :, :]
assert combined.shape == (2, 3, 4)
assert jnp.array_equal(combined[0], combined[1])

normal(stddev=0.02) samples an ordinary normal distribution with mean zero and standard deviation 0.02. It is not truncated normal. The scale is part of this exercise's specification, not a guarantee about training speed or stability. Keep the same initializer to match the expected output. Linen derives the parameter key from the key given to init, so sampling directly from the original key is not equivalent.

Select rows without changing the table

table[:T] selects rows 0 through T - 1. table[T] selects only row T, which is a different operation. Changing T during apply does not recreate the parameters: it selects a different prefix of the same initialized table.

If T exceeds the table length, slicing stops at the end of the table. It does not create the missing position vectors. Adding that shorter result to a longer sequence would then have incompatible shapes. This exercise keeps T within the table's capacity.

Problem

Implement learned_pos_embed_forward(seed, T, d_model):

  1. Convert seed, T and d_model to Python integers.
  2. Use a Linen module with max_T = 32 and a parameter named pos_embed of shape (32, d_model).
  3. Initialize it with nn.initializers.normal(stddev=0.02).
  4. Have the module return pos_embed[:T].
  5. Initialize the module with jax.random.PRNGKey(seed), then apply it using the returned variables and the same T.
  6. Return the result flattened with .reshape(-1).

d_model is positive and 0 <= T <= 32. The output is a one-dimensional array of length T * d_model, with each position's features together. Always initialize all 32 rows, even when T is smaller. If you wrap this function in jax.jit, make all three arguments static: they control the random key construction, parameter shape or slice length.

Loading visualization…