Approaches to biomedical knowledge

Overview

Important: - https://arxiv.org/pdf/2305.17026

• Transformer models (recap)
• Decoder models
• Multihead attention
• Output generation
• Training LLMs
• KV Cache

Transformer models

Transformer models

• The original paper, “Attention is all you need…”, described an encoder-decoder transformer
• This architecture is well suited for sequence-to-sequence tasks such as machine translation
• Last week, we focused on an introduction to the encoder model
• This week, we will provide additional detail and focus mainly on decoder-only models and on training

Decoder models

Recap

• Recall encoder algorithm from the previous lecture:

Multihead attention

Multihead attention

• Instead of performing a single attention function with \(d_{hidden}\)-dimensional keys, queries, and values
• It is beneficial to linearly project the queries, keys, and values \(h\) times with different linear Projections
• The results are concatenated and once again projected resulting in the final values

\[ \mathrm{MultiHead}(Q,K,V) = \mathrm{Concat}(\mathrm{head}_1, \ldots,\mathrm{head}_h)\mathrm{W}^{O} \] where \[ \mathrm{head}_i = \mathrm{Attention}(QW_i^Q, KW_i^K, VW_i^V) \]

• In the original paper, the authors use \(h=8\) (8 heads)
• Then, \(d_k=d_v=d_{hidden}/h = 64\) because \(d_{hidden} =512\) (In the paper \(d_{hidden}\) is called \(d_{model}\))
• dimension check: \(W_i^Q\) and \(W_i^K\) are \(d_{hidden}\times d_k\), \(W_i^V\) is \(d_{hidden}\times d_v\), and \(\mathrm{W}^{O}\) is \(hd_v \times d_{hidden}\)

Multihead attention

MHA step by step

• Gray background: Identical to single-head attention case

Output generation

• Unlike input generation, output generation is a sequential process
  • The answer is generated one token at a time
  • Each generated token is appended to the previous input
  • Process repeated until we get an \(\langle \mathrm{EOS}\rangle\) token¹ or the maximum length is reached

Autoregressive text generation

Text generation: step by step

Token decoding: Deterministic Methods

• Deterministic methods may lead to model degeneration
  • output becomes “unnatural”
  • output is marked repetitive and overly predictable language.
  • text lacks variety and fails to reflect natural human expression.

Token decoding: greedy search

• Greedy search selects the word with the highest probability as its next word at each timestep

\[ \arg \max_w P(w\mid w_1:w_{t-1}) \]

• Greedy search continues in this way until reaching either an end-of-sequence (EOS) token or a maximum timestep \(T\).
• Greedy search may overlook high-probability words that are accessible only after a lower-probability word.
• As a result, greedy decoding does not formally guarantee a global optimum of the decoding objective, as its choices are only locally optimal.

Token decoding: Stochastic methods

• Stochastic approaches introduce randomness during the decoding process, leading to more diverse and natural text generation.

Token decoding: Top-\(k\) sampling

• The \(k\) most probable next words are selected¹
• The probability of the \(k\) top words is redistributed to form a new distribution, i.e., the probabilities of these k tokens are then normalized to sum to 1, resulting in a truncated distribution.
• A token is then randomly sampled from this distribution and appended to the current sequence.
• This process is repeated iteratively until a termination condition is satisfied.

Token decoding: Top-\(p\) (Nucleus) Sampling

• Selects the smallest set of words whose cumulative probability meets or exceeds a predefined threshold \(p\).
• Rather than sampling exclusively from the top \(k\) most probable words, Top-\(p\) redistributes the probability mass across this dynamic set.
• This adaptive approach allows the size of the word set (i.e., the number of included words) to increase or decrease based on the shape of the next-token probability distribution.
• Rest of approach is identical

Token decoding: Temperature Sampling

• The core idea of temperature sampling is to control the “sharpness” of the probability distribution by introducing a temperature parameter \(t\).
• This parameter is applied in the softmax function after the transformer’s final layer to compute token probabilities.
• The temperature t directly influences the level of randomness in the sampling process, with higher values increasing randomness and lower values reducing it.

Token decoding: Temperature Sampling

• The core idea of temperature sampling is to control the “sharpness” of the probability distribution by introducing a temperature parameter \(t\).
• This parameter is applied in the softmax function after the transformer’s final layer to compute token probabilities.
• The temperature \(t\) directly influences the level of randomness in the sampling process, with higher values increasing randomness and lower values reducing it.¹

\[ P(x_i) = \frac{\exp(z_i / t)}{\sum_j \exp(z_j / t)} \]

• \(t=1\): standard softmax (no change)
• \(t<1\): sharper distribution — closer to greedy/argmax as \(t\to 0\)
• \(t>1\): flatter distribution — closer to uniform as \(t\to\infty\)

Temperature: worked example

• Let the following be the logits for the next token after “The cat sat on the”:

• At \(t=0.5\): “mat” dominates (≈0.86) — the model behaves almost deterministically
• At \(t=1.0\): the raw model distribution (≈0.64 for “mat”)
• At \(t=2.0\): much flatter (≈0.45 for “mat”) — “couch” and “floor” become far more likely than under the original distribution

Combining temperature with top-\(k\) / top-\(p\)

• Temperature and truncation (top-\(k\), top-\(p\)) are typically applied together, in sequence:
  1. Scale logits by \(1/t\)
  2. Apply softmax to get the temperature-adjusted distribution
  3. Truncate: keep only the top-\(k\) tokens, or the smallest set whose cumulative probability exceeds \(p\) (nucleus sampling)¹
  4. Renormalize the remaining probabilities to sum to 1
  5. Sample from the renormalized distribution
• Example: applying top-\(k=2\) after \(t=0.5\) scaling keeps only {mat: 0.862, floor: 0.117}, renormalized to {mat: 0.880, floor: 0.120} — “bed” and “couch” are excluded entirely
• Intuition:
  • temperature reshapes how peaked the whole distribution is
  • top-\(k\)/top-\(p\) decide how much of the tail is eligible for sampling at all.
• Many chat APIs expose both temperature and top_p as independent parameters.

Training LLMs

• The next sections will cover several aspects of how to train LLMs
• For the most part, LLMs are synonymous with decoder-only transformers; from now on we will say LLM for conciseness.

Training LLMs: Causal (Masked) Self-Attention

• This would make the training loss trivially low — the model learns nothing about actually predicting the future, since the future is already visible.
• We need a mechanism that lets us train on the full sequence in parallel (for efficiency), while still preventing this kind of “leakage.”

Masked Self-Attention

• Recall: \(\text{Attention}(Q,K,V) = \text{softmax}\left(\dfrac{QK^\top}{\sqrt{d_k}}\right)V\)
• We modify this by adding a mask matrix \(M\) to the scores, before the softmax step:

\[ \text{MaskedAttention}(Q,K,V) = \text{softmax}\left(\frac{QK^\top}{\sqrt{d_k}} + M\right)V \]

• \(M\) is an \(n\times n\) matrix, defined entrywise as:

\[ M_{i,j} = \begin{cases} 0 & j \le i \quad \text{(allowed: current or past token)} \\ -\infty & j > i \quad \text{(disallowed: future token)} \end{cases} \]

• Adding \(-\infty\) to a score, then applying softmax, transforms the entry tp exactly \(0\).

Worked Example: The Mask Matrix

• For our 4-token sentence “I am an automaton” (\(n=4\)), the mask matrix is:

\[ M = \begin{bmatrix} 0 & -\infty & -\infty & -\infty \\ 0 & 0 & -\infty & -\infty \\ 0 & 0 & 0 & -\infty \\ 0 & 0 & 0 & 0 \end{bmatrix} \]

• Row \(i\) corresponds to the query token; column \(j\) to the key token.
• Row 0 (“I”): only column 0 is unmasked — the first token can only attend to itself.
• Row 3 (“automaton”): all four columns are unmasked — the last token can attend to the entire sequence, since nothing comes after it.
• Causal masking is often just called “triangular masking.”

Worked Example: From Raw Scores to Masked Attention Weights

Suppose the raw attention scores \(S =\frac{QK^\top}{\sqrt{d_k}}\) for our sentence. - The masking procedure is as follows

• Row “I”: weight \([1.000, 0, 0, 0]\) — with only itself unmasked, softmax trivially assigns it all the probability mass.
• Row “am”: weight \(\approx[0.310, 0.690, 0, 0]\) — splits attention only between “I” and “am”.
• Row “automaton”: weight \(\approx[0.095, 0.156, 0.233, 0.517]\) — the only row using the full softmax, since nothing is masked.
• Notice the upper triangle is exactly \(0\) in every row

Pretraining

• Some ML tasks are similar. For instance
  • Spam detection
  • sentiment Analysis
  • machine translation
• All three tasks require the model to understand natural language
• The transfer learning approach reused a trained model for a new task by “tuning” the previously trained model
• LLMs take this approach to an extreme: Pretraining

Pretraining

• Pre-training is the foundation for the power of Large Language Models (LLMs)
• Train by predicting the next work (token) on large collections of text data
• Most LLMs are trained on general data (“all” web pages, “all” books, etc.)
• Conversation text may help LLMs get better at chatting (train on social media data, e.g., Reddit)
• Books: May help LLMs “understand” complicated topics better than webpages/social media
• Stack exchange/GitHub: Help LLMs “understand” code

Pretraining

• Size of training data

Pretraining: Size matters

• Language modeling performance improves smoothly as we increase the model size, datasetset size, and amount of compute2 used for training.
• Coversely, the authors found very weak dependence on many architectural and optimization hyperparameters.

Optimal model size

• On the other hand, what if we have a fixed compute budget and cannot increase the number of parameters or the amount of data without limit?
• The authors investigatef the optimal model size and number of tokens for training a transformer language model under a given compute budget.
• They found that current large language models are significantly undertrained, a consequence of the recent focus on scaling language models whilst keeping the amount of training data constant.

Pretraining: Learning parameters

• The model doesn’t predict the whole sentence at once — it predicts one conditional distribution per position, and the chain rule of probability guarantees these conditionals multiply out to the joint probability of the full sequence.
• This factorization is exactly why causal masking (Section 2) is needed: each factor \(P(x_t \mid x_1,\ldots,x_{t-1})\) must only depend on the past, or the factorization isn’t valid.

The Loss Function

• Training uses cross-entropy loss between the predicted next-token distribution and the true next token.
• For a single position \(t\), with true token \(x_t\) and predicted probability \(\hat{P}(x_t \mid x_{<t})\):

\[ \mathcal{L}_t = -\log \hat{P}(x_t \mid x_1,\ldots,x_{t-1}) \]

• The total loss for a sequence is the sum (or mean) over all positions:

\[ \mathcal{L} = -\sum_{t=1}^n \log \hat{P}(x_t \mid x_1,\ldots,x_{t-1}) \]

Teacher forcing

• Teacher forcing: during training, the true previous tokens are fed as input at every position — not the model’s own (possibly wrong) predictions.
• This is what makes it possible to compute the loss for all \(n\) positions in a single parallel forward pass, using the causal mask from Section 2 to prevent each position from seeing its own answer.
• At inference time there is no ground truth to feed back in — this is exactly the “Append and reprocess” loop from our decoding-steps diagram, where each new token is the model’s own prediction, fed back as input for the next step.

Worked Example: Cross-Entropy Loss

Continuing “I am an automaton” — suppose after “I am an” the model’s predicted distribution over the next token is:

• The true next token is “automaton,” with predicted probability \(0.55\).
• The loss at this position is \(-\log(0.55) \approx 0.598\).
• If the model had instead assigned “automaton” a probability of, say, \(0.01\), the loss would jump to \(-\log(0.01)\approx 4.6\) — cross-entropy penalizes confident wrong predictions much more heavily than uncertain ones.

KV Cache

• Let us return to the input processing (prefill) for decoder-only transformers
  • Inputs (tokenized prompt) are embedded and positionally encoded
  • Multihead attention (MHA) computes keys and values (KV)
  • Lots of matrix multiplications, high usage of hardware accelerator on the GPU
• Output is sequential (one token at a time)
  • Generated token is appended to previous input, …, again and again
  • MHA process is repeated each time
  • sequential means we cannot parallelize much!

KV Cache

• Can we avoid recomputing KV values again and again?
• We only really need to compute KV values for the newly generated token
• The KV Cache stores the keys and values for all tokens we have previously generated

KV Cache

• When processing \(\mathrm{token}[k]\), we only need the \(k^{th}\) row of \(\mathbf{Q}\).
• When processing \(\mathrm{token}[k]\), we need the full \(K\) and \(V\) tensons, but we can mostly reuse the cached values
• Add the value for \(k+1\) as we go along

KV Cache: Size

• The Cache size (FP16) can be calculated as

\[ 2\times 2\times\text{batch-size}\times\text{seq-length}\times\text{num-layers}\times\text{embedding dimension} \] - We need to multiply by 2 because we have \(K\) and \(V\) - We need to multiply by 2 again because we have 16 bits (FP16) - Usually gigabytes of memory required - Many different attempts to shrink cache size to allow the batch size to be increased

Multihead latent attention

• Introduced by DeepSeek for version 2

• Some previous approaches to reduce size of KV Cache used fewer keys and values (CQA, MQA)
• MLA
  • Does not cache K and V
  • uses a low-rank representation that is learned during training and cached
  • also learn a projection matrix to shring K and V during training
  • Much lower KV cache usage (90%+ savings)
  • 5-6x inference speedup
  • Hgher accuracy than vanilla MHA

Multihead latent attention

LORA

• We will return to low-rank approximation (LORA) in a subsequent lecture

Sources

• Vaswani A et al. (2017) Attention Is All You Need. arXiv 1706.03762
• Turner RE (2026) An Introduction to Transformers. arXiv 2304.10557
• Allamar J (Blog) The Illustrated GPT-2 (Visualizing Transformer Language Models)
• Simon J Decoder-only inference: A step-by-step dive (video)