Tensor Shapes in Transformers: Batch, Sequence, and Hidden Size

Identity statement: a tensor is an n-dimensional array of numbers, and the tensor carrying text through a transformer is 3-D with axes (batch, sequence, hidden).

When it matters: every time you read a model card, size a deployment, interpret a shape error, or reason about why a longer prompt costs more.

The critical asymmetry: hidden size is a property of the trained model and cannot be changed, while batch and sequence vary per request. A model with a hidden size of 4,096 has that hidden size forever; you cannot serve it "at 2,048" without retraining. But you choose the batch, and the caller's input determines the sequence length.

Rank, shape, and dimension — three words that are not synonyms

Shape talk goes wrong when these three get swapped, and they are swapped constantly in casual writing.

Because "dimension" carries both meanings, prefer axis when you mean the slot and width or length when you mean the count. "A 4,096-dimensional embedding" is the width sense; "reduce along the sequence dimension" is the axis sense.

A vocabulary ladder is worth having too, because the exam and the literature both use the lower-rank names freely:

L1 — The one-sentence picture

A transformer is a stack of identical-shaped rooms. Text enters as integers, the embedding layer turns each integer into a row of numbers, that block of numbers walks through every layer coming out the same shape it went in, and only at the very end does the shape change — into one score per possible next token.

If you only remember that, you can answer the shape-adjacent items. Everything below makes it precise.

L2 — The end-to-end trace

Follow one request end to end. Take a batch of 2 texts, 10 tokens each, through a model with hidden size 4,096 and vocabulary 32,000.

raw text            "The capital of France is..."
tokenize        →   (2, 10)              integer token ids, no hidden axis yet
embedding       →   (2, 10, 4096)        each id becomes a 4096-vector
+ positional    →   (2, 10, 4096)        shape unchanged
attention block →   (2, 10, 4096)        shape unchanged
feed-forward    →   (2, 10, 4096)        shape unchanged
... × N layers  →   (2, 10, 4096)        shape unchanged
output head     →   (2, 10, 32000)       hidden axis becomes vocabulary scores

Three things to take from that trace.

The embedding layer is where the hidden axis is born. Before it you have integers — token ids, shape (batch, sequence). After it you have vectors. This is the conversion of text to numbers that the whole field rests on.

The residual stream is shape-invariant. Every transformer block takes (B, S, H) and returns (B, S, H). That is not an accident of implementation; it is what makes the blocks stackable, and it is what lets residual connections add a block's output back to its input. When you read "24 layers," it means this identical shape passed through 24 times.

Only the last step changes the final axis, projecting hidden size to vocabulary size so you get one score per possible next token — the distribution from 01-01. Note where the next-token prediction lives: at the last sequence position of that output. During generation you read row (b, S-1, :) and discard the rest.

L2 — Where the third axis splits for multi-head attention

Inside attention, the hidden axis is temporarily reshaped into heads: (B, S, H) becomes (B, heads, S, H/heads). With H = 4,096 and 32 heads, each head works in 128 dimensions. Then the heads are concatenated back to (B, S, H). You need to recognise this so a 4-D tensor in a diagram does not surprise you. You do not need to reproduce the attention arithmetic — candidate reports consistently name detailed attention math as depth that did not appear on the exam.

The reshape is worth one more sentence because it explains a config constraint you will see enforced: hidden size must be divisible by head count. 4096 / 32 = 128 works; a hypothetical 4,096 with 30 heads does not, and a library will refuse it. When a config fragment lists both numbers, dividing them tells you the per-head width immediately.

L3 — The scaling consequence you can predict from the shape alone

The shape tells you the cost. Activation memory scales with batch × sequence × hidden, so it is linear in each. But attention compares every token to every other token, producing an intermediate of shape (B, heads, S, S) — quadratic in sequence length.

That single asymmetry explains a family of exam-adjacent facts: why doubling the batch roughly doubles memory, why doubling the prompt length can roughly quadruple the attention term, why context windows were historically small and expensive to extend, and why long-context efficiency is an active engineering problem rather than a config flag.

Make it concrete with an illustrative construction. Take B = 1, 32 heads, hidden 4,096, and compare two sequence lengths:

S = 512     residual stream   1 × 512 × 4096      =      2.10M values
            attention scores  1 × 32 × 512 × 512  =      8.39M values

S = 2048    residual stream   1 × 2048 × 4096     =      8.39M values   (4× — linear)
            attention scores  1 × 32 × 2048 × 2048 =   134.2M values    (16× — quadratic)

Quadruple the sequence length: the residual stream grows fourfold, the attention score term grows sixteenfold. That is the whole story of why long context is expensive, and it fell out of reading shapes rather than from any measurement. (Production attention kernels avoid materialising that full S × S block, which is exactly why such kernels exist — but the compute the shape implies does not go away.)

This lesson stops at L3 on purpose. You will not be asked to write the reshape, tune a kernel, or derive attention complexity formally. Reading a shape and predicting a cost direction is the whole ask.

Every one of these is a number in a model config, and every pair of them gets confused by someone. This table is the highest-value asset on the page.

Two rows are worth reading against each other. Sequence length and context window are the same axis, one actual and one maximum — configs list the maximum as max_position_embeddings, and exceeding it truncates or errors. Hidden size and embedding dimension are the same kind of quantity in different products: a generative model's hidden size is internal and you never see it, while an embedding model's output dimension is the thing you must match to your vector index forever.

Here is a config fragment of the kind you meet on any model hub:

hidden_size:            4096
num_hidden_layers:      32
num_attention_heads:    32
max_position_embeddings: 4096
vocab_size:             32000

Read it off:

• Hidden size 4,096 — each token is a 4,096-number vector throughout the stack. Fixed forever.
• 32 layers — the (B, S, 4096) tensor passes through 32 shape-preserving blocks.
• 32 heads — attention splits into 32 subspaces of 128 dimensions each (4,096 ÷ 32).
• max_position_embeddings 4,096 — the sequence axis cannot exceed 4,096 tokens. This is the context window, and it is a different number from the hidden size that happens to match here. Conflating them is a real and easy error.
• vocab_size 32,000 — the output head produces 32,000 scores per position.

Now predict. You send a batch of 8 requests of 512 tokens each. The tensor entering layer 1 is (8, 512, 4096). The tensor leaving layer 32 is (8, 512, 4096). The logits are (8, 512, 32000), and for generation you care about (8, 4096) — the last position of each sequence, projected. If someone sends a 5,000-token prompt, it fails or is truncated, because 5,000 exceeds the sequence cap. None of that required running anything.

One more prediction from the same config, because it is the one people find surprising. How big is the logits tensor for that batch?

logits = 8 × 512 × 32000 = 131,072,000 values
at 2 bytes each (BF16)   ≈ 262 MB

A quarter of a gigabyte of scores, of which — during generation — you use 8 rows. That waste is the direct motivation for the implementation trick of only projecting the final position, and it is a nice demonstration that shapes tell you where the cost is before you profile anything.

Shape errors are the most common failure you will hit doing this work, and they are readable once you know which axis the library is complaining about. Here are the four archetypes, written as an illustrative reconstruction rather than a quote from any specific library version:

The reading habit that turns all four into thirty-second fixes: name the axis, then name where it should have come from. Axis 0 is batch and comes from how you assembled the request. Axis 1 is sequence and comes from the tokenizer. Axis 2 is hidden and comes from the model config. If the number in the error is a vocabulary-sized number, you are past the output head; if it is a small power of two times a hundred-ish, you are looking at a per-head width.

Padding deserves its own note because it is the practical reason batches are annoying. A batch must be rectangular — every sequence in a (B, S, H) tensor has the same S. Real requests do not. So short sequences are padded to the longest one in the batch and an attention mask marks which positions are real, so the padding cannot influence the result. Two consequences: batching wildly different lengths wastes compute on padding, and length-bucketing requests is a standard throughput optimisation.

The shape-preserving rule holds inside the stack. It stops holding at three boundaries you will meet constantly.

Embedding models collapse the sequence axis. A text embedding model reads (B, S, H) and returns (B, H) — one fixed-length vector per input text, regardless of how many tokens went in. That collapse is called pooling, and it is why a two-word query and a 400-word paragraph produce vectors of the same length and can be compared at all. That comparison is 01-04, and the vector-database lessons depend on it entirely. The three pooling styles you should be able to name: mean pooling over positions, taking a designated [CLS]-style position, and last-token pooling — all doing the same job of turning a variable-length matrix into one vector.

Classification heads collapse to class count. An encoder-only classifier goes (B, S, H) → pool → (B, H) → head → (B, C) for C classes. Sentiment with three classes returns (B, 3). Note the family resemblance to the generative output head, which is just the same move with C = vocabulary size.

The KV cache is a fourth shape entirely. During generation, keys and values from previous positions are cached so they need not be recomputed. Its shape is roughly (layers, 2, B, heads, S, H/heads) — the 2 being keys and values. What you should take from that is not the exact ordering but the multiplication: the KV cache grows with layers × batch × sequence length, which is why long conversations and large batches, not big weight files, are usually what exhausts a serving GPU's spare memory. The weights are a fixed cost from 01-02; the KV cache is the variable one.

Tensor shapes serve objectives 1.5 (ML fundamentals) and 1.10 (using Python packages such as NumPy and Keras to implement analyses) — you cannot use NumPy meaningfully without shape fluency. Core ML is 30% of the blueprint.

You will almost certainly not see an item that says "what is the shape of the residual stream." You will see items where the shape is the reason the answer is the answer:

• Why a longer context costs disproportionately more (quadratic attention term).
• Why batch size is a throughput lever and hidden size is not a lever at all.
• Why an embedding model outputs a fixed-length vector regardless of input length — a pooled (B, H), with the sequence axis collapsed.
• Why increasing batch size raises throughput but also raises per-request latency.
• Why you cannot swap in an embedding model with a different output dimension without re-indexing your corpus.

How the question tends to be phrased

Two shapes recur. The first is a lever question: "throughput on a serving endpoint is too low; which change is most likely to help?" — where the keyed answer is the runtime lever (batch size, or batching strategy) and the distractors offer architecture numbers you cannot change without retraining. The second is a cost-direction question: "a team doubles the average prompt length; what happens?" — where the keyed answer is that cost grows faster than linearly because of attention, and the distractors offer "unchanged," "doubles exactly," or "depends only on output length."

What the distractors typically look like

Expect hidden size offered as something you tune per request; context window offered as a fix for a model not knowing something (an empty bigger window changes nothing — that is the 01-02 point); vocabulary size offered where hidden size belongs; and "reduce the number of layers" offered as a serving-time option, which it is not for a given checkpoint.

Calibration, not fact: candidate reports describe the exam as general-level and specifically name deep attention math as depth that did not appear. NVIDIA publishes no item-level detail. So treat this lesson as enabling rather than tested. Its value is that four or five later answers stop requiring memorisation once you can see the shape.

For a correct implementation, no — batching is a throughput arrangement, not a semantic one. Each sequence in a batch is processed independently; nothing in a transformer lets row 3 attend to row 5.

Two honest caveats, both worth knowing because they cause real confusion. First, floating-point arithmetic is not perfectly associative, so different batch sizes can take different kernel paths and produce results that differ in the last bits. That is numerical noise, not a behaviour change, though it can make "identical" runs fail an exact-match test. Second, if padding is not masked correctly, batch composition genuinely does change the answer — which is a bug, and one that only appears when you batch inputs of different lengths.

The practical rule: batch size is a cost and latency lever. If a scenario claims batching changed the meaning of a response, suspect masking.

Because hidden size is baked into every weight matrix in the checkpoint. Every attention projection is a d × d matrix and every feed-forward matrix is d × 4d; change d and none of the stored numbers have the right shape any more. There is no operation that "widens" a trained model — you would be initialising a new, differently shaped model and training it from scratch.

This is the concrete reason for the architecture-versus-runtime split in section 3. Hidden size, layer count, head count, and vocabulary are decisions frozen at pretraining. Batch size, sequence length used, and precision are the things you actually control at deployment. When an exam scenario asks what a deployment team can change, the answer lives in the runtime column.

Because pooling collapses the sequence axis. The model still builds (B, S, H) internally — a 3-token query produces 3 hidden vectors, a 300-token paragraph produces 300 — and then reduces along the sequence axis to a single (B, H) vector per text.

That is not a detail; it is what makes semantic search possible. Two texts of wildly different length end up as two vectors of identical length in the same space, and identical length is the precondition for taking a dot product or a cosine between them. It also explains a known limitation: squeezing a long document into one fixed-length vector loses detail, which is exactly why RAG pipelines chunk documents before embedding them rather than embedding a whole document at once.

• Text flows as (batch, sequence, hidden): how many texts, how many tokens, how many numbers per token.
• Hidden size is fixed by the architecture. Batch and sequence vary per request. That split decides what a deployment team can actually change.
• Transformer blocks are shape-preserving: (B, S, H) in, (B, S, H) out, N times. That is what makes them stackable and what residual connections require.
• The embedding layer creates the hidden axis; the output head replaces it with vocabulary scores.
• Multi-head attention temporarily reshapes hidden into (B, heads, S, H/heads) and concatenates back; hidden size must be divisible by head count.
• Memory is linear in batch and hidden but quadratic in sequence length through attention — the one cost fact worth carrying forward.
• Pooling collapses the sequence axis to (B, H), which is what makes embeddings comparable across texts of any length.
• The KV cache, not the weight file, is usually what consumes spare serving memory as conversations and batches grow.

You can now say what shape a token's representation has. What you cannot yet do is compare two of those representations — and comparison is the operation underneath embeddings, semantic search, vector databases and retrieval scoring, which together account for a substantial slice of the exam.

Next: 01-04 Vectors, dot products, and cosine similarity — the one piece of linear algebra this course asks you to know airtight, because six later groups of lessons reduce to it.