Transformers and Embeddings: Mapping Words to Cell Type Space

After completing this topic

In parts #2-4, we explored neural networks, which are essentially dials with 33.62 million parameters, and how to train them. In this part, we will learn about the Transformer, an architecture that reassembles these machines to be language-aware. We will cover five key components in order: tokenization, embedding, positional encoding, residual connection, and layer normalization.

The key operating principle, attention, will be covered separately in part #6. Here, we will simply indicate where attention will be placed and focus on the remaining structure.

Three Linguistic Headaches, Different from Images

In part #2, we looked at neural networks for pathology slides (images). Images are data that are relatively easy for neural networks to handle. The pixel values are already real numbers, the image has a natural 2D grid structure, and the neighborhood relationships are clear.

Language is fundamentally different. There are three headaches.

First, it is discrete. The word "cat" is not a real number, but a single discrete symbol. To handle this with a neural network, it must be represented numerically. However, assigning arbitrary real numbers, such as "cat = 0.42, dog = 0.51," creates meaningless relationships, such as "cats are 82% dogs."

Second, the order matters. "The dog bit the person" and "The person bit the dog" are completely different news stories. The meaning is reversed by changing the order of the words. Images are somewhat robust to rotation and translation, but language is destroyed by even a single change in order.

Third, the meaning changes depending on the context. As mentioned in part #1, the word "liver" can have multiple meanings, such as "liver (organ)," "the intensity of saltiness," or "-ed (past tense)." The meaning depends on the surrounding context.

The Transformer solves these three problems with the following components:

• Discreteness → Tokenization + Embedding
• Order → Positional Encoding
• Context dependence → Attention (Part #6)

In this part, we will first cover the first two problems.

Tokenization: Breaking Down Language into Discrete Units

Before feeding language into a neural network, we must first define a "unit." This is called tokenization.

Let's start with the simplest method.

Option A: Word-based. Split "The cat likes fish" into [The, cat, likes, fish]. Problem: The vocabulary of language is in the tens or hundreds of thousands. Moreover, new words (neologisms, technical terms, typos) continue to appear. If an unseen word appears during training, it cannot be processed.

Option B: Character-based. Split into [T, h, e, , c, a, t, , l, i, k, e, s, , f, i, s, h, ...]. The vocabulary is smaller (tens of thousands, including common Chinese characters), but each character has almost no meaning, so the sequence becomes very long. The computational cost explodes.

Option C: Subword-based - the modern standard. Frequently occurring sub-word parts are treated as a single token, and rare words are broken down further. If "cat" appears frequently, it is a single token, and if "metarepresentationism" is rare, it is [meta, representation, ism].

Algorithms that automatically learn Option C are Byte Pair Encoding (BPE) or SentencePiece. Most LLMs in practice use this family. The vocabulary size is typically 30,000 to 150,000.

Bio Analogy: When analyzing genomes, there is a similar problem. If we deal with single nucleotide (A, T, G, C) units, the sequence becomes too long, and if we deal with genes, we cannot handle new genes. In practice, we use a k-mer (e.g., a fragment of 6-8 nucleotides) as a sub-sequence unit. This is exactly the same idea as Option C in tokenization. In fact, recent genomic language models (e.g., Nucleotide Transformer) use exactly k-mer tokenization.

The result of tokenization is a sequence of integer IDs.

"The cat likes fish"
   ↓ Tokenization
[The, cat, likes, fish]
   ↓ Vocabulary lookup
[234, 1234, 5678, 9012]

Now, this sequence of integers needs to be converted into a sequence of real-valued vectors that the neural network can handle. This is embedding.

Embedding: Let's Create a Cell Type Space

Embedding is the process of mapping discrete tokens to high-dimensional real-valued vectors. Each token has a learned vector.

Mathematically, it is very simple. If there are V tokens in the vocabulary and the embedding dimension is d (e.g., 4096), there is an embedding matrix E of size V × d. The embedding of token ID i is the corresponding row E[i].

V = 100000   # Vocabulary size
d = 4096     # Embedding dimension
embedding = torch.nn.Embedding(V, d)

token_ids = torch.tensor([234, 1234, 5678, 9012])
vectors = embedding(token_ids)   # shape: (4, 4096)

This embedding matrix is initialized randomly at the beginning of training and is learned together with the neural network during training through backpropagation. After training, each token resides at a specific location in this 4096-dimensional space.

The geometry of this space is surprising. In a well-trained embedding space, semantically similar tokens are placed close to each other. And semantic relationships are reproduced by vector arithmetic.

Bio Analogy - UMAP space of scRNA-seq. If you have done single-cell RNA sequencing, you will get expression profiles of tens of thousands of genes for each cell. It is impossible to plot this directly, but by mapping it to a 2-50 dimensional embedding space using UMAP, t-SNE, or PCA, the cells form an interesting arrangement in that space.

• Cells of the same type (e.g., CD8+ T cells) are clustered together.
• Cells close in the developmental trajectory form a continuous trajectory.
• Direction vectors between cell types have biological meaning - for example, the direction from "naïve T cell" to "effector T cell" and the direction from "naïve B cell" to "plasma cell" are similar.

The LLM's token embedding space is exactly the same structure. Instead of cell types, words are arranged, and instead of gene expression axes, semantic axes of words (concreteness, gender, positivity/negativity, etc.) appear as directions in the embedding space.

An example of the amazing vector arithmetic in the embedding space (reproduced in actual word2vec/GloVe/LLM embeddings):

vec("queen") - vec("woman") + vec("man") ≈ vec("king")

By removing the "female" component from "queen" and adding the "male" component, we obtain a vector that is approximately "king." This is because the embedding space has learned the gender axis. It is important that this learned vector arithmetic is not programmed. The neural network discovers these relationships on its own from the training data.

Bio Analogy: Similar vector arithmetic is observed in the scRNA-seq embedding space.

vec("helper T cell") - vec("CD4+") + vec("CD8+") ≈ vec("cytotoxic T cell")

In the T cell embedding space, CD4/CD8 marker expression forms one axis, and moving along this axis changes the cell type from helper to cytotoxic. The reason that neural network embeddings and biological embeddings have the same conceptual structure is that they both compress the statistical relationships of the training data into a low-dimensional Euclidean space.

Solving the Order Problem: Positional Encoding

After embedding, each word becomes a vector. However, the second headache remains. The order of the words.

The Transformer (as we will see in detail in Part #6) treats the input sequence as an unordered set. That is, there is a risk that "The dog bit the person" and "The person bit the dog" will yield the same result.

To prevent this, we explicitly add positional information to the embedding. This is called positional encoding.

The simplest method: have a learned vector for each position and add it to the embedding.

pos_embedding = torch.nn.Embedding(max_length, d)
positions = torch.arange(len(token_ids))
final = token_embed(token_ids) + pos_embedding(positions)

This is called learned positional embedding. This method was used in the original GPT-2.

The original Transformer paper (2017) uses a different method: sinusoidal positional encoding. For each position p and each dimension i:

PE(p, 2i)   = sin(p / 10000^(2i/d))
PE(p, 2i+1) = cos(p / 10000^(2i/d))

A waveform that is calculated deterministically without parameters. The relationship between positions p and p+k can be expressed by trigonometric identities, making it easy for the model to learn relative positions.

Bio Analogy - circadian rhythm phase. In cell physiology, when expressing the time of day in a 24-hour period, we often use the combination of sin(2πt/24) and cos(2πt/24). This naturally places noon and midnight in opposite directions in vector space and moves morning and afternoon continuously. The sinusoidal PE of the Transformer is exactly this principle. By combining different periods (day, hour, minute), a more precise time representation is possible. The different periods of the PE (10000^(2i/d) is different for each dimension) automatically implement this multi-period representation.

Modern standard - RoPE and ALiBi. New methods that have emerged since 2020 have become the standard in practice.

• RoPE (Rotary Position Embedding) - Rotates the embedding vector by a certain angle depending on the position. Relative positions are naturally expressed as the difference in vector angles. Adopted by most modern LLMs such as LLaMA, GPT-NeoX, and Qwen. The elegance of being able to differentiate relative positions.
• ALiBi (Attention with Linear Biases) - Applies a penalty to the attention score proportional to the positional distance. Advantage of being able to generalize to some extent even to contexts longer than those seen during training.

Both methods have the advantage of being able to be extended to lengths longer than the context seen during training, which is why the sinusoidal PE is more widely used in practice. The expansion of the context window of large models such as GPT-3/4 and LLaMA (2K → 8K → 32K → 128K) is due to this series of positional encodings.

The Skeleton of the Transformer Block

After embedding and positional encoding, each word becomes a vector with positional information. These vectors now pass through a layered structure called the Transformer block.

The skeleton of the Transformer block (Pre-LN structure, modern standard):

Input x
   ↓
LayerNorm
   ↓
Multi-head Attention ←── (Detailed in Part #6)
   ↓
+ x                    ←── Residual connection
   ↓
LayerNorm
   ↓
FeedForward Network (FFN)
   ↓
+ (Previous result)          ←── Residual connection
   ↓
Output

This block is stacked 12 times in GPT-2, 96 times in GPT-3, and 100+ times in the latest large models. The deeper the layer, the more complex patterns can be captured, but training stability becomes difficult. The residual connection and layer normalization ensure this stability.

Residual Connection – A Shortcut Across Layers

A residual connection, or skip connection, is the process of directly adding the input to the output of a layer.

output = F(input) + input

F represents the actual transformation of the layer (e.g., attention, FFN).

Why is this important? It addresses the vanishing gradient problem discussed in Part 4.

During backpropagation, the gradient passing through 100 layers is multiplied by the local derivative at each layer, causing it to decay exponentially. After 100 multiplications, the gradient at most layers is effectively zero, preventing training.

Residual connections elegantly solve this problem.

∂output/∂input = ∂F/∂input + 1  ←── "+ input" ensures that 1 always remains

Thanks to this "+1," the gradient, even as it passes through the layers, maintains at least the original value. This allows gradients to propagate through even 100 layers.

The great success of the original Transformer paper, and even earlier, the success of ResNet (2015), are both due to this idea. This is the decisive reason why neural networks with 100+ layers can be trained today.

Biological Analogy – Cellular Homeostasis. Cells, even when exposed to external stimuli, do not immediately and completely change. They possess a homeostatic system that maintains a certain baseline state while gradually adjusting. This maintained baseline is conceptually similar to the "+ input" in a neural network. If a cell were to completely reset its state with each stimulus and start anew, it would lose stability. In cell signaling systems, strong signals also flow while maintaining the baseline. This aligns with the concept of residual flow in Transformers.

In a Transformer, each layer creates a "change" through attention or an FFN, and this change is added to the residual. When viewed from the perspective of a residual stream, it can be understood that the embedding forms a single stream from the first layer to the last, with each layer adding information to this stream. This perspective is a core framework in recent interpretability research.

Layer Normalization – Stabilizing Activation Distributions

Layer Normalization (LayerNorm) is an operation that normalizes the activation distribution of a layer at each step.

For each vector:

normalized_vector = (vector - mean) / standard_deviation

Then, it is multiplied by a learned scale γ and added to a learned bias β:

result = γ ⊙ normalized_vector + β

Why is it necessary? As activations pass through multiple layers, their distribution shifts and expands unpredictably. In some layers, activations become large, while in others, they become small. This instability makes training difficult.

LayerNorm stabilizes training by re-adjusting the mean to 0 and the variance to 1 for each layer's activations. γ and β provide the flexibility to rescale and shift if needed.

Biological Analogy – Intracellular pH and Ion Homeostasis. Cells maintain their cytosolic pH around 7.2, sodium concentration at 12mM, and potassium concentration at 140mM, even when exposed to external stimuli. Without this homeostasis, each signaling pathway would operate in unexpected states, disrupting cellular function. The function of LayerNorm in a neural network layer conceptually aligns with this. It ensures that each layer operates within an expected activation range, providing a constant condition.

Pre-LN vs. Post-LN.

• Post-LN (original paper): Output = LayerNorm(F(x) + x). Normalization after the residual connection.
• Pre-LN (modern standard): Output = F(LayerNorm(x)) + x. Normalization before the layer.

It has been found that Pre-LN leads to much more stable training, and most modern LLMs use Pre-LN. This includes GPT-2, GPT-3, LLaMA, Qwen, and Claude.

RMSNorm. A variation of LayerNorm. It omits the mean subtraction and normalizes only using the RMS (root mean square). It is slightly faster to compute and has similar performance, making it the de facto standard after LLaMA.

RMSNorm(x) = γ ⊙ x / sqrt(mean(x^2) + ε)

The Big Picture – Transformer Architecture

Putting the components we've discussed so far together gives us the overall picture of a Transformer.

Sequence of token IDs [8422, 15903, 421, 892]
   ↓
Token Embeddings (lookup in a V × d matrix)
   ↓
+ Positional Encoding (or RoPE rotation)
   ↓
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
Transformer Block 1
  LayerNorm → Multi-head Attention → + Residual
  LayerNorm → FFN → + Residual
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
Transformer Block 2
  ... (same structure)
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
...
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
Transformer Block 96 (in GPT-3)
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
   ↓
Final LayerNorm
   ↓
Language Modeling Head (d × V matrix)
   ↓
Probability distribution over vocabulary size V (predicts the next token)

Each Transformer block is a combination of the components discussed in this section – normalization, attention, FFN, and residual connections. The nature of attention will be the topic of Part 6. For now, let's understand it as "a component that allows each token to absorb information from other tokens in a context-aware manner."

Scale. In today's large LLMs, most of the parameters in each block are actually in the FFN. The attention parameters are relatively small. For example, in each block of GPT-3, the attention parameters are ~50M, and the FFN parameters are ~200M. In other words, attention is a component for information transfer, and the FFN is a component for information storage and processing. Recent research suggests that a significant portion of the LLM's "knowledge" is stored in the FFN parameters.

Biological Application Scenarios

Scenario 1 – ESM Protein Language Model

ESM (Evolutionary Scale Modeling), mentioned in Part 4, is exactly the Transformer architecture discussed in this section. The differences are:

• Tokens: Instead of words, it uses 20 amino acids.
• Embeddings: Each amino acid has a learned vector. If training is successful, this vector will reproduce the physicochemical properties of the amino acid (hydrophobicity, charge, size).
• Positional Encoding: Represents the position in the sequence. Crucial for predicting secondary and tertiary structures.
• Attention: Learns interactions between amino acids that are far apart in the sequence (e.g., disulfide bonds, hydrophobic core, active site).

In the trained ESM embedding space, proteins with similar functions cluster together, and amino acid pairs that are in contact in the sequence are naturally emphasized in the attention. This attention map is used in the training of AlphaFold.

Scenario 2 – Enformer Genome Sequence Model

DeepMind's Enformer is a Transformer that takes a genome sequence (A/T/G/C) as input and predicts cell-type-specific gene expression.

• Tokens: DNA bases (or k-mers).
• Positional Encoding: Represents the position in the genome. Necessary for learning local features such as promoters, enhancers, and CTCF binding sites.
• Long Context: Enhancers can be hundreds of kb away from genes. Enformer's context window is 100kb+.
• Residual Flow: Essential for maintaining gradient flow in very deep networks (tens of blocks).

Scenario 3 – ChemBERTa Chemical Structure Model

A model that represents molecular structures as SMILES strings and feeds them into a Transformer. After training, molecules with similar chemical properties cluster together in the embedding space. Used in drug discovery for searching candidate molecules and generating similar molecules.

Key Takeaways

• The three challenges of language (discreteness, order, and context) are addressed by tokenization + embedding, positional encoding, and attention, respectively.
• Embeddings map discrete tokens to high-dimensional real-valued spaces. If training is successful, similar tokens are placed close together, and semantic relationships are represented by vector arithmetic. This is similar to the UMAP space of scRNA-seq data.
• Positional encoding informs the neural network of the order of words. Examples include sinusoidal, learned, RoPE, and ALiBi. RoPE is the modern standard.
• Transformer blocks are combinations of attention, FFN, residual connections, and LayerNorm.
• Residual connections solve the vanishing gradient problem, enabling the training of 100+ layer networks. This is conceptually similar to cellular homeostasis.
• Layer normalization stabilizes activation distributions and helps with training convergence. Pre-LN and RMSNorm are the modern standards.
• Most of the parameters in large LLMs are in the FFN, and there is evidence that a significant portion of the knowledge is stored in these parameters.

📐 Appendix — Mathematical Formulas for Experts

Difficulty: Very Hard Target Audience: Readers with a university-level understanding of linear algebra, probability, and optimization theories.

A.1 Mathematics of the Embedding Layer

Vocabulary size V, embedding dimension d.

Embedding Matrix: E ∈ ℝ^{V × d}

Token ID i → Embedding Vector: e_i = E[i, :] (i-th row)

Training: E is the trainable parameter. Each row is updated individually via backpropagation.

Shared Embedding (Weight Tying): In language modeling heads (the last layer d → V), the weights are often shared with E^T. This saves parameters and improves training stability.

A.2 Sinusoidal Positional Encoding

Original Transformer paper (Vaswani et al., 2017):

PE(p, 2i)   = sin(p / 10000^(2i/d))
PE(p, 2i+1) = cos(p / 10000^(2i/d))

Key Property: The PE at position p+k can be expressed as a linear function of the PE at position p:

PE(p+k) = M_k · PE(p)

M_k is a rotation matrix that depends only on k. This property makes it easy for the model to learn relative positions.

Frequency Spectrum: As the dimension i increases, the period of the waveform increases (10000^(2i/d)). Lower dimensions represent short distances, and higher dimensions represent long distances.

A.3 RoPE (Rotary Position Embedding)

At each position p, the embedding vector is divided into pairs of 2D vectors and rotated by a certain angle.

The embedding vector x ∈ ℝ^d is divided into d/2 pairs of 2D vectors (x_{2i}, x_{2i+1}).

Each pair is rotated by a position-dependent angle θ_{p,i} = p · 10000^{-2i/d}:

[x'_{2i}  ]   [cos(θ_{p,i})  -sin(θ_{p,i})]   [x_{2i}  ]
[x'_{2i+1}] = [sin(θ_{p,i})   cos(θ_{p,i})] · [x_{2i+1}]

Key Property: The inner product of two vectors depends only on the position difference p - q:

⟨RoPE(x, p), RoPE(y, q)⟩ = f(x, y, p-q)

This property naturally represents relative positional information in the attention score. Continues from Attention #6.

A.4 ALiBi (Attention with Linear Biases)

Instead of adding positional encoding to the embeddings, a positional distance penalty is added to the attention scores.

After calculating the Q, K attention scores:

score(i, j) = ⟨q_i, k_j⟩ / sqrt(d) - m · |i - j|

m is a slope that is different for each attention head (e.g., 2^{-8/h}, h = number of heads).

Advantages: During training, it can generalize to contexts longer than the original context. No computational cost.

A.5 LayerNorm Formula

For a vector x ∈ ℝ^d:

μ = (1/d) Σ_i x_i                    (Mean)
σ^2 = (1/d) Σ_i (x_i - μ)^2          (Variance)
x̂_i = (x_i - μ) / sqrt(σ^2 + ε)      (Normalization)
y_i = γ_i · x̂_i + β_i                 (Scale and Shift)

γ, β ∈ ℝ^d are the trainable parameters. ε ≈ 1e-5 is for numerical stability.

Important: Normalization is done for each vector, not for the batch. This is the difference from BatchNorm. BatchNorm has batch size dependency and is unsuitable for sequence models.

A.6 RMSNorm Formula

LayerNorm with the mean subtraction part omitted:

y_i = γ_i · x_i / sqrt((1/d) Σ_j x_j^2 + ε)

Half the number of parameters (no β), slightly faster computation. Adopted by LLaMA, Mistral, and Qwen.

A.7 Partial Derivative of Residual Connection

Block y = F(x) + x.

∂y/∂x = ∂F/∂x + I

I is the identity matrix. The fact that 1 always remains is the key to solving the vanishing gradient problem.

Passing through multiple blocks:

y = F_L(F_{L-1}(...F_1(x)...)) + skip path

During backpropagation, the skip path is added, not multiplied, so the gradient does not decay exponentially.

A.8 Pre-LN vs Post-LN

Post-LN (original paper):

y = LayerNorm(F(x) + x)

Pre-LN (modern standard):

y = F(LayerNorm(x)) + x

Advantages of Pre-LN:

• The residual flow does not pass through the normalization, so it always maintains the original scale.
• Training is stable without warm-up.
• Training of very deep networks (100+ layers) is possible.

Disadvantages of Pre-LN:

• There are reports that the performance may be slightly worse. Recently, hybrid approaches such as DeepNet have also emerged.

A.9 FFN (Feed-Forward Network) Formula

FFN inside the Transformer block. Typically a 2-layer MLP:

FFN(x) = W_2 · σ(W_1 · x + b_1) + b_2

W_1 ∈ ℝ^{4d × d}, W_2 ∈ ℝ^{d × 4d} (intermediate dimension is 4 times larger). Convention of the GPT family.

σ is the activation function. Up to GPT-2, GELU was used, and recent large models use SwiGLU:

SwiGLU(x) = (W_2 · x) ⊙ Sigmoid(W_1 · x) · V · x

A structure where two gates are multiplied. Adopted by LLaMA, PaLM, and Qwen.

A.10 Approximate Calculation of the Number of Parameters

Embedding dimension d, number of layers L, vocabulary size V:

• Embedding: V · d (shared if weight tying is used)
• Attention in each block: 4 · d^2 (Q, K, V, Out each d × d)
• FFN in each block: 8 · d^2 (d → 4d → d)
• LayerNorm in each block: 4d (almost negligible)

Total per block: ~12 · d^2

Total parameters (excluding embeddings): ~12 · L · d^2

GPT-3 175B: L = 96, d = 12288. 12 · 96 · 12288^2 ≈ 174B. As the name suggests.

In training, consider the optimizer state (Adam requires 2 times the memory per parameter) and the storage of activation values, and the actual GPU memory is about 8 to 16 times the number of parameters.

References

All the content, scenarios, analogies, and figures in this part are developed in-house by BioPlayground, and the following are external references that will help you learn the concepts.

• Original Transformer paper: Vaswani et al., "Attention Is All You Need" (NeurIPS 2017)
• RoPE original paper: Su et al., "RoFormer: Enhanced Transformer with Rotary Position Embedding" (2021)
• ALiBi original paper: Press et al., "Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation" (ICLR 2022)
• Pre-LN analysis: Xiong et al., "On Layer Normalization in the Transformer Architecture" (ICML 2020)
• RMSNorm original paper: Zhang & Sennrich, "Root Mean Square Layer Normalization" (NeurIPS 2019)
• Residual connection original paper: He et al., "Deep Residual Learning for Image Recognition" (CVPR 2016)
• Word embedding original: Mikolov et al., "Efficient Estimation of Word Representations in Vector Space" (word2vec, 2013)
• Residual flow interpretation: Elhage et al., "A Mathematical Framework for Transformer Circuits" (Anthropic 2021)
• Deep learning visualization education: 3Blue1Brown "Deep Learning" Ch 5·6 (YouTube) — for pedagogical reference
• BPE original paper: Sennrich et al., "Neural Machine Translation of Rare Words with Subword Units" (ACL 2016)
• Enformer paper: Avsec et al., "Effective gene expression prediction from sequence by integrating long-range interactions" (Nature Methods 2021)

This part assembled the skeleton of the Transformer. In Part #6, we will explore how its heart, attention, works.

Next Concepts

• Ep. #6 attention-mechanism — The heart of the Transformer. The principle of how words exchange information with each other.
• Ep. #7 prompt-engineering — How to effectively utilize a trained Transformer.
• Ep. #12 pytorch-basics — Implement the components from this episode in code.