Section 8.4
The Transformer

We are almost ready to present the complete transformer architecture. There is just one more component that we need to discuss, layer normalisation. Unlike batch-norm that we saw in earlier lectures, layer normalisation is applied to feature values within a single sample. Samples within a batch are treated independently and thus the operation can be applied at test time without violating any independence assumptions or leaking information from one sample to another.

Let \(y \in \reals^n\) and \(x \in \reals^n\) be the input and output for a single sample processed by layer normalisation. Then,

where \(\mu\) and \(\sigma\) are the mean a standard deviation of \(x_i\), respectively. Quantities \(\gamma\) and \(\beta\) are learnable parameters. We have written the expressions in vector form for simplicity. Extending to tensor feature maps is straightforward, and performed by computing the mean and standard deviation over every element in the feature map (but still treating each sample in the batch independently).

We have now seen all the building blocks used in the transformer model. It’s time to put the blocks together. The full transformer architecture [103] is shown in Figure 114. Originally designed for language translation, the transformer accepts two sequences, an input sequence and a target sequence.¹ These can be two different languages, for example, English and French. The target sequence may be a partial sequence, and just like language models discussed in the previous lecture, the goal is to predict the (probability distribution over the) next token. The important thing about transformers is that the prediction is conditioned on a very long sequence of previous tokens, not the simple five or six word sentences that we saw previously, making them very powerful AI language models.²

The input sequence goes through various layers of multi-head scaled dot-product attention. Here queries, keys and values are obtained by self-attention. The output of each multi-head attention block is added back to the input via a skip connection and layer normalisation applied. Further processing by a two-layer perceptron with ReLU activation function and residual connection, followed by layer normalization then prepares the output for the next layer. This stack, typically consisting of 6–12 blocks is called the transformer encoder.

On the target sequence side, similar processing is applied, in what we call the transformer decoder. In the decoder stack, each block consists of masked multi-head attention to ensure that generated sequences cannot look forward in time. The block also uses cross-attention in a second stage of multi-head scaled dot-product attention processing, the keys and values being obtained from the corresponding layer from the encoder side. The same architecture of a residual single-layer perceptron and layer normalisation is used to produce the output for the next block.

At the top of the decoder is a linear layer followed by softmax. This gives output probability vectors associated with each token. As we will see, the distribution is typically trained to predict the next token in the target sequence, but other training objectives are also possible.

Since the introduction of the original transformer model, a large number of architectural variations and improvements have been suggested in the literature and industry. State-of-the-art large language models are decoder-only architectures with no cross-attention. They include features such as group query attention (GQA) and mixture-of-experts (MoE) to reduce inference costs and improve model expressiveness.

8.4.1 Training and Inference

Training of transformers for language translation works in the same way as training a RNN language model. The first language is tokenized to produce an input sequence for the encoder, \(\langle w_1, \ldots, w_m, \texttt{<pad>}, \ldots \rangle\), which may include some special padding tokens if the output sequence is longer than the input. Likewise, the second language is tokenized to produce a target sequence for the decoder, \(\langle \texttt{<sos>}, v_1, \ldots, v_n, \texttt{<eos>} \rangle\). The output of the decoder is then trained to minimize the next-token cross-entropy loss where \(t\)-th output after passing through softmax is a probability distribution of the next work given the full first language sequence and the partially completed second language sequence up to position \(t-1\), i.e., \(p(v_t \mid w_1, \ldots, w_m, v_1, \ldots, v_{t-1})\).

There is another way to train transformers for tasks such as (conditional) language generation (e.g., ChatGPT), where we only require the decoder stack. This proceeds as above but without any input tokens or cross-attention, i.e., with the encoder removed. Other variants include bi-directional prediction and masked token prediction where we predict a token somewhere in the middle of a sequence. The resulting models are then used as generic text encoders for some downstream task.³

Inference for language translation or language generation is done using auto-regressive next token prediction as we did with RNN language models. Here the target sequence is updated each time a new token is sampled. The input sequence (sometimes called the prompt) is left unchanged.

Sometimes we are interested in using transformer models as classifiers (e.g., to classify the sentiment of a piece of text or, as we will see later, classifier an image that has been appropriately tokenized). Here we introduce a special <cls> token, usually added at the start of the sequence. We only care about the output associated with this token, which is further processed through a multi-layer perceptron head and trained using cross-entropy loss. Rather than auto-regressive prediction, the transformer now behaves as a sequence classifier as shown in Figure 115.

8.4.2 Prompt Engineering

Large language models (LLMs) excel in domains where reasoning over language is central to the task and where there is abundant text or text-like data (e.g., software source code). Since they sample outputs, the domain should also be tolerant to probabilistic outputs rather than guaranteed correctness. For example, LLMs have been successfully deployed as summarization agents for large bodies of text, code generation and completion agents, customer support and helpdesks, and in training and education. Domains where LLMs still struggle include numerical and algorithmic computation, formal verification, long-horizon planning and reasoning, regulated decision making (such as in medical, legal, and financial applications), and real-world physical interactions.

To be effective the prompt given to an LLM needs to be crafted carefully since the way a prompt is frames strongly influences the LLMs response. Agents build on pre-trained LLMs include a system prompt that is prepended to the user prompt. Both are designed and refined to encourage the desired behaviour and reduce runtime costs, the former by the company deploying the agent and the later by the user. Typical prompts include a role, a description of the task, context and constraints, the required output format, and optimal examples in some applications. For example,

You are a computer science professor. Explain the concept of back propagation in neural networks.
Assume the reader is a third-year university student. Use simple language and equations to explain
the concept. Respond in four paragraphs or less.

Chain-of-thought prompting or asking to refine previous answers often improves performance on reasoning tasks:

Solve the problem step by step.

Note that prompting cannot fix weak models and what works for one model may not work for others. With the advent of LLM-based agents in a number of domains, prompt engineering has emerged as an important skill for agentic application developers.

8.4.3 Dealing with Tabular Data

Transformers were not originally designed for tabular (or any non-sequential) data. However, there are a few techniques for adapting them to processing tables, such as representing rows in the table as sequences or tokens and adding column embeddings (similar to positional encoding) to tokens. In addition, some specialized architectures have been proposed for handling tabular data [34, 49]. Further discussion of these is beyond the scope of this course.