The Stack: From Token Order to Next-Token Probabilities

The Transformer is often described through its attention equation, but that equation alone cannot turn tokens into predictions. This chapter traces the complete stack that makes the operation usable: what the encoder computes, what the decoder adds, how information crosses their boundary, and how token order and next-token probabilities are represented. The goal is a system-level view in which each component has a distinct responsibility rather than a catalogue of parts.

A useful way to read the architecture is to follow two questions at every boundary. What representation enters this component, and what information is allowed to leave it? The encoder builds a contextual representation of the input. The decoder builds an output representation under a causal restriction while consulting the encoder. Feed-forward layers transform each position after attention has mixed information; embeddings and positional encodings make symbols and order available; the final projection turns the decoder’s representation into a probability distribution.

The two stacks and their interfaces

The Transformer keeps the encoder–decoder shape familiar from earlier sequence-to-sequence systems, but changes the machinery inside the stacks. The encoder receives a sequence of input symbols and maps it to a sequence of continuous representations. “Continuous” here means vectors of numbers rather than discrete words or tokens. The decoder receives the already generated portion of the target sequence and uses the encoder’s representations while producing the next target symbol.

In the configuration described in the paper, the encoder is a stack of six identical layers. Each layer has two sublayers. The first is multi-head self-attention, which allows every source position to build context from the output of the preceding layer at all source positions. The second is a position-wise feed-forward network, which transforms each resulting position independently. The same kind of operation is applied at every position in a layer, but the values at those positions differ because their contextual inputs differ.

Each sublayer is surrounded by a residual connection and followed by layer normalization. In the paper’s notation, the output has the form LayerNorm(x + Sublayer(x)). The residual path carries the sublayer’s input into the combined representation, while normalization is applied after the addition. It is useful to treat this as a repeated layer interface: a sublayer receives a representation of the common model width, contributes an update, and returns a representation that remains compatible with the next sublayer. The paper specifies this organization; it does not, in this description, isolate the causal contribution of residual connections or normalization in a separate experiment.

The decoder has the same repeated-layer idea but one additional relationship. Each decoder layer contains masked multi-head self-attention, multi-head attention over the encoder output, and a position-wise feed-forward network. The first attention operation relates a target position to the permitted target prefix. The second lets that target-side representation query the source-side representations produced by the encoder. Residual connections and layer normalization surround each of the decoder’s sublayers as well.

This gives the decoder three different kinds of context. Its masked self-attention provides information from earlier target positions without exposing later answers. Its encoder-decoder attention provides access to the complete encoded input. Its feed-forward sublayer then applies a nonlinear transformation at each position. The decoder does not receive a single undifferentiated “context” object. It receives representations shaped by separate interfaces, each with a different information boundary.

The paper uses a common model dimension, d_model = 512, for the sublayers and embeddings in the base model. That common width is an interface contract: attention outputs, feed-forward outputs, residual additions, and embedding representations can be combined because they have the same model dimension. The decoder’s extra attention sublayer does not create a separate output language; it contributes another update to the same kind of representation.

The direction of information flow matters as much as the list of components. The encoder’s source representations can be built for all input positions using self-attention. The decoder’s target-side representations are constrained by the causal mask, but they can still use encoder output through cross-attention. At the end of a decoder layer, the representation at a target position reflects a permitted target prefix, source information selected through encoder-decoder attention, and position-wise transformations applied along the way.

A layer does not replace a position with a single summary of the whole sequence. It returns one updated vector per position. In the encoder, the vector at each source position carries the result of that position’s attention-weighted interaction and its subsequent feed-forward transformation. The next encoder layer receives all of those updated vectors and can calculate another set of relationships. Depth therefore gives the model repeated opportunities to refine the same positions, while the attention operation determines which other representations can influence each refinement.

The decoder has a related but asymmetric flow. A target position first receives an update from masked self-attention, so its target-side context is limited by the causal rule. The resulting representation then supplies a query to encoder-decoder attention, where the source-side keys and values come from the encoder output. The source representation crosses into the decoder through that attention interface; future target symbols do not cross the masked target-side boundary. This distinction is why the encoder output is better understood as a sequence of source-side memory vectors than as one fixed sentence summary, and why each decoder position can select a different mixture of source information.

The common model width makes these handoffs composable. Attention may use internal projections, and the feed-forward network may temporarily expand a vector to its inner dimension, but the sublayer output returns to the shared width before it is added to the residual path. A decoder layer can therefore combine its target-side update, its source-side lookup, and its position-wise transformation without changing the representation type expected by the next layer. The architecture’s global behavior comes from composing these compatible local interfaces across the stack.

classDiagram class TokenAndPosition class EncoderStack class EncoderLayer class EncoderSelfAttention class DecoderStack class DecoderLayer class MaskedSelfAttention class CrossAttention class FeedForward class ResidualNorm class OutputProjection class Softmax TokenAndPosition --> EncoderStack : supplies source vectors EncoderStack *-- EncoderLayer : repeats EncoderLayer *-- EncoderSelfAttention : mixes source context EncoderLayer *-- FeedForward : transforms each position EncoderLayer --> ResidualNorm : wraps each sublayer TokenAndPosition --> DecoderStack : supplies target-prefix vectors DecoderStack *-- DecoderLayer : repeats DecoderLayer *-- MaskedSelfAttention : preserves target causality DecoderLayer *-- CrossAttention : reads encoder output DecoderLayer *-- FeedForward : transforms each position DecoderLayer --> ResidualNorm : wraps each sublayer DecoderStack --> OutputProjection : produces next-token scores OutputProjection --> Softmax : produces probabilities TokenAndPosition --> OutputProjection : shares embedding weights

The diagram represents responsibilities and interfaces rather than a software class hierarchy. The encoder stack owns the transformation of source representations. The decoder stack owns target-side representation under a causal rule and receives the encoder’s output through cross-attention. The residual-and-normalization boundary is repeated around each sublayer, and the output projection converts the final decoder representation into scores that the softmax turns into probabilities.

That last boundary is also where the decoder’s temporal behavior becomes visible. During training, the known target sequence can be presented with a one-position offset and processed under a mask. During generation, the model predicts a token, adds that known token to the usable prefix, and predicts again. The architecture’s internal layers may be highly parallelizable for known positions, but the output interface still asks for a next-token distribution conditioned on what has already been generated.

The missing ingredients: local transformation, token interfaces, and order

Attention mixes information across positions, but mixing is not the only computation a representation needs. The Transformer places a feed-forward network after each attention sublayer. For an input vector x, the paper defines it as:

FFN(x) = max(0, xW₁ + b₁)W₂ + b₂

The first linear transformation expands the representation into an inner dimension, the ReLU sets negative intermediate values to zero, and the second linear transformation returns it to the model dimension. In the described base model, the input and output dimension is 512 and the inner dimension is 2048. The network is applied separately and identically at every position in a layer. “Identically” means that positions use the same layer parameters; it does not mean that they receive the same output. Each position arrives with a different context from attention, so the same nonlinear transformation can produce different results.

This division of labor is central. Attention supplies cross-position mixing: it changes a position by combining values associated with other positions. The feed-forward network supplies position-wise nonlinear transformation: it changes each position using its own vector after that mixing has occurred. A Transformer layer is therefore not a choice between global interaction and local computation. It composes them, with residual paths and normalization preserving a common interface between the updates.

The stack begins before either operation. Input and output tokens are discrete symbols, so the model first maps them into learned vectors of dimension d_model. These embeddings provide the numeric representations on which attention and feed-forward layers operate. The output side uses the same kind of embedding interface for the previously known target symbols, with the positional information added at the bottom of the decoder stack.

At the other end, a learned linear transformation maps the final decoder representation to scores over possible output tokens. A softmax converts those scores into next-token probabilities. The paper shares the same weight matrix between the two embedding layers and the pre-softmax linear transformation, and multiplies embedding outputs by √d_model. Weight sharing is a choice in the described model, not a requirement that follows from the definition of every Transformer. Its role here is to connect the token representation interface with the output scoring interface.

The word “sequence” creates one more requirement. A set of vectors by itself does not tell the model which vector came first. If the architecture had recurrence, order would be partly expressed by the order of state updates. If it had convolution, local neighborhoods and their arrangement would contribute positional structure. The Transformer uses neither recurrence nor convolution, so it adds positional encodings to the token embeddings at the bottoms of the encoder and decoder stacks.

The paper’s chosen positional encoding is fixed and sinusoidal. For position pos and dimension index i, its even and odd coordinates are defined as:

PE(pos, 2i) = sin(pos / 10000^(2i/d_model))

PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model))

The dimensions use different frequencies, with wavelengths forming a geometric progression. The result is added to the token embedding rather than placed in a separate stream. A representation entering the first encoder or decoder layer therefore contains both what symbol is present and where it occurs according to the chosen positional basis.

The motivation has two parts. First, the model needs some explicit signal that distinguishes order; otherwise an arrangement of tokens could be confused with another arrangement containing the same symbols. Second, the authors hypothesize that for a fixed offset, the encoding at one position can be represented as a linear function of the encoding at another position. That could make relative-position relationships easier for the model to use. This is a design rationale, not a guarantee that the representation automatically solves every positional problem.

The paper also tests learned positional embeddings instead of the sinusoidal functions and reports nearly identical results in its English–German development-set comparison. It chooses the sinusoidal version because it may allow extrapolation to sequence lengths longer than those encountered during training. The wording matters: the near-equivalence is a reported result under one experiment, while the longer-sequence extrapolation is a hypothesis motivating the choice. The paper does not present a systematic study that turns the hypothesis into a general guarantee.

Putting the interfaces together produces a complete path. An input token becomes a learned vector, receives a positional vector, and enters the encoder stack. Encoder self-attention mixes source information; the feed-forward network transforms each position; residual connections and normalization carry the representation between sublayers. The resulting encoder output becomes the source-side memory available to the decoder. On the target side, known output tokens are embedded and positioned, masked self-attention limits target-side information to the permitted prefix, cross-attention lets the decoder query the encoder output, and the feed-forward network transforms the result. The final decoder vector is projected to token scores and normalized into probabilities.

No single component owns the whole meaning of the sequence. Position encodings supply order but do not contextualize tokens. Attention supplies relationships but does not replace nonlinear per-position transformation. Feed-forward layers transform representations but do not create direct access to other positions. The output softmax creates a probability distribution, but it does not decide which future token will actually exist until the decoding procedure selects one and feeds it back as known context.

This division of responsibility also clarifies what the Transformer’s architectural claim does and does not remove. It removes sequence-aligned recurrence and convolution from the described representation architecture, changing how positions exchange information. It does not remove the need for order, the need to transform each position, or the causal boundary of next-token generation. The full stack makes global attention operational; the cost and parallelism of that stack can only be judged once its information paths are visible.

References