Architecture — decoder-only Transformer

The goal, stated concretely

A language model predicts the next token given the previous ones. More precisely: for every position t in a sequence, it produces a probability distribution over the vocabulary:

$$ P(x_t \mid x_1, x_2, \dots, x_{t-1}) $$

Training minimizes the cross-entropy between that predicted distribution and the true next token.

A decoder-only Transformer differs from the original "encoder + decoder" design in Attention Is All You Need in three ways:

• There is one stack of Transformer blocks, not two.
• Attention is causal — each position only attends to earlier positions, enforced with a lower-triangular mask.
• Training is straightforward: take a long sequence, shift it, and make the model predict each next token at every position in parallel (one forward pass learns T predictions).

This is the architecture behind GPT-2, Llama, and every modern chat LM. Lavender-2 follows the Llama recipe closely.

One Transformer block, end to end

For a single block with hidden size d_model:

x_in  (B, T, d_model)
  │
  ├──── RMSNorm ──── Attention (GQA + RoPE) ──── residual add
  │                                                │
  │                                                ▼
  │                                              x_mid
  │                                                │
  ├──── RMSNorm ──── FFN (SwiGLU MoE) ─────────── residual add
  │                                                │
  │                                                ▼
  │                                              x_out
  └────────────────────────────────────────────────┘

B = batch size, T = sequence length. Pre-norm placement (normalization before each sub-layer, not after) is from Xiong et al. 2020 and is the standard choice today — it trains more stably than the original post-norm placement.

The model stacks n_layers of these blocks, then a final RMSNorm and a linear output head.

Layer-by-layer

Token embedding

nn.Embedding(vocab_size, d_model)

Each token ID indexes into a vocab_size × d_model lookup table to get a learned d_model-dimensional vector. Position is not injected at this step — that is RoPE's job inside attention.

RMSNorm

Standard LayerNorm centers and scales:

$$ \text{LayerNorm}(x) = \gamma \cdot \frac{x - \mu}{\sigma} + \beta $$

RMSNorm drops the mean centering and the learnable bias:

$$ \text{RMSNorm}(x) = \gamma \cdot \frac{x}{\sqrt{\text{mean}(x^2) + \varepsilon}} $$

where γ is a learned scale vector. Empirically this gives similar quality to LayerNorm with slightly less compute and one fewer parameter per dimension.

Paper: Zhang & Sennrich, Root Mean Square Layer Normalization, NeurIPS 2019. https://arxiv.org/abs/1910.07467.

Rotary Position Embedding (RoPE)

The model needs to know which position each token occupies, but decoder-only models generally avoid adding a positional vector to the embedding (that tends to hurt long-context generalization). RoPE instead rotates the query and key vectors inside each attention head by an angle that depends on the position:

q_m = R(m · θ) · q
k_n = R(n · θ) · k

where R(φ) is a 2D rotation matrix applied to pairs of dimensions. The attention dot product then becomes:

q_m · k_n  =  q · R((n − m) · θ) · k

which depends only on the relative position n − m. That is exactly what attention over a sequence should care about, and no explicit position embedding ever has to be stored in the checkpoint.

v0 precomputes the cos/sin lookup tables once at model init and applies them per layer. Paper: Su et al., RoFormer, 2021. https://arxiv.org/abs/2104.09864.

Attention with Grouped-Query Attention (GQA)

Standard multi-head attention splits the representation into n_heads heads, each with its own Q, K, and V projections. Grouped-Query Attention keeps n_heads query projections but reduces the number of key/value projections: n_kv_heads < n_heads, and query heads are split into groups that share a single K/V pair.

Config	n_heads	n_kv_heads
Dev	8	4
70B	40	8

Why bother? The KV cache (see inference) dominates memory at generation time, and its footprint scales with n_kv_heads × d_head × T. Cutting n_kv_heads by 5× in the 70B config cuts per-token cache memory by 5× with almost no quality loss. Modern LLMs essentially all use GQA at scale.

Paper: Ainslie et al., GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints, EMNLP 2023. https://arxiv.org/abs/2305.13245.

Feed-forward sub-layer

The second half of each block is a feed-forward network. In v0 this is a SwiGLU expert wrapped by a Mixture-of-Experts router. That has enough depth on its own to deserve its own page:

→ SwiGLU & Mixture-of-Experts.

Final output head

After the last Transformer block:

nn.Linear(d_model, vocab_size, bias=False)

This produces logits of shape (B, T, vocab_size). Training applies cross-entropy against the shifted input token IDs; inference applies softmax + sampling at the last position.

Two shipped configs, side by side

Field	Dev	70B reference
d_model	small	5120
n_heads / n_kv_heads	8 / 4	40 / 8
n_layers	4	40
d_ff (per expert)	small	13824
num_experts / top_k	4 / 2	8 / 2
vocab_size	32,000	128,256

The dev config runs end-to-end in v0.1.0. The 70B column is a design target, not a measurement.

Reading the code

All of the above lives in the model cell of full_llm_pipeline.ipynb. The file is a single notebook on purpose — every layer is readable in one scroll without jumping between files.

Next stage: SwiGLU & MoE — the feed-forward sub-layer.