Add the Batch Axis Without Mixing Sequences

Batch-major attention preserves separate batch, position, and feature axes while sharing projection parameters. Trace batched matrix multiplication, mask broadcasting, padding, and an independence test that detects cross-sequence mixing.

A batch stores several sequences in one tensor so the same operations can be carried out together. For batch-major input,

X\in\mathbb R^{B\times T\times d_{model}}.

The first axis selects a sequence, the second selects a position within that sequence, and the third selects a feature. Keeping these meanings explicit is more useful than memorizing a framework call.

Projections Preserve Batch and Position Axes

Multiplication by a projection matrix contracts only the final feature axis:

\begin{aligned} Q=XW_Q&:B\times T\times d_k,\\ K=XW_K&:B\times T\times d_k,\\ V=XW_V&:B\times T\times d_v. \end{aligned}

The same parameter matrices are used at every position and in every batch item. Sharing parameters does not mix the data. Sequence $b$ still produces $Q_b$ , $K_b$ , and $V_b$ only from $X_b$ .

Transpose Only the Final Two Axes

For each batch item,

S_b=\frac{Q_bK_b^\top}{\sqrt{d_k}}.

In tensor notation,

(B,T,d_k)(B,d_k,T)\longrightarrow(B,T,T).

The transpose swaps the position and key-feature axes inside each batch item. Moving the batch axis would change which sequences are paired and can create a silent data-leakage bug even when a matrix multiplication still runs.

A Shared Causal Mask Can Broadcast

When every batch item has the same sequence length and causal rule, one mask with shape $(1,T,T)$ can be broadcast across $B$ :

A=\operatorname{softmax}_{last}(S+M),\qquad S:B\times T\times T,\quad M:1\times T\times T.

Broadcasting means that the same mask values are applied to every batch item; it does not combine their scores.

Sequences of different lengths are often padded to a common $T$ . A padding mask removes artificial padding positions, while the causal mask removes real future positions. These masks express different rules and may be combined before softmax. Every query row that is used must retain at least one allowed key, or softmax has no valid distribution to normalize.

Two Sequences Stay Independent

Let the batch contain

X_1=\begin{bmatrix}1&0\\0&1\end{bmatrix},\qquad X_2=\begin{bmatrix}2&0\\0&2\end{bmatrix}.

With identity Q/K/V projections, the first sequence forms scores from entries in $X_1$ only. Doubling $X_2$ changes the second sequence's dot products but cannot change $S_1$ , $A_1$ , or $H_1$ . This independence is a useful test: alter one batch item and assert that every output in the other items remains fixed.

Check independence across a batch

The function computes identity-projection attention separately for each batch item. Change the second sequence and verify that the first output remains fixed.

Q1. Derive a batched score shape

$Q$ and $K$ both have shape $(8,20,6)$ under the batch-major convention. What is the shape of $QK^\top$ when the transpose swaps only the final two axes?

Choose the score-tensor shape

(8, 20, 20)

(20, 8, 8)

(8, 6, 6)

Select one choice, then check.

Hint

For each batch item,

(20,6)(6,20)=(20,20)

Solution

The result has shape

(8,20,20)

. It contains one

20\times20

score matrix for each of the eight sequences.

Name Every Axis Before Reshaping

Many attention bugs are valid tensor operations with the wrong semantic axes. Before a transpose, broadcast, or matrix multiplication, write the axis names beside the shapes. Then test independence by changing one batch item. A passing shape check alone does not prove that batch boundaries were preserved.