ViT with Mixup Augmentation

Implement vit_train_with_mixup: mix pairs of images and their labels, run a Vision Transformer classifier, then take one SGD step using gradients you compute by hand. Return the updated parameters, not predictions or loss.

Mixup blends an image with another image from the batch. The target blends their one-hot labels by the same amount. For example, mixing class 0 and class 2 with lam=0.75 gives target [0.75, 0, 0.25] in a three-class task. Do not average the integer class IDs.

Inputs

• x: (N, C, H, W), a batch of images.
• y: (N,), class IDs in [0, num_classes). These arrive as floats; cast to integers for indexing.
• w_proj: (CPP, D), the patch projection.
• cls_token: (D,), the shared classification token.
• pos_embed: (M + 1, D), where M = (H/P)*(W/P).
• blocks_weights: (L, 6, D, D), packed as [w_q, w_k, w_v, w_o, w_mlp1, w_mlp2] for each block.
• w_head: (D, num_classes), the classifier matrix.
• num_heads: a positive integer dividing D.
• lr: the nonnegative learning rate; alpha: the positive Beta distribution parameter.
• seed: the integer random seed; num_classes: the positive number of classes.

Shapes are consistent. Image dimensions, P and D are positive; N is at least one and L may be zero. P divides H and W. Infer P from w_proj.shape[0] = CPP, or from the image area and patch count. Floating-point tensors are finite and share a dtype and device. Leave all inputs unchanged.

1. Draw the mixing coefficient and pair the examples

Use one scalar lam from Beta(alpha, alpha) for the whole batch, and one permutation idx of the batch indices. A permutation can pair an example with itself. With N=1, the image and target are unchanged by mixing.

Use the recipe for your framework below. The frameworks have separate expected results because their random generators differ. Draw the coefficient in float32 even if the model uses another dtype.

In PyTorch

Use fresh CPU generators and draw the two gamma samples in order:

g = torch.Generator(device="cpu").manual_seed(seed)
a = torch.tensor([alpha], dtype=torch.float32, device="cpu")
ga = torch._standard_gamma(a, generator=g)
gb = torch._standard_gamma(a, generator=g)
lam = (ga / (ga + gb)).item()
g_perm = torch.Generator(device="cpu").manual_seed(seed + 1)
idx = torch.randperm(N, generator=g_perm, device="cpu").to(x.device)

In JAX

The runtime uses its Threefry generator with partitionable mode enabled:

lam = float(jax.random.beta(jax.random.PRNGKey(seed), alpha, alpha,
                            dtype=jnp.float32))
idx = jax.random.permutation(jax.random.PRNGKey(seed + 1),
                             jnp.arange(N, dtype=jnp.int32))

Mix images and targets

Do not reseed global random state or substitute a different sampling sequence. Convert labels to one-hot vectors of width num_classes, then compute:

x_mixed = lam * x + (1 - lam) * x[idx]
y_mixed = lam * one_hot(y) + (1 - lam) * one_hot(y[idx])

The same idx and lam must be used for both. Every row of y_mixed sums to one.

2. Run the classifier

Extract non-overlapping square patches, visiting patch rows then columns. Flatten each patch by channel, row within patch, then column within patch. Project with w_proj, prepend cls_token to each image's sequence, and add pos_embed. Position 0 belongs to CLS; the remaining positions follow patch order.

Apply each block in order:

h = LN(seq)
q, k, v = h @ w_q, h @ w_k, h @ w_v
z = seq + multi_head_attention(q, k, v) @ w_o
seq = z + GELU(LN(z) @ w_mlp1) @ w_mlp2

LN normalizes each token across its D features using population variance: (token - mean) / sqrt(variance + 1e-5), with no learned scale or bias. Attention splits features into consecutive heads of width D/num_heads, scales scores by sqrt(D/num_heads), and applies softmax over key positions. Concatenate head outputs in their original order. All tokens in an image attend to all tokens in that image, with no causal mask.

Use the tanh approximation:

GELU(u) = 0.5 * u * (1 + tanh(sqrt(2/pi) * (u + 0.044715 * u**3)))
logits = seq[:, 0, :] @ w_head

There are no biases, dropout or final normalization. With zero blocks, logits depend only on cls_token, pos_embed[0] and w_head.

3. Compute gradients and update

Use cross-entropy with the soft targets. Sum over classes, then average over the N examples. Its gradient with respect to logits is:

loss = -sum(y_mixed * log_softmax(logits)) / N
dlogits = (softmax(logits) - y_mixed) / N

Backpropagate through the classifier, every block, the patch projection, the CLS token and the position embeddings. Treat the mixed images and targets as fixed data. Sum gradients for parameters shared across examples or positions. Use stable softmax computations by subtracting each row's maximum first.

Compute all gradients using the original parameters, then update each parameter with new_parameter = parameter - lr * gradient. Do not use automatic differentiation, backward(), jax.grad, or an optimizer helper in either framework. Cache forward activations for the manual backward pass.

Return value

Return one flat tensor, preserving the input floating-point dtype and device. Flatten each updated parameter in row-major order, then concatenate in this order:

[w_proj_flat, cls_token_flat, pos_embed_flat, blocks_weights_flat, w_head_flat]

Update all five parameter groups, including CLS and position embeddings. When lr=0, return the original parameters concatenated in that same order.

See mixup: Beyond Empirical Risk Minimization for the augmentation method.