DeiT-Style Transformer with Two Output Tokens

Two output tokens, two heads

A class-token vision transformer reads one token to produce class logits. This exercise adds a second learned token and a separate classification head. Both tokens pass through the same transformer blocks and can attend to all image patches, each other and themselves.

The design is inspired by the distillation token in DeiT. During distillation, a student learns from a teacher model as well as labeled examples. The class head is supervised by the true label; the distillation head is supervised by the teacher. Teacher supervision can use class labels or soft predictions, depending on the loss.

This task implements only a small forward model. It has no teacher, training loss or weight updates. Adding a token alone does not perform distillation or make a model data-efficient. The returned logits come from random initialization, so neither head has learned its intended role.

Keep the token order explicit

Let T be the number of image patches and D the model width. Build this sequence before adding positions:

row 0:       class token        (1, D)
row 1:       distillation token (1, D)
rows 2..T+1: patch tokens       (T, D)

The class and distillation tokens are separate parameters, each shared across images. A position table of shape (T + 2, D) gives every sequence position its own vector, including both output tokens.

For an image (4, 4, 3), patch size 2 and D=8, the patch convolution produces (2, 2, 8). Reshaping gives four patch tokens. Prepending the two output tokens gives (6, 8), which is also the position-table shape. All transformer blocks preserve that shape.

Read the heads separately

After a final LayerNorm, apply cls_head to row 0 and dist_head to row 1. Each head has its own Dense kernel and bias. The task returns the class logits first, followed by the distillation logits:

import jax.numpy as jnp

cls_logits = jnp.array([2., -1.])
dist_logits = jnp.array([0., 3.])
result = jnp.concatenate([cls_logits, dist_logits], axis=0)
# [2., -1., 0., 3.], not their average [1., 1.]

This example shows output ordering, not a trained model. Do not average the heads or apply softmax in the graded function. Concatenating two K-class outputs gives length 2 * K; it does not define a new 2K-class classification problem.

Use explicit names cls_head and dist_head to match the seeded model. Separate unnamed Dense instances could also have separate parameters, but their different names would change random initialization. Reusing one Dense instance for both tokens would share its weights.

Transformer block

Each compact ViTBlock has two residual branches:

h = LayerNorm(x)
x = x + SelfAttention(h)
h = LayerNorm(x)
x = x + Dense(D)(ReLU(Dense(d_ff)(h)))

Use an unnamed nn.MultiHeadDotProductAttention with num_heads=num_heads and qkv_features=D, passing the normalized sequence as its sole input. Every token may attend to every token. Use no attention mask or dropout. Retain other Linen defaults, including biases and LayerNorm settings. Each block has its own parameters.

Within a block, the submodule names are LayerNorm_0, MultiHeadDotProductAttention_0, LayerNorm_1, Dense_0 and Dense_1. Their nesting and creation order matter for seeded initialization.