See also previous posts on transformers:

• Transformers from first principles in Julia.
• Generative transformer from first principles in Julia.

All code available at github.com/LiorSinai/TransformersLite.jl/tree/feature/mla.

Table of Contents

1 Introduction

In January 2025, DeepSeek unveiled their new DeepSeek-V3 and DeepSeek R1 models. It took the world by storm. Users were impressed with its abilities on top of their claims that is was up to 50× more efficient to train and run than their competitors.

They also released multiple papers (DeepSeek-V2, DeepSeek-V3, DeepSeek-R1) with an impressive array of new techniques across the whole machine learning pipeline, from high level theory to intricate implementation details. Most of it built on existing ideas in innovative ways. They include:

• Theory
  • Multi-Head Latent Attention (MLA): compress vectors during attention, which reduces the cache size during inference.
  • DeepSeekMoE: segmented and isolated mixture of experts.
  • Multi-token prediction.
  • Reinforcement learning with Group Relative Policy Optimization but without supervised data.
  • Improved chain-of-thought reasoning.
• Implementation
  • DualPipe: accelerate training by overlapping forward and backward computation communication phases.
  • Exponential moving average on the CPU.
  • Mixed precision floating point numbers during training. FP8 and FP32 are used.
  • Low-precision storage and communication.

The aim of this post is to explore the first of these ideas in depth, namely Multi-Head Latent Attention (MLA). It is actually a combination of 3 ideas:

1. Attention and multi-head attention from Attention is all you need (2017).
2. KV Caching.
3. Low-Rank Adaption matrices (LoRA) from LoRA: Low-Rank Adaptation of Large Language Models (2021).

The idea behind MLA is simple: it compresses the input matrix into a single matrix for caching during inference, as opposed to standard KV caching which caches two intermediate matrices. This is a novel innovation by DeepSeek. However, MLA has side effects which DeepSeek barely address. The compression also results in a performance boost but most likely comes with a qualitative penalty. That these are not discussed is a notable omission in their paper.

DeepSeek also adds two unwieldly enhancements to MLA. While they only speak good of these, they complicate the mathematics and require specialised optimised code to see performance gains:

1. Weight absorption.
2. Decoupled rotary position embeddings (RoPE), a modification of RoFormer: Enhanced Transformer with Rotary Position Embedding (2021).

The original source code is written in Python with the PyTorch machine learning framework. However, my favourite language is Julia and continuing with my previous posts on transformers, all the code here is written in Julia using the Flux.jl machine learning framework.

The mathematics therefore follows Julia’s columnar major format and not Python’s row major format. In this format each column represents a single sample. The second dimension is sequence length and all remaining dimensions represent batches.

2 KV Caching

2.1 Theory

There is an inefficiency in using transformers for generation. In classification use cases, the model only needs to calculate attention over the input sentence once (per layer) before making its prediction. But in text generation, it needs to recalculate it over the entire sentence. One analogy I read is that this is like having to reread the entire sentence so far in order to produce the next word and then rereading again to produce the next word, and so on.

In mathematical terms, this is computing the attention between the first token and itself, between the first two tokens, between the first 3 tokens and so on until it’s between all the tokens, and repeating this whole process each time for every new token. This was the case for the generator I made based on Andrej Karpathy’s Zero to Hero course.

It would be better to have a “memory” of what has been generated so far. This is where the KV cache comes in, for Key-Value cache. We can store the previous keys and values and use them to calculate the new attention value. This will give the exact same value as recalculating the entire attention.

Here is a mathematical proof. For a visual interpretation, please see João Lages’ excellent article.

The attention equation is

where

where $X \in \mathbb{R}^{d\times n \times B}$ and $W^K, W^Q, W^V \in \mathbb{R}^{d_hH \times d}$.

Focusing on one head in one batch with dimension $d_h$, the input matrix $X$ is split into the first $n-1$ columns and the $n$th column. The first multiplication can then be written as (dimensional analysis on right):¹

Looking at the final line, the first $(n-1)$ columns of the query can be safely dropped without affecting the $n$th column. (We would do this anyway in generation.) The $n$th column only depends on $K_{1:n-1}$, $K_n$ and $Q_n$. Of these, $K_{1:n-1}$ will come from the cache and the other two will be calculated from $X_n$.

Similarly for the next multiplication we have ($Z=\text{softmax}(S)$):

However, we said we are dropping the first $(n-1)$ columns. Without them only the $n$th column is calculated. Hence we have:

which depends on $V_{1:n-1}$, which will come from the cache, and $V_n$, which will be calculated from $X_n$.

There will be two caches each with size $d_h H \times N \times B$ for $H$ heads, a maximum sequence length of $N$ and a maximum batch size of $B$. The total cache size can grow very large for large transformers with many multi-head attention layers. The primary aim of MLA is to reduce the size of this cache. That will be covered in the next section.

2.2 Code

Building on my code in TransformersLite.jl, it is straightforward to create a new MultiHeadAttentionKVCache layer with two caches:

struct MultiHeadAttentionKVCache{
    Q<:Dense, K<:Dense, V<:Dense, O<:Dense, C<:Array{T, 3}  where T
    }
    nhead::Int
    denseQ::Q
    denseK::K
    denseV::V
    denseO::O
    cache_k::C
    cache_v::C
end

Flux.@layer trainable=(denseQ, denseK, denseV, denseO)

The forward pass calculates the current q, k and v values from the input and gets the rest from the cache. It then continues without any additional modifications from the original code:

function (mha::MultiHeadAttentionKVCache)(
    query::A3, key::A3, value::A3
    ; start_pos::Int=1, use_cache::Bool=true, kwargs...
    ) where {T, A3 <: AbstractArray{T, 3}}
    q = mha.denseQ(query) # size(q) == (dh, 1, B)
    k = mha.denseK(key)
    v = mha.denseV(value) # size(k) == size(v) == (dh, 1, B)
    if use_cache
        dim, seq_length, batch_dim = size(query)
        end_pos = start_pos + seq_length - 1
        mha.cache_k[:, start_pos:end_pos, 1:batch_dim] = k
        mha.cache_v[:, start_pos:end_pos, 1:batch_dim] = v
        K = mha.cache_k[:, 1:end_pos, 1:batch_dim]
        V = mha.cache_v[:, 1:end_pos, 1:batch_dim]
    else
        K = k
        V = v
    end
    A, scores = multi_head_scaled_dot_attention(mha.nhead, q, K, V; kwargs...)
    mha.denseO(A), scores
end

Here is a small example of it in action. (For the full code, see test/MultiHeadAttention.jl.)

Create the layer and inputs:

using TransformersLite
using TransformersLite: MultiHeadAttention, MultiHeadAttentionKVCache
using TransformersLite: make_causal_mask, clone_add_kv_cache
nhead, dim_model, dim_out = 4, 32, 13
mha = MultiHeadAttention(nhead, dim_model, dim_out) 
mha = clone_add_kv_cache(mha, 64, 8)
X = randn(Float32, 32, 10, 5)

Fill the cache:

mask = make_causal_mask(ones(10, 10))
A, scores = mha(X, X, X; mask=mask, start_pos=1, use_cache=true)
size(A) # (13, 10, 5)
size(scores) # (10, 10, 4, 5)

Use the cache with a new vector:

x = randn(Float32, 32, 1, 5)
mask = repeat([true], inner=(11, 1))
Ax, scoresx = mha(x, x, x; mask=mask, start_pos=11, use_cache=true)
size(Ax) # (13, 1, 5)
size(scoresx) # (11, 1, 4, 5)

Compare without the cache:

Xx = cat(X, x, dims=2)
mask = make_causal_mask(ones(11, 11))
AXx, scoresXx = mha(Xx, Xx, Xx; mask=mask, start_pos=1, use_cache=false)
isapprox(AXx[:, end, :], Ax[:, end, :]) # true

3 Multi-Head Latent Attention

3.1 C cache

We’ve seen that the KV cache has size $2d_h H \times N \times B$ elements per multi-head attention layer. (Each element is 1-4 bytes depending if FP8, FP16 or FP32 is used.) The aim of MLA is to reduce this, specifically to $d_c \times N \times B$ elements per multi-head attention layer. Therefore we will choose $d_c < 2d_h H$.

DeepSeek’s innovation is to introduce a weight matrix $W^{DKV} \in \mathbb{R}^{d_c\times d}$ to compress the input $X \in \mathbb{R}^{d\times n}$ to a lower rank matrix $C^{KV} \in \mathbb{R}^{d_c \times n}$. This $C^{KV}$ matrix is then stored in the cache. Then two other weight matrices $W^{UK}$ and $W^{UV} \in \mathbb{R}^{d_h H\times d_c}$ uncompress the same $C^{KV}$ matrix to the key $K$ and value $V$ respectively. The above figure shows this visually.

The KV cache is now replaced with a $C^{KV}$ cache of size $d_c \times N \times B$. DeepSeek theorises that this compression also results in a regularization effect that improves performance. This is supported by other LoRA research. However, the lossy compression might instead adversely affect quality. DeepSeek provides no evidence towards either claim.

The compression also results in a significant performance boost which DeepSeek strangely does not mention in their paper. Note that in MLA there are three matrix multiplications to perform to create $K$ and $V$ instead of two matrix multiplications in MHA. However the three multiplications comprise of less scalar operations:² \(\begin{align} \frac{\text{# MLA ops}}{\text{# MHA ops}} &= \frac{2(2d_h H + d)d_c nB}{2(2d_h H)d n B} \\ &= \frac{2\frac{d_h H}{d} + 1}{2 \tfrac{d_h H}{d_c}} \\ &= \frac{3}{2r} \end{align} \tag{3.1.2} \label{eq:mla_ops}\) with the standard $d = d_h H$ and a compression ratio $r=\tfrac{d_h H}{d_c}$. This requires $r > 1.5$ for performance gains. The only performance penalty is the memory required for the $n$th $c^{KV}_n$ vector before it is transferred to the cache, which is $d_c \times 1 \times B$.

DeepSeek-V3 uses $d_h = 128$, $H=128$ and $d_c=4 d_h = 512$ which means it has a compression ratio of $32$ and a 20× speed up!

To reduce the activation memory during training, DeepSeek also applies the same strategy to the query:

In total, five matrix multiplications are needed to create $Q$, $K$ and $V$ instead of three. The overall ratio of scalar operations is $\tfrac{5}{3r}$, which requires $r>1.67$ for performance gains.

DeepSeek give further enhancements to this which will be described shortly. They also apply layer normalisation to $C^Q$ and $C^{KV}$ which I will ignore in this article. For now, lets see this basic version of MLA in action.

3.2 Code

First create a struct similar to the MultiHeadAttentionKVCache layer.

struct MultiHeadLatentAttention{D1<:Dense, D2<:Dense, A<:AbstractArray{T, 3} where T} 
    nhead::Int
    denseDQ::D1
    denseUQ::D1
    denseDKV::D1
    denseUK::D1
    denseUV::D1
    denseO::D2
    cache_ckv::A
end

Flux.@layer MultiHeadLatentAttention trainable=(denseDQ, denseUQ, denseDKV, denseUK, denseUV, denseO)

Here is a convenience constructor to construct it from the various input dimensions:

function MultiHeadLatentAttention(;
    nhead::Int, dim_in::Int, dim_head::Int, dim_lora, dim_out::Int,
    max_seq_length::Int, max_batch_size::Int
    )
    denseDQ = Dense(dim_in => dim_lora; bias=false)
    denseUQ = Dense(dim_lora => dim_head * nhead; bias=false)
    denseDKV = Dense(dim_in => dim_lora; bias=false)
    denseUK = Dense(dim_lora => dim_head*nhead; bias=false)
    denseUV = Dense(dim_lora => dim_head*nhead; bias=false)
    denseO = Dense(dim_head*nhead => dim_out; bias=false)
    cache_ckv = Array{Float32, 3}(undef, dim_lora, max_seq_length, max_batch_size)
    MultiHeadLatentAttention(
        nhead,
        denseDQ, denseUQ,
        denseDKV, denseUK, denseUV,
        denseO,
        cache_ckv
    )
end

The forward pass is:

function (mla::MultiHeadLatentAttention)(query::A3, key::A3
    ; start_pos::Int=1, use_cache::Bool=true, mask::Union{Nothing, M}=nothing
    ) where {T, A3 <: AbstractArray{T, 3}, M <: AbstractArray{Bool}}
    dm, seq_length, batch_dim = size(key)
    cq = mla.denseDQ(query) # size(cq) == (dc, dq, B)
    ckv = mla.denseDKV(key) # size(ckv) == (dc, dkv, B)
    if use_cache
        end_pos = start_pos + seq_length - 1
        mla.cache_ckv[:, start_pos:end_pos, 1:batch_dim] = ckv
        ckv = mla.cache_ckv[:, 1:end_pos, 1:batch_dim]
    end
    K = mla.denseUK(ckv) # size(k) == (dh*nhead, dkv, B)
    V = mla.denseUV(ckv) # size(v) == (dh*nhead, dkv, B)
    Q = mla.denseUQ(cq)  # size(q) == (dh*nhead, dq, B)
    A, scores = multi_head_scaled_dot_attention(mla.nhead, Q, K, V; mask=mask)
    A = mla.denseO(A)
    A, scores
end

Create a test layer with compression ratio $r=\tfrac{d_hH}{d_c}=4$:

nhead, dim_head, dim_lora, dim_out = 8, 64, 128, 8*64
dim_model = nhead * dim_head
N, max_seq_length, batch_dim = 20, 32, 8
mla = MultiHeadLatentAttention(
    nhead=nhead, dim_in=dim_model, dim_head=div(dim_model, nhead),
    dim_lora=dim_lora, dim_out=dim_out,
    max_seq_length=max_seq_length, max_batch_size=batch_dim
    )
X0 = randn(Float32, dim_model, N, batch_dim)

Fill the cache:

mask = make_causal_mask(ones(N, N));
A, scores = mla(X0, X0; mask=mask, use_cache=true); 
size(A) # (512, 20, 8)
size(scores) # (20, 20, 8, 8)

Use the cache with a new vector:

x = randn(Float32, dim_model, 1, batch_dim)
mask = repeat([true], inner=(N + 1, 1))
Ax, scoresx = mla(x, x; mask=mask, start_pos=N+1, use_cache=true)
size(Ax) # (512, 1, 8)
size(scoresx) # (21, 1, 8, 8)

3.3 Absorption

DeepSeek suggests a way to further decrease the computational cost by absorbing weight matrices into each other. To quote DeepSeek directly:

In addition, during inference, since $W^{UK}$ can be absorbed into $W^{Q}$ , and $W^{UV}$ can be absorbed into ${W^O}$, we even do not need to compute keys and values out for attention.

What they mean is that the weight matrices can be multiplied to produce a single weight matrix. This can only be done during inference because during training they need to be kept separate so that the gradients flow properly backwards through each matrix.

This technique can be used independently of MLA.

To show why this works, rewrite the attention equation $\ref{eq:attention}$ as follows:

This looks straightforward but there are further complications with the dimensions. Here is a dimensional analysis of the above equation following the two rules of batch matrix multiplication:

1. The inner matrix dimensions must match. That is, the second dimension of the first matrix must match the first dimension of the second.
2. All the batch dimensions (dimensions 3 and greater) must be equal.

where

• Line 2 reshapes the weight matrices from $d_h H \times d_c$ to $d_h \times d_c \times H$. This is necessary because the non-linear softmax function must be applied independently over each head dimension.
• Line 3 shows that $W^{KQ} \in \mathbb{R}^{d_c \times d_c \times H}$.
• Line 4 adds extra broadcast dimensions to make the batch dimensions match.

Broadcasting is a technique where the smaller array is replicated along all dimensions of size 1 to match the size of the larger array. For broadcasted batched multiplication this only needs to be done for the 3rd and higher dimensions. Pseudo-code for this is:

This same absorption technique can be applied to the value and output matrices: \(\begin{align} Y &= W^O V Z \\ &= W^O (W^{UV} C^{KV} Z) \\ &= W^{OV} C^{KV} Z \end{align} \label{eq:absorbOV} \tag{3.3.3}\)

The dimensional analysis here is similar:

This shows that the $C^{KV} Z$ multiplication is a broadcasted batched multiplication. However where $W^{KQ} \in \mathbb{R}^{d_c \times d_c \times H}$, $W^{OV} \in \mathbb{R}^{d_o \times d_c H}$ is a typical 2D matrix. Therefore the usual matrix multiplication can be applied by reshaping the $C^{KV} Z$ result from a 3D $d_c H \times n \times B$ array to a $d_c H \times nB$ matrix.

3.4 Broadcasted batched multiplication

I have written an implementation in Julia directly based on the pseudo code. It can be seen here: broadcasted_batched_mul.jl. However, it is inefficient and uses scalar indexing which is extremely slow on a GPU.

My solution instead is to physically replicate the broadcasted dimensions. This is of course inefficient compared to virtual replication but it makes the function viable on a GPU. The downside is it can be up to 4× slower than the naive code on a CPU.

using Flux: batched_mul
function broadcasted_batched_mul(x::AbstractArray{T, N}, y::AbstractArray{T, N}) where {T, N}
    batch_dims_x = Tuple(size(x, idx) == 1 ? size(y, idx) : 1 for idx in 3:N)
    dims_x = (1, 1, batch_dims_x...)
    batch_dims_y = Tuple(size(y, idx) == 1 ? size(x, idx) : 1 for idx in 3:N)
    dims_y = (1, 1, batch_dims_y...)
    xb = repeat(x; outer=dims_x)
    yb = repeat(y; outer=dims_y)
    batched_mul(xb, yb)
end

The DeepSeek source code meanwhile uses torch.einsum and carries out the multiplications right to left instead of creating a new matrix.³ Here is the relevant code. As far as I know, this has the same drawbacks with scalar indexing with none of the advantages of absorption as described in their paper.

q = self.wq_b(self.q_norm(self.wq_a(x)))
q = q.view(bsz, seqlen, self.n_local_heads, self.qk_head_dim)
kv = self.wkv_a(x)
wkv_b = self.wkv_b.weight 
wkv_b = wkv_b.view(self.n_local_heads, -1, self.kv_lora_rank)
q = torch.einsum("bshd,hdc->bshc", q, wkv_b)
self.kv_cache[:bsz, start_pos:end_pos] = self.kv_norm(kv)
scores = torch.einsum("bshc,btc->bsht", q, self.kv_cache[:bsz, :end_pos])

Presumably fast MLA implementions for GPU kernels would make use of virtual replication. (It would great if someone can clarify what FlashMLA does.)

3.5 Code

I will detail some of the code here. For the full code, see MultiHeadLatentAttentionV2.jl.

The first step is create the $W^{KQ}$ and $W^{OV}$ matrices.

First, $W^{KR}=(W^{UK})^T W^{UQ}$ while reshaping from $d_h H \times d_c$ to $d_h \times d_c \times H$.

function _absorb_WUK_WUQ(nhead::Int, W_UK::AbstractMatrix, W_UQ::AbstractMatrix)
    dh = div(size(W_UK, 1), nhead)
    dim_lora = size(W_UK, 2)
    W_UQ = permutedims(reshape(W_UQ, dh, nhead, dim_lora), (1, 3, 2)) 
    W_UK = permutedims(reshape(W_UK, dh, nhead, dim_lora), (1, 3, 2))
    W_UKT = permutedims(W_UK, (2, 1, 3)) # (dh, dc, nhead)^T => (dc, dh, nhead)
    batched_mul(W_UKT, W_UQ) 
end
W_KQ = _absorb_WUK_WUQ(nhead, denseUK.weight, denseUQ.weight)

Then $W^{OV}=W^O W^{UV}$ while preserving the head dimension through reshaping:

function _absorb_WO_WUV(nhead::Int, W_O::AbstractMatrix, W_UV::AbstractMatrix)
    dh = div(size(W_UV, 1), nhead)
    dim_lora = size(W_UV, 2)
    dout = size(W_O, 1)
    W_UVh = permutedims(reshape(W_UV, dh, nhead, dim_lora), (1, 3, 2)) # (dh*nhead, dc) => (dh, dc, nhead)
    W_Oh = reshape(W_O, dout, dh, nhead) # (dout, dh*nhead) => (dout, dh, nhead)
    W_OVh = batched_mul(W_Oh, W_UVh) # (dout, dh, nhead) * (dh, dc, nhead)
    reshape(W_OVh, dout, dim_lora*nhead) # (dout, dc, nhead) => (dout, dc*nhead)
end
W_OV = _absorb_WO_WUV(nhead, denseO.weight, denseUV.weight)

The forward pass is the same as in the original code until the end of caching:

function mla_absorb(
    mla::MultiHeadLatentAttention, query::A3, key::A3
    ; start_pos::Int=1, use_cache::Bool=true, mask::Union{Nothing, M}=nothing
    ) where {T, A3 <: AbstractArray{T, 3}, M <: AbstractArray{Bool}}
    dm, seq_length, batch_dim = size(key)
    dh = div(dm, mla.nhead)
    cq = mla.norm_cq(mla.denseDQ(query))  # size(cq) == (dc, dq, B)
    ckv = mla.norm_ckv(mla.denseDKV(key)) # size(ckv) == (dc, dkv, B)
    if use_cache
        end_pos = start_pos + seq_length - 1
        mla.cache_ckv[:, start_pos:end_pos, 1:batch_dim] = ckv
        ckv = mla.cache_ckv[:, 1:end_pos, 1:batch_dim]
    end

Then add the broadcast dimensions:

    ckv_ = Flux.unsqueeze(ckv, dims=3)
    keyT = permutedims(ckv_, (2, 1, 3, 4)) # (dkv, dc, B) => (dkv, dc, 1, B)
    cq_ = Flux.unsqueeze(cq, dims=3) # (dkv, dc, B) => (dkv, dc, 1, B)
    W_KQ = Flux.unsqueeze(mla.W_KQ, dims=4); # (dc, dc, nhead) => (dc, dc, nhead, 1)

Then apply the equations as before except using broadcasted_batched_mul instead of batched_mul:

    atten_base = broadcasted_batched_mul(keyT, broadcasted_batched_mul(W_KQ, cq_))
    atten = one(T)/convert(T, sqrt(dh)) .* (atten_base)
    atten = apply_mask(atten, mask)
    scores = softmax(atten; dims=1)
    A = broadcasted_batched_mul(ckv_, scores) # (dc, dq, nhead, B)
    # (dc, dq, nhead, B) => (dc*nhead, dq, B)
    A = permutedims(A, [1, 3, 2, 4])
    A = reshape(A, :, size(A, 3), size(A, 4))
    mla.denseOV(A), scores 
end

Test that these give the same result:

X0 = randn(Float32, dim_model, N, batch_dim)
mask = make_causal_mask(ones(N, N));
A_naive, scores_naive = mla_naive(mla, X0, X0; mask=mask, use_cache=true);
A_absorb, scores_absorb = mla_absorb(mla, X0, X0; mask=mask, use_cache=true);
isapprox(A_absorb, A_naive) # true

3.6 Decoupled RoPE

The last enhancement DeepSeek adds is RoPE. One issue with RoPE however is that it breaks the absorption property described above. To prove this, RoPE can be represented as a series of matrix multiplications on each column in an input matrix $X$:

Applying this to the scores equation: \(\begin{align} S &= \text{RoPE}(K)^T\text{RoPE}(Q) \\ &= \text{RoPE}(W^{UK}C^{KV})^T \text{RoPE}(W^{UQ}C^Q) \\ &= \begin{bmatrix} R_1 W^{UK}_1 C^{KV}_1 & ... & R_n W^{UK}_n C^{KV}_n \end{bmatrix}^T \\ &\phantom{=x} \begin{bmatrix} R_1 W^{UQ}_1 C^Q_1 & ... & R_n W^{UQ}_n C^Q_n \end{bmatrix} \\ \implies S_{ij} &= (C^{KV}_{i})^T (W^{UK}_i)^T R_i^T R_j W^{UQ}_j C^{Q}_j \end{align} \label{eq:absorb_RoPE} \tag{3.4.2}\)

which shows that the rotation matrices will appear right in the middle of the product.

DeepSeek’s solution is to concatenate another matrix to the bottom of the key $K$ and query $Q$ respectively and only apply RoPE to these matrices. Furthermore, because the new matrix $K^R$ will also need to be cached, they share it across all heads. To put it another way, $K^R$ will be broadcasted across the head dimension during multiplication. So for each head $h$:

where $W^{KR} \in \mathbb{R}^{d_R \times d}$ and $W^{QR} \in \mathbb{R}^{d_R H \times d_c}$. This means that $K,Q \in \mathbb{R}^{(d_h + d_R) \times n \times H \times B}$. The cache will consist of both $C^{KV}$ and $K^R$ for a total size of $(d_c + d_R) \times N \times B$ elements per layer.

Very conveniently, this results in an addition between the original and embedded scores:

which means that these results can be calculated separately. Note that $S_h$ must now be scaled by $1/\sqrt{d_h + d_R}$ instead of $1/\sqrt{d_h}$.

3.7 Code

I will detail some of the code here. For the full code, see MultiHeadLatentAttentionV2.jl.

Here is a Julia implementation of RoPE:

struct RoPE{T}
    base::Int
    dim::Int
    seq_length::Int
    freqs_complex::Matrix{Complex{T}}
end

RoPE(dim::Int, max_seq_length::Int; base::Int=10_000) = RoPE(Float32, dim, max_seq_length; base=base)

function RoPE(T::DataType, dim::Int, max_seq_length::Int; base::Int=10_000)
    @assert dim % 2 == 0 "Require even dim"
    θ = 1 ./ (base .^ ((0:2:(dim - 2)) / dim))
    angles = θ * transpose(0:(max_seq_length-1))
    freqs = map(x -> reverse(sincos(x)), angles)
    freqs_complex = map(cs -> Complex(cs...), freqs)
    RoPE{T}(base, dim, max_seq_length, freqs_complex)
end

The forward pass can be calculated with matrices, but the RoPE authors gave a more efficient implementation with complex numbers:

(r::RoPE)(x::AbstractArray) = apply_rope(x, r.freqs_complex[:, 1:size(x, 2)])
(r::RoPE)(x::AbstractArray, indices) = apply_rope(x, r.freqs_complex[:, indices])

function apply_rope(x::AbstractArray{T}, freqs_complex::AbstractMatrix{<:Complex{T}}) where T
    x_complex = reinterpret(Complex{T}, x)
    rx_complex = freqs_complex .* x_complex
    T.(reinterpret(T, rx_complex))
end

Then add the embedding, denseQR and denseKR layers to the MultiHeadLatentAttention struct.

The embeddings are applied as follows:

function _apply_embeddings(mla::MultiHeadLatentAttention, key::A3, cq::A3, idx::UnitRange{Int}) where {T, A3 <: AbstractArray{T, 3}}
    dim_lora, dq, batch_dim = size(cq)
    kr = mla.denseKR(key)
    qr = mla.denseQR(cq)
    kr = mla.embedding(kr, idx) # size(kr) == (dr, dkv, B)
    qr = permutedims(reshape(qr, :, mla.nhead, dq, batch_dim), (1, 3, 2, 4)) # (dr*nhead, dq, B) => (dr, dq, nhead, B)
    qr = mla.embedding(qr, idx)
    kr, qr
end
kr, qr = _apply_embeddings(mla, key, cq, start_pos:end_pos)

Note that embedding is done per head, hence the reshaping of qr.

For the naive method, concatenate along the head dimension. This requires reshaping for qr and repeating kr.

    Q, K = _cat_decoupled_embedding(mla.nhead, Q, qr, K, kr)

where:

function _cat_decoupled_embedding(
    nhead::Int, Qin::A3, Qr::A4, Kin::A3, kr::A3
    ) where {T, A3 <: AbstractArray{T, 3}, A4 <: AbstractArray{T, 4}}
    dhq, dq, B = size(Qin)
    dhk, dkv, B = size(Kin)
    Q = reshape(
        cat(reshape(Qin, :, nhead, dq, B), permutedims(Qr, (1, 3, 2, 4)), dims=1),
        : , dq, B)
    Kr = repeat(Flux.unsqueeze(kr, dims=2), outer=(1, 2, 1, 1))
    K = reshape(
        cat(reshape(Kin, :, nhead, dkv, B), reshape(Kr, :, nhead, dkv, B), dims=1),
        :, dkv, B)
    Q, K
end

Then continue as before.

With absorption, broadcast batched multiply kr and qr and add to the original attention:

    krT = Flux.unsqueeze(permutedims(kr, (2, 1, 3)), dims=3) # (dr, dkv, B) => (dkv, dr, 1, 1B)
    atten_base = broadcasted_batched_mul(keyT, broadcasted_batched_mul(W_KQ, cq_))
    atten_embed = broadcasted_batched_mul(krT, qr)
    atten = one(T)/convert(T, sqrt(dh + dr)) .* (atten_base + atten_embed)

Conclusion

Overall, I think MLA is a smart and useful idea. However, after having explored it in depth, I am more critical of their enhancements.

The basic premise of Multi-Head Latent Attention is simple. It compresses the input matrix so that a single smaller $C^{KV}$ matrix can be stored instead of the key $K$ and value $V$ matrices. It then uncompresses this matrix into $K$ and $V$ with two additional weight matrices. This also results in a significant performance increase which scales with the compression ratio $\frac{d_h H}{d_c}$ by a factor of $\tfrac{2}{3}$ - so the compression needs to be greater than a modest 1.5 for gains to be realised - and requires no further modifications to existing MHA code. However, it is unclear what the qualitative effects of the compression are, and it is strange that DeepSeek did not discuss the performance benefits.

To this DeepSeek adds weight absorption and decoupled RoPE. We have seen this complicates the mathematics and requires careful dimensional analysis. True performance gains only come with an optimised broadcasted_batched_mul function. Their own open source code does not even have such optimisations. Personally, I see no benefit to this and would recommend the naive method with RoPE applied normally. That is, apply each of the weight matrices to their inputs individually and then apply RoPE to the entire $K$ and $Q$ matrices.

While I am impressed with the ingenuity behind MLA, DeepSeek’s omissions coupled with extra, unwieldly enhancements makes me more skeptical of their methodology. If I examine their other techniques, I will do so with more caution.

Table of Contents

1 Introduction

I recently updated my PolygonAlgorithms.jl package to use the Martinez-Rueda algorithm for boolean operations between polygons. I had originally implemented a version of the Weiler-Atherton algorithm, explained in detail in an earlier blog post. However, that algorithm can only calculate intersections between polygons, whereas Martinez-Rueda simultaneously calculates the intersection as well as unions, differences and XORs between the polygons. See the above example and the table below for a brief comparison between the algorithms.

The Martinez-Rueda algorithm is more versatile because it fundamentally operates at a segment level whereas Weiler-Atherton which operates a point level, and so it has a “bigger picture” view of the polygons. A disadvantage of the Martinez-Rueda algorithm is that it is more sensitive to numerical inaccuracies - for reasons that will be described shortly - such as a line that is almost vertical or tiny regions of intersection. In practice I found their runtimes similar, with Martinez-Rueda running faster in some situations and slower in others. For the spiral-star example, it is about 1.5 slower.

The original paper can be found here, but I followed the guide at https://sean.fun/a/polygon-clipping-pt2/.

The core idea behind the Martinez-Rueda algorithm is to calculate fill annotations for each segment for each polygon: is this segment filled above and below by this polygon, and is it filled above and below by the other polygon? Once these are known, it is easy to select the relevant segments to the given operation, and to link them up again into polygons.

The genius of the Martinez-Rueda algorithm is to extend upon the Bentley-Ottmann algorithm for segment intersections to do this. It does a vertical line sweep from left to right, bottom to top. At any any given moment, we can imagine having a stack of all the lines that intersect the vertical line, ordered from top to bottom. According to the Bentley-Ottmann algorithm, to find intersections through a segment, we only need to check for intersections with the segments immediately above and immediately below it in stack. At the same time, we can propagate the fill annotations from the segment below, or empty space if nothing is below it. Hence, finding the exact segments that are above and below a segment is paramount to this algorithm, and even slight mistakes can cause errors that propagate to other segments.

This is the gist of the algorithm. The practicalities of handling the event queue and many edge cases such as handling coincident lines and tricky annotation situations are described in the article. From now, I will focus only on the is_above? algorithm to determine if a segment is above another segment. I spent a long time debugging the whole Martinez-Martinez algorithm against a variety of test cases, and I always seemed to land back at this is_above? function. Getting this function right solved most of my problems.

2 Is above?

For reference, the function is called statusCompare in the article.

The goal of this algorithm is to sort lines by height. This will then give a sweep status like:

Given line segments $AB$ and $CD$, it is tempting to sort them only by the starting point coordinate:

This will work in most cases. However already in the figure we can see an example where it does not. Line 2’s starting point is below line 3’s, but it makes more sense to consider line 2 as “above” line 3.

A better definition for “above” is needed. Instead, we will consider one segment above another if its starting point is above its projection on the other line:

That is, $y_p \leq y_A$, where $y_p$ is:

However one problem is this equation is not symmetrical. (This is the case in the original statusCompare.) In the figure above, both projections are below the other line. Hence is_above will return false for both segments. Yet one must be above the other. Therefore to maintain symmetry, the function will always only consider the right segment. If it is the segment of interest, we check if its starting point is above its projection on the left line. Otherwise if we are checking if the left segment is above, we check if the right segment’s starting point’s is below its projection on the left segment.

There are two other special cases. The first is if the starting point is colinear or coincident with the other line:

In this case, the endpoint is used instead.

The second and final case is a vertical line:

As implied by equation $\ref{eq:projection}$, if the line is vertical the projection equation is indeterminant. In fact, if the line were slight sloped towards the left or towards the right, the answer would differ. Here instead we will simply compare y-values. That is, fallback to $\ref{eq:greater}$. (The original statusCompare did not account for this case.)

Hence the is_above algorithm is:

Is segment AB above CD?
inputs: $AB$, $CD$
if colinear($A, C, D$)
$\quad$ return point_above_line($B, CD$)
if $x_C < x_A$
$\quad$ return point_above_line($A, CD$)
else
$\quad$ return not point_above_line($C, AB$)

where point_above_line is:

point_above_line
inputs: $P$, $CD$
if $x_C = x_D$
$\quad$ return $y_p \geq \text{min}(y_C, y_D)$
return $ (y_P - y_C)(x_D - x_C) - (y_D - y_C)(x_P - x_C) \geq 0$

This final algorithm is simple, but absolutely crucial for the algorithm.

3 Compare events

For reference, the function is called eventCompare in the article.

It sorts segment events from left to right, bottom to top. There are two events per segment: a start event and an end event. An example ordering is:

The algorithm is:

1. If the points are not the same, the smaller event is to the left, or the lower one if they are on a vertical line.
2. If the other points are also the same, this event is not smaller. (Equal segments.)
3. If the one is a start event and the other an end event, the end event is considered smaller. (Common points.)
4. The smaller event is below the other one according to not is_above, unless the segment of interest is vertical, then the smaller event is “not above” if it is to the right. (Common start/end points.)

For example, in the picture:

• Event 1 is smaller than event 2 by step 4: lower event is to right of a vertical segment.
• Event 2 is smaller than event 5 by step 1: they are on the same segment, but event 2 is defined by the lower start point.
• Event 3 is smaller than event 4 by step 4: common start point but segment 3 is lower than segment 4.
• Event 5 is smaller than event 6 by step 3: same point but event 5 is an end event, while event 6 is a start event.

And so on.

4 Conclusion

This was a short post to address minor issues and some improvements to two parts of the Martinez-Rueda implementation from https://sean.fun/a/polygon-clipping-pt2/. Otherwise that article did a very good job at explaining this algorithm and I highly recommend it.

This is part of a series. The other articles are:

• Part 1: ChainRules.
• Part 2: Automation with expressions.
• Part 3: Automation with IR.
• Part 4: Extensions.

All source code can be found at MicroGrad.jl.

Table of Contents

1 Introduction

The previous four sections have developed a minimal automatic differentiation package. The aim of this part is to demonstrate how it can be used as the backbone for a machine learning framework like Flux.jl. In this post we will create a multi-layer perceptron also known as a fully connected neural network. This is an extremely popular and powerful machine learning model. New code will be needed for the forward pass and for some extra rrules. Otherwise, the rest is handled by code from the previous parts.

2 Moons dataset

The moons dataset is a toy dataset for testing and visualising classification algorithms. While clearly distinct, the curved nature of the two classes requires a non-linear algorithm to discern them. This was the dataset chosen by Karpathy to demonstrate his micrograd package, and so it will be used here too.

This dataset can be reconstructed in Julia as follows, based on the Scikit-Learn function:

using Random

function make_moons(rng::AbstractRNG, n_samples::Int=100; noise::Union{Nothing, AbstractFloat}=nothing)
    n_moons = floor(Int, n_samples / 2)
    t_min = 0.0
    t_max = π
    t_inner = rand(rng, n_moons) * (t_max - t_min) .+ t_min
    t_outer = rand(rng, n_moons) * (t_max - t_min) .+ t_min
    outer_circ_x = cos.(t_outer)
    outer_circ_y = sin.(t_outer)
    inner_circ_x = 1 .- cos.(t_inner)
    inner_circ_y = 1 .- sin.(t_inner) .- 0.5

    data = [outer_circ_x outer_circ_y; inner_circ_x inner_circ_y]
    z = permutedims(data, (2, 1))
    if !isnothing(noise)
        z += noise * randn(size(z))
    end
    z
end

make_moons(n_samples::Int=100; options...) = make_moons(Random.default_rng(), n_samples; options...)

Creating the moons and labels:

n = 100
X = make_moons(2n; noise=0.1) # 2×200 Matrix 
y = vcat(fill(1, n)..., fill(2, n)...) # 200-element Vector{Int64}

3 Layers

3.1 ReLU

The Rectified Linear Unit (ReLU) is a common activation function in machine learning. It is defined as follows:

This can be realised as a broadcast of the max function:

relu(x::AbstractArray) = max.(0, x)

The derivative is:

In code:

function rrule(::typeof(relu), x::AbstractArray)
    relu_back(Δ) = (nothing, ifelse.(x .> 0, Δ, 0))
    relu(x), relu_back
end

3.2 Dense layer

The fully connected layer equation is:

This is the code from Flux.jl to create this fully connected layer (source):

using Random
struct Dense{M<:AbstractMatrix, B<:AbstractMatrix, F}
    weight::M
    bias::B
    activation::F
end

function (a::Dense)(x::AbstractVecOrMat)
    a.activation(a.weight * x .+ a.bias)
end

Dense((in, out)::Pair; activation=relu) = Dense(glorot_uniform(in, out), zeros(out, 1), activation)

function glorot_uniform(rng::AbstractRNG, fan_in::Int, fan_out::Int)
    scale = sqrt(24 / (fan_in + fan_out))  # 0.5 * sqrt(24) = sqrt(1/4 * 24) = sqrt(6)
    (rand(rng, fan_out, fan_in) .- 0.5) .* scale
end

glorot_uniform(fan_in::Int, fan_out::Int) = glorot_uniform(Random.default_rng(), fan_in, fan_out)

Also add a method to paramaters:

parameters(a::Dense) = (;weight=a.weight, bias=a.bias)

Create and test:

X = rand(2, 4)
layer = Dense(2 => 3; activation=relu)
layer(X) # 3×3 Matrix{Float64}

3.3 Reverse broadcast

Inspect the IR @code_ir layer(X):

1: (%1, %2)
  %3 = Base.getproperty(%1, :activation)
  %4 = Main.:+
  %5 = Base.getproperty(%1, :weight)
  %6 = %5 * %2
  %7 = Base.getproperty(%1, :bias)
  %8 = Base.broadcasted(%4, %6, %7)
  %9 = Base.materialize(%8)
  %10 = (%3)(%9)
  return %10

From part 1 and part 4 we have rrules for getproperty (getfield), matrix multiplication (*) and for the activation (relu). We still need rrules for broadcasted and materialize.

Creating rules for broadcasting in general is complex¹, so instead create a specific rule for the broadcast invoked here:

function rrule(::typeof(Broadcast.broadcasted), ::typeof(+), A::AbstractVecOrMat{<:Real}, B::AbstractVecOrMat{<:Real})
    broadcast_back(Δ) = (nothing, nothing, unbroadcast(A, Δ), unbroadcast(B, Δ))
    broadcast(+, A, B), broadcast_back
end

function unbroadcast(x::AbstractArray, x̄)
    if length(x) == length(x̄)
        x̄
    else
      dims = ntuple(d -> size(x, d) == 1 ? d : ndims(x̄)+1, ndims(x̄))
      dx = sum(x̄; dims = dims)
      check_dims(size(x), size(dx))
      dx
    end
end

function check_dims(size_x, size_dx) # see ChainRulesCore.ProjectTo
    for (i, d) in enumerate(size_x)
        dd = i <= length(size_dx) ? size_dx[i] : 1 # broadcasted dim
        if d != dd 
            throw(DimensionMismatch("variable with size(x) == $size_x cannot have a gradient with size(dx) == $size_dx"))
        end
    end
end

Testing:

X = rand(2, 4)
b = rand(2)
Z, back = rrule(Base.broadcasted, +, X, b) # (2×4 Matrix{Float64}, broadcast_back)
back(ones(2, 4)) # (nothing, nothing, ones(2, 4), [4.0; 4.0;;])

The definition for Base.Broadcast.materialize is:

@inline materialize(bc::Broadcasted) = copy(instantiate(bc))
materialize(x) = x

Hence we need rrules for copy and instantiate (source):

function rrule(::typeof(copy), bc::Broadcast.Broadcasted)
    uncopy(Δ) = (nothing, Δ)
    return copy(bc), uncopy
end

function rrule(::typeof(Broadcast.instantiate), bc::Broadcast.Broadcasted)
    uninstantiate(Δ) = (nothing, Δ)
    return Broadcast.instantiate(bc), uninstantiate
end

Now the pullback for the Dense layer works:

Y, back = pullback(layer, X) # (3×4 Matrix, Pullback)
back(ones(3, 4)) # ((;weight=...,bias=...,activation=nothing), 2×4 Matrix)
Y, back = pullback(m->m(X), layer) # (3×4 Matrix, Pullback)
back(ones(3, 4)) # (nothing, (;weight=...,bias=...,activation=nothing))

3.4 Chain

Here is the Flux code to create a generic chain (source):

struct Chain{T<:Tuple}
    layers::T
end
  
Chain(xs...) = Chain(xs)

(c::Chain)(x) = _apply_chain(c.layers, x)

@generated function _apply_chain(layers::Tuple{Vararg{Any,N}}, x) where {N}
  symbols = vcat(:x, [gensym() for _ in 1:N])
  calls = [:($(symbols[i+1]) = layers[$i]($(symbols[i]))) for i in 1:N]
  Expr(:block, calls...)
end

Add a method to parameters:

parameters(c::Chain) = (;layers = map(parameters, c.layers))

We will need an rrule for getindex:

world = Base.get_world_counter()
pr1 = _generate_pullback(world, typeof(_apply_chain), Tuple{typeof(cos), typeof(sin)}, Float64)

It is as follows (source):

function rrule(::typeof(getindex), x::T, i::Integer) where {T<:Tuple}
    function getindex_back_1(Δy)
        dx = ntuple(j -> j == i ? Δy : nothing, length(x))
        return (nothing, (dx...,), nothing)
    end
    return x[i], getindex_back_1
end

Test (compare the results in part 1):

model = Chain(cos, sin)
model(0.9) # 0.5823
z, back = pullback(model, 0.9)
back(1.0) # ((layers=(nothing, nothing),), -0.6368)

Test a multi-layer perceptron:

model = Chain(
    Dense(2 => 16, activation=relu),
    Dense(16 => 16, activation=relu),
    Dense(16=>2, activation=relu)
)
model(X) # 2×4 Matrix
Z, back = pullback(m->m(X), model)  # (2×4 Matrix, Pullback)
back(ones(2, 4)) # (nothing, (layers=((weight=...), (weight=...), (weight=...))))

4 Loss

4.1 Cross entropy

The output of the machine learning model will be a probability $p_j$ for a sample $j$ being in a certain class. This will be compared to a probability for a known label $y_j$, which is either 1 if that sample is in the class or 0 if it is not. An obvious value to maximise is their product:

with range $[0, 1]$.

However most machine learning optimisation algorithms aim to minimise a loss. So instead $p_j$ is scaled as $-\log(p_j)$, so that the loss ranges from $[0, \infty)$ with the goal to minimise it at 0. This is called the cross entropy loss:

4.2 Logit cross entropy

The outputs of the neural network are not probabilities but instead a vector of logits containing $N$ real values for $N$ classes. By convention these logits are scaled to a probability distribution using the softmax function:

Combining equations $\ref{eq:cross_entropy}$ and $\ref{eq:softmax}$ and taking a mean across samples gives the mean logit cross entropy loss:

where $z_{ij}$ is the output of the logsoftmax function. Assuming that $y_{ij}$ is 1 for exactly one value of $i$ and 0 otherwise, this can be simplified to:

In Julia this can be implemented as follows (source):

using StatsBase
logsoftmax(x::AbstractArray) = x .- log.(sum(exp.(x), dims=1))
function logit_cross_entropy(x::AbstractVecOrMat, y::AbstractVecOrMat)
    mean(-sum(y .* logsoftmax(x), dims=1))
end

According to the multivariable chain rule, the derivative with respect to one logit $x_{ij}$ in the vector for sample $j$ is (gradients come from the main case $k=i$ case as well as the sum in the softmax for $k\neq i$):

where $\delta_{ij}$ is the Kronecker delta. Assuming that $y_{kj}$ is 1 for one value of $k$ and 0 otherwise, this simplifies too:

In Julia this can be implemented as follows (source):

function rrule(::typeof(logsoftmax), x::AbstractArray)
    expx = exp.(x)
    Σ = sum(expx, dims=1)
    function logsoftmax_back(Δ)
        (nothing, Δ .- sum(Δ; dims=1) .* expx ./ Σ)
    end
    x .- log.(Σ), logsoftmax_back
end

function rrule(::typeof(logit_cross_entropy),  x::AbstractVecOrMat, y::AbstractVecOrMat)
    ls, logsoftmax_back = rrule(logsoftmax, x)
    function logit_cross_entropy_back(Δ)
        size_ls = size(ls)
        n = length(size_ls) > 1 ? prod(size(ls)[2:end]) : 1
        ∂x = logsoftmax_back(-y * Δ/n)[2]
        ∂y = -Δ/n .* ls
        return nothing, ∂x , ∂y
    end
    mean(-sum(y .* ls, dims = 1)), logit_cross_entropy_back
end

Testing:

y1, y2 = rand(4), rand(4)
l, back = pullback(logit_cross_entropy, y1, y2) # (2.69, logit_cross_entropy_back)
back(1.0) # (nothing, [0.4,...], [1.37,...] )
X = rand(2, 4) 
Y = [1.0 1.0 0.0 0.0 ; 0.0 0.0 1.0 1.0] # one hot encoded
l, back = pullback(logit_cross_entropy, X, Y)
back(1.0) # (nothing, 2×4 Matrix, 2×4 Matrix)

5 Train and Evaluate

5.1 Train

Create the moons data and labels:

n = 100
X = make_moons(2n; noise=0.1) # 2×200 Matrix 
y = vcat(fill(1, n)..., fill(2, n)...) # 200-element Vector{Int64}

Convert the labels to a one hot presentation:

function onehot(y::AbstractVector, labels)
    num_classes = maximum(labels)
    Y = zeros(num_classes, length(y))
    for (j, label) in enumerate(y)
        Y[label, j] += 1
    end
    Y
end
Y = onehot(y, 1:2)

Create the model:

model = Chain(
    Dense(2 => 16, activation=relu),
    Dense(16 => 16, activation=relu),
    Dense(16=>2, activation=relu)
)

Test the loss function:

l, back = pullback(m->logit_cross_entropy(m(X), Y), model); # (0.69, Pullback{...}(...))
back(1.0) # (nothing, layers=((weight=...),(weight=...),(weight=...),))

Use the exact same gradient_descent! function from part 4:

history = gradient_descent!(
    model, logit_cross_entropy, X, Y
    ; learning_rate=0.9, max_iters=200
)

5.2 Evaluate

Plot the history:

Calculate accuracy:

Y_pred = model(X)
y_pred = vec(map(idx -> idx[1], argmax(Y_pred, dims=1)))
mean(y_pred .== y) # 100%

Plot decision boundary:

using Plots
xmin, xmax = extrema(X[1, :])
ymin, ymax = extrema(X[2, :])
h = 0.01
xrange = (xmin-0.1):h:(xmax+0.1)
yrange = (ymin-0.1):h:(ymax+0.1)

x_grid = xrange' .* ones(length(yrange))
y_grid = ones(length(xrange))' .* yrange
Z = similar(x_grid)
for idx in eachindex(x_grid)
    logits = model([x_grid[idx], y_grid[idx]])
    Z[idx] = softmax(logits)[1]
end
canvas = heatmap(xrange, yrange, Z, size=(800, 500))

Plot points over the boundary:

scatter!(
    X[1, :], X[2, :], color=y, label="", aspectratio=:equal,
    xlims = xlims(canvas),
    ylims = ylims(canvas),
)

The result:

6 Conclusion

That was a long and difficult journey. I hope you understand how automatic differentiation with Zygote.jl works now!

This is part of a series. The other articles are:

• Part 1: ChainRules.
• Part 2: Automation with expressions.
• Part 3: Automation with IR.
• Part 5: MLP.

All source code can be found at MicroGrad.jl.

Table of Contents

1 Introduction

By end of part 3 we had code that could automatically differentiate many functions as long as we had rrules and there was no control flow.

However, the code failed for the polynomial model:

struct Polynomial{V<:AbstractVector}
    weights::V
end
(m::Polynomial)(x) = evalpoly(x, m.weights)
(m::Polynomial)(x::AbstractVector) = map(m, x)
model = Polynomial([3.0, 2.0, -3.0, 1.0])
x = [1.0, 2.0, 3.0, 4.0]
pullback(model, x) # ERROR: No method found for Tuple{typeof(fieldtype) ....}

Calling @code_ir model(x), we can see that code is lowered as follows:

1: (%1, %2)
  %7 = Main.map(%6, %2)
  return %7

And further that model(1.0) is lowered to:

1: (%1, %2)
  %3 = Base.getproperty(%1, :weights)
  %4 = Main.evalpoly(%2, %3)
  return %4

We could have also defined the map using an anonymous function:

(m::Polynomial)(x::AbstractVector) = map(x->evalpoly(x, m.weights), x)

In which case it would have been lowered to:

1: (%1, %2)
  %3 = Main.:(var"#43#44")
  %4 = Core.typeof(%1)
  %5 = Core.apply_type(%3, %4)
  %6 = %new(%5, %1)
  %7 = Main.map(%6, %2)
  return %7

The calls to Core.typeof and Core.apply_type are in the list of ignored functions. However we need to handle map, getproperty and %new. These sort of functions do not have formal mathematical derivatives and so they do not have rrules in ChainRules.jl. Instead, Zygote.jl handles these functions with their own custom pullbacks. Zygote also replaces some low level functions like new, getproperty and getindex entirely with custom code.

2 Extending pullback

2.1 map

The pullback for map is fairly complex. What will be presented here is a simplified version. It might also help to look at the less generic code in the example in part 1.

Consider the following code:

f(x) = sin(x)
x = [0.1, 0.2, 0.5]
map(f, x)

The pullback for map should return 3 values: $\text{s̄elf}$ for map, $\bar{f}$ for the function f and $\bar{x}$ for each value in x.

The code will start by getting pullbacks for each value in x:

ys_and_backs = map((xs...) -> pullback(f, xs...), x) # ((0.099, Pullback), (0.198, Pullback), (0.479, Pullback))

This list is in a “zipped” format: there are $n$ entries of $(y_i, \mathcal{B}_i)$ for an array length $n$. This will be unzipped into two lists each of length $n$: $(y_1,…,y_n), (\mathcal{B}_1,…,\mathcal{B}_n)$:

Δ = ones(length(x))
ys = map(first, ys_and_backs) # (0.099, 0.198, 0.479)
∂f_and_∂x_zipped = map(((_, pb), δ) -> pb(δ), ys_and_backs, Δ) # ((nothing, 0.995), (nothing, 0.980), (nothing, 0.877))

The gradients list of $n$ entries

needs to be further unzipped into $k+1$ lists for $\text{s̄elf}$ and $k$ arguments:

This is done with an unzip function which generalises first to any index i (source):

struct StaticGetter{i} end
(::StaticGetter{i})(v) where {i} = v[i]
(::StaticGetter{i})(::Nothing) where {i} = nothing

function _unzip(tuples, ::Val{N}) where {N}
  getters = ntuple(n -> StaticGetter{n}(), N)
  map(g -> map(g, tuples), getters)
end

function unzip(tuples)
  N = length(first(tuples))
  _unzip(tuples, Val(N))
end

The result:

∂f_and_∂x = unzip(∂f_and_∂x_zipped) # [nothing, nothing, nothing], [0.995, 0.98, 0.877]

As a final step, all the gradients for the function are accumulated into one value:

∂f = reduce(accum, ∂f_and_∂x[1]) # nothing

Putting all this code in a single function (source):

function pullback(::typeof(map), f::F, args::Vararg{Any, N}) where {F, N}
    ys_and_backs = map((xs...) -> pullback(f, xs...), args...)
    ys = map(first, ys_and_backs)
    function map_pullback(Δ)
      # technically should apply f in reverse and reverse back afterwards in case f is stateful
      ∂f_and_∂x_zipped = map(((_, pb), δ) -> pb(δ), ys_and_backs, Δ)
      ∂f_and_∂x = unzip(∂f_and_∂x_zipped) 
      ∂f = reduce(accum, ∂f_and_∂x[1])
      ∂args = ∂f_and_∂x[2:end]
      return (nothing, ∂f, ∂args...)
    end
    ys, map_pullback
end

Testing:

x = [0.1, 0.2, 0.5]
z, back = pullback(map, sin, x) 
back(ones(length(x))) # (nothing, nothing, [0.995, 0.98, 0.877])

And also:

f(a,b)=a/(a+b*b)
z, back = pullback(map, f, [2.0, 4.0], [3.0, 5.0]) 
back([1.0, 1.0]) # (nothing, nothing, [0.074, 0.029], [-0.099, -0.047])

2.2 Instrument

Zygote.jl modifies some of the source code before creating the primal and reverse passes. Here is a simplified version of this instrument function which only replaces new and getfield (source):

function instrument(ir::IR)
    pr = Pipe(ir)
    for (v, st) in pr
        ex = st.expr
        if isexpr(ex, :new)
            pr[v] = xcall(Main, :__new__, ex.args...)
        elseif is_literal_getfield(ex)
            pr[v] = xcall(Main, :literal_getfield, ex.args[2], Val(unwrapquote(ex.args[3])))
        end
    end
    finish(pr)
end

iscall(x, m::Module, n::Symbol) = isexpr(x, :call) && x.args[1] == GlobalRef(m, n)
unwrapquote(x) = x
unwrapquote(x::QuoteNode) = x.value

is_literal_getfield(ex) =
  (iscall(ex, Core, :getfield) || iscall(ex, Base, :getfield)) &&
  ex.args[3] isa Union{QuoteNode,Integer}

Modify the existing _generate_pullback_via_decomposition and _generate_callable_pullback functions to call it:

function _generate_pullback_via_decomposition(T, world)
    m = meta(T; world=world)
    isnothing(m) && return nothing
    ir = IR(m)
    length(blocks(ir)) == 1 || error("control flow is not supported")
    ir = instrument(ir) # new
    pr, calls = primal(ir, T)
    m, pr, calls
end

function _generate_callable_pullback(j::Type{<:Pullback{S, T}}, world, Δ) where {S, T}
    m = meta(S; world=world)
    ir = IR(m)
    isnothing(ir) && return :(error("Non-differentiable function ", repr(args[1])))
    length(blocks(ir)) == 1 || error("control flow is not supported")
    ir = instrument(ir) # new
    back = reverse_differentiate(ir)
    back = slots!(inlineable!(back))
    ci = build_codeinfo_(back)
    ci.slotnames = [Symbol("#self#"), :Δ]
    ci
end 

Now we need to define literal_getfield and __new__ and their pullbacks.

2.3 getfield

Calls to getproperty default to getfield, where a field is is declared in a struct’s declaration. The getfield function is substituted with literal_getfield (source):

literal_getfield(x, ::Val{f}) where f = getfield(x, f)

The pullback will return a NamedTuple for each field, where the gradient is Δ for the relevant field and nothing for the others (source):

@generated nt_nothing(x) = Expr(:tuple, [:($f=nothing) for f in fieldnames(x)]...)
@generated pair(::Val{k}, v, _=nothing) where k = :($k = v,)

function pullback(::typeof(literal_getfield), x, ::Val{f}) where f
  val = getfield(x, f)
  function literal_getfield_back(Δ)
    if isimmutable(x)
      dx = (; nt_nothing(x)..., pair(Val(f), Δ)...)
      (nothing, dx, nothing)
    else
      error("multable stucts not supported")
    end
  end
  val, literal_getfield_back
end

pullback(::typeof(getfield), x, field_name::Symbol) = pullback(literal_getfield, x, Val(field_name))

For example:

struct Foo
    a
    b
    c
end
foo = Foo(1.0, 'a', "hello")
z, back = pullback(getfield, foo, :b) # ('a', literal_getfield_back)
back(1.0) # (nothing, (a = nothing, b = 1.0, c = nothing), nothing)

And for the polynomial model:

z, back = pullback(model, 1.0)
back(2.3) # ((weights = [2.3, 2.3, 2.3, 2.3],), -2.3)

For the first time we have a value $\text{s̄elf}$, which is the named tuple for the fields.

2.4 new

The code now works with:

(m::Polynomial)(x::AbstractVector) = map(m, x)

It returns $\text{s̄elf}$ and $\bar{x}$:

model = Polynomial([3.0, 2.0, -3.0, 1.0])
x = [1.0, 2.0, 3.0, 4.0]
z, back = pullback(model, x)
back(ones(4)) # ((weights = [4.0, 10.0, 30.0, 100.0],), [-1.0, 2.0, 11.0, 26.0])

However with an anonymous function:

(m::Polynomial)(x::AbstractVector) = map(x->evalpoly(x, m.weights), x)

nothing is returned for $\text{s̄elf}$:

z, back = pullback(model, x)
back(ones(4)) # (nothing, [-1.0, 2.0, 11.0, 26.0])

If we inspect the primal(ir), we see that it’s because no pullbacks and hence no gradients are recorded against variable %1 (self):

1: (%1, %2)
  %3 = Main.:(var"#74#75")
  %4 = Core.typeof(%1)
  %5 = Core.apply_type(%3, %4)
  %6 = %new(%5, %1)
  %7 = Main.pullback(Main.map, %6, %2)
  %8 = Base.getindex(%7, 1)
  %9 = Base.getindex(%7, 2)
  %10 = Base.tuple(%9)
  %11 = (Pullback{Any})(%10)
  %12 = Base.tuple(%8, %11)
  return %12

The solution is to swap %new with a call to a custom function __new__ with a pullback. This function is as follows (source):

macro __splatnew__(T, args)
  esc(Expr(:splatnew, T, args))
end

@inline __new__(T, args...) = @__splatnew__(T, args)

And the pullback is (source):

using Base: RefValue
struct Jnew{T,G}
  g::G
end

Jnew{T}(g) where T = Jnew{T,typeof(g)}(g)

function pullback(::typeof(__new__), T, args...)
  x = __new__(T, args...)
  g = !ismutabletype(T) || fieldcount(T) == 0 ? nothing : grad_mut(x)
  x, Jnew{T,typeof(g)}(g)
end

@generated function (back::Jnew{T,G})(Δ::Union{NamedTuple,Nothing,RefValue}) where {T,G}
  !ismutabletype(T) && Δ == Nothing && return :nothing
  Δ = G == Nothing ? :Δ :
      Δ <: RefValue ? :(back.g[]) :
      :(accum(back.g[], Δ))
  quote
    x̄ = $Δ
    $(G == Nothing || :(back.g[] = nt_nothing($Δ)))
    (nothing, nothing, $(map(f -> :(x̄.$f), fieldnames(T))...))
  end
end

Now if we try the following (after redefining @generated function pullback and function (methodinstance::Pullback)) we should get the same results:

z, back = pullback(model, x)
back(ones(4)) # ((weights = [4.0, 10.0, 30.0, 100.0],), [-1.0, 2.0, 11.0, 26.0])

3 Gradient Descent revisited

3.1 Generic Gradient Descent

Now that we have an automatic differentiation engine, it is possible to create a much more generic gradient descent function than in part 1:

function gradient_descent!(
    model,
    loss,
    X::AbstractVecOrMat,
    Y::AbstractVecOrMat
    ; learning_rate::AbstractFloat=0.1,
    max_iters::Integer=100
    )
    losses = Float64[]
    for i in 1:max_iters
        loss_iter, back = pullback(model) do m
            result = m(X)
            loss(result, Y)
        end 
        Δf, Δm = back(1.0)
        update_params!(parameters(model), Δm; learning_rate=learning_rate)
        push!(losses, loss_iter)  
    end
    losses
end

Note that pullback(m->f(m), model) is directly equivalent to pullback(model) do f(m) end.

The update_params! function is defined as follows:

function update_params!(params::NamedTuple, grads::NamedTuple; options...)
    for key in keys(params)
        update_params!(params[key], grads[key]; options...)
    end
end

function update_params!(params::Tuple, grads::Tuple; options...)
    for (p, g) in zip(params, grads)
        update_params!(p, g; options...)
    end
end

function update_params!(params, grads; learning_rate::AbstractFloat=0.1)
    params .-= learning_rate .* grads # must broadcast to edit elements and not copies!
end

The parameters function is defined per model. (Flux uses the generic Functors.jl library to accomplish something similar.)

3.2 Polynomial curve fitting revisited

Let’s create the exact same data set from part 1:

using StatsBase
target_weights = [15.0, -2.1, 13.9, 1.5]
noise_factor = 0.2
xs = (rand(100) .- 0.5) .* 10
ys = map(x -> evalpoly(x, target_weights), xs)
scale_factor = mean(abs.(ys))
ys .+= randn(length(ys)) * scale_factor * noise_factor

The Polynomial model is defined in the introduction. We also need a custom method for parameters:

parameters(m::Polynomial) = (;weights=m.weights)

Define the model:

model = Polynomial(rand(4))

Some sanity checks:

x = [1.0, 2.0, 3.0]
z, back = pullback(model, x) # ([1.68, 7.21, 21.2], Pullback) 
back([1.0, 1.0, 1.0]) # ((weights = [3.0, 6.0, 14.0, 36.0],), [-1.0, 2.0, 11.0])
z, back = pullback(m->m(x), model) 
back([1.0, 1.0, 1.0]) # (nothing, (weights = [3.0, 6.0, 14.0, 36.0],))
y = [2.0, 4.0, 8.0]
z, back = pullback(m->mse(m(x), y), model) 
back(1.0) # (nothing, (weights = [10.7 30.5, 87.6, 254.6],))

Train the model:

history = gradient_descent!(model, mse, xs, ys; learning_rate=1e-5, max_iters=2000)

This works just as well as before.

4 Conclusion

We now have a fully working AD package. It has some limitations, such as it cannot handle control flow or keyword arguments. However it can already work on a wide variety of code. All that might be needed is new rrule definitions. The next and final part of this series is a demonstration of exactly that.

This is part of a series. The other articles are:

• Part 1: ChainRules.
• Part 2: Automation with expressions.
• Part 4: Extensions.
• Part 5: MLP.

All source code can be found at MicroGrad.jl. The code here is based on the example at IRTools.jl.

Table of Contents

1 Introduction

Part 1 introduced the rrule for implementing chain rules and Part 2 defined a @generated pullback function for inspecting and decomposing complex code. The goal here is to replicate the results of Part 2 except in a more robust manner using the IRTools.jl package.

For example, from part 1 there are rrules for +, * and /. The goal is then to automatically differentiate the following:

like so:

f(a, b) = a / (a + b*b)
z, back = pullback(f, 2.0, 3.0) # (0.1818, ∂(f))
back(1.0) # (nothing, 0.0744, -0.099)

where pullback is a @generated function that inspects the Intermediate Representation (IR) code for f:

using IRTools
ir = @code_ir f(2, 3)
#= 1: (%1, %2, %3)
  %4 = %3 * %3
  %5 = %2 + %4
  %6 = %2 / %5
  return %6
=#

This is an advanced use of the Julia programming language. You should be comfortable with the language before reading this post. At the very least, the Julia documentation page on metaprogramming is required for this post and will be considered assumed knowledge, especially the sections on “Expressions and evaluation”, “Code Generation” and “Generated Functions”. I also suggest going through the IRTools.jl documentation first.

This post can be read independently to Part 2 and will repeat parts of it. However it is advised to read Part 2 first because it is easier to understand than this post.

2 Differentiating Wengert Lists

The Zygote.jl automatic differentiation (AD) package is a realisation of the paper Don’t Unroll Adjoint: Differentiating SSA-Form Programs (2019) by Michael J Innes.
The paper works with Wengert lists, also known as tapes, and a generalisation of it called Static Single Assignment (SSA) form. The aim here is to develop a minimal AD package, so this series only focuses on the sections on Wengert lists. A consequence is that the code will not be to handle any non-linear logic in Julia, for example any control flow like if, while or for blocks.

The paper uses the same example as the introduction:

This can be broken down into smaller steps where each intermediate variable is saved. This is known as a Wengert list, or tape, or (backpropagation) graph:

To differentiate this, all function calls are wrapped with a differentiation function $\mathcal{J}$ which returns both the output $y$ and a pullback function $\mathcal{B}$. This is called the primal form:

The pullback function $\mathcal{B}$ takes as input the gradient of a scalar $l$ (typically a loss function) to a function $y(x)$ and returns the gradient with regards to the variable $x$. This partial gradient $\frac{\partial l}{\partial x}$ is written as $\bar{x}$.

so we can write in this mathematical notation as:

where $\bar{y}=\frac{\partial l}{\partial y}$ and $J=\frac{\partial y}{\partial x}$ is the Jacobian (gradient) for arrays.

The various partial gradients are calculated by reversing the list. Each pullback function $\mathcal{B}_i$ takes as input the previous gradient $\bar{y}_i$. The input is an existing gradient $\Delta$. At the start this is usually set to 1:

The final step is to accumulate the gradients for variables which are used multiple times:

This end result is equivalent to rolling everything up into one function using the multivariable chain rule:

3 Pullback

3.1 Definition

The goal is to generate code which automatically implements the equations of section 2.

To start, define a pullback function (source):

function pullback end

This will be turned into a generated function.

Julia changed the behaviour of generated functions in version 1.10. Before 1.10, they always had access to the world age counter. This is a single number that is incremented every time a method is defined, and helps optimise compilations. However from version 1.10 generated functions Base.get_world_counter() will only return typemax(UInt). This is to prevent reflection - code inspection - in generated functions.¹ However the code here relies on reflection. Thankfully, there is a hack that Zygote.jl uses to access the world age in pullback. Because of this, the definition of pullback is different based on the version, but both will forward to a common internal _generate_pullback function.

@generated function pullback(f, args...)
        _generate_pullback(nothing, f, args...)
end

function _pullback_generator(world::UInt, source, self, f, args)
        ret = _generate_pullback(world, f, args...)
        ret isa Core.CodeInfo && return ret
        stub = Core.GeneratedFunctionStub(identity, Core.svec(:methodinstance, :f, :args), Core.svec())
        stub(world, source, ret)
end

@eval function pullback(f, args...)
        $(Expr(:meta, :generated, _pullback_generator))
        $(Expr(:meta, :generated_only))
end

3.2 ChainRules

The first goal of _generate_pullback will be to forward the function and its arguments to a matching rrule if it exists. For now it will throw an error if it cannot find one.

function _generate_pullback(world, f, args...)
    T = Tuple{f, args...}
    if (has_chain_rrule(T, world))
        return :(rrule(f, args...))
    end
    :(error("No rrule found for ", repr($T)))
end

In part 1 the most generic method of rrule was defined for an Any first argument, so if the compiler dispatches to this method it means no specific rrule was found.²

using IRTools: meta
function has_chain_rrule(T, world)
    Tr = Tuple{typeof(rrule), T.parameters...}
    meta_T = meta(Tr; world=world)
    if isnothing(meta_T)
        return false
    end
    method_ = meta_T.method
    sig = method_.sig
    !(sig isa DataType) || (sig.parameters[2] !== Any)
end

Let’s test all this code from bottom to top for a function with an rrule and one without: + and f(a,b)=a/(a+b*b). As a reminder, generated functions only have access to a variables types, so to test the _generate_pullback and all functions under it, we can only work with the types.

Firstly, for + acting on floats (redefine @generated pullback if necessary):

world = Base.get_world_counter()
T = Tuple{typeof(+), Float64, Float64}
has_chain_rrule(T, world) # true
_generate_pullback(world, typeof(+), Float64, Float64) # :(rrule(f, args...))
pullback(+, 1.0, 2.0) # (3.0, var"#add_back#5"())

Now for f, also acting on floats:

world = Base.get_world_counter()
T = Tuple{typeof(f), Float64, Float64}
has_chain_rrule(T, world) # false
_generate_pullback(world, typeof(f), Float64, Float64) # :(error(...))
pullback(f, 1.0, 2.0) # ERROR: No rrule found for ...

The more interesting task is to inspect f and apply the equations of section 2 to fully differentiate with respect to all input parameters.

3.3 IR

The first step is to create a Wengert list for f in Intermediate Representation (IR) form. Julia already does this as part of the compilation process. IRTools.jl mimics this internal IR form with its own custom IR struct. It can be generated as follows:

using IRTools: IR, meta
T = Tuple{typeof(f), Float64, Float64}
m = meta(T; world=Base.get_world_counter())
ir = IR(m)
#=
1: (%1, %2, %3)
  %4 = %3 * %3
  %5 = %2 + %4
  %6 = %2 / %5
  return %6
=#

The returned object corresponds exactly to $\ref{eq:f_wengert}$.

Using this knowledge, we can now create a new function _generate_pullback_via_decomposition which will be called if no rrule exists. It uses the IR to create the primal (equation $\ref{eq:primal}$) (source).

using IRTools: meta, IR, blocks
function _generate_pullback_via_decomposition(T, world)
    m = meta(T; world=world)
    isnothing(m) && return nothing
    ir = IR(m)
    length(blocks(ir)) == 1 || error("control flow is not supported")
    pr, calls = primal(ir, T)
    m, pr, calls
end

3.4 Primal

The goal here is to create an IR for equation $\ref{eq:primal}$. This is what it will look like:

1: (%1, %2, %3)
  %4 = Main.pullback(Main.:*, %3, %3)
  %5 = Base.getindex(%4, 1)
  %6 = Base.getindex(%4, 2)
  %7 = Main.pullback(Main.:+, %2, %5)
  %8 = Base.getindex(%7, 1)
  %9 = Base.getindex(%7, 2)
  %10 = Main.pullback(Main.:/, %2, %8)
  %11 = Base.getindex(%10, 1)
  %12 = Base.getindex(%10, 2)
  %13 = Base.tuple(%6, %9, %12)
  %14 = (Pullback{Tuple{typeof(f), Float64, Float64}})(%13)
  %15 = Base.tuple(%11, %14)
  return %15

Although harder to read, this code represents the same code as the expressions in part 2.

The primal function first wraps the existing IR with Pipe to make inserts more efficient. It defines two arrays to store information (source):

using IRTools: block, isexpr, finish, Pipe, Variable, return!, returnvalue, stmt, xcall
function primal(ir::IR, T=Any)
    pr = IRTools.Pipe(ir)
    calls = []
    pullbacks = []

The calls array stores the subset of variables that require a pullback. Because the IR is a dictionary - ir[Variable(i)] returns statement i - this creates a direct link to the statement called. These will be used to generate the reverse code (equation $\ref{eq:reverse}$) in the next section.

Next, iterate over each statement in the IR. For each statement if it is an expression :call and not part of a special ignored list, replace it with three calls: the first is to pullback and then two calls to getindex to get the output variable v and back function J from the tuple t:

    for (v, st) in pr
        ex = st.expr
        if isexpr(ex, :call) && !ignored(ex)
            t = insert!(pr, v, stmt(xcall(Main, :pullback, ex.args...), line = st.line))
            pr[v] = xcall(Base, :getindex, t, 1)
            J = push!(pr, xcall(:getindex, t, 2))
            push!(calls, v)
            push!(pullbacks, J)
        end
    end

After working through all the statements, a final statement is added which returns a tuple with the output of the function and a Pullback struct which stores all the pullbacks. In the last step the pipe is converted back into an IR.

    pb = Expr(:call, Pullback{T}, xcall(:tuple, pullbacks...))
    return!(pr, xcall(:tuple, returnvalue(block(ir, 1)), pb))
    finish(pr), calls
end

This code requires a definition for the Pullback struct as well as the ignored function.

There are no closures in lowered Julia code, so instead Zygote.jl stores the pullbacks in a generic struct:

struct Pullback{S,T}
    data::T
end
Pullback{S}(data) where S = Pullback{S,typeof(data)}(data)

In the next section this struct will be turned into a callable struct. That is, for back=Pullback{S}(data), we will create a generated function that dispatches on itself: (j::Pullback)(Δ) so that we can call back(Δ). This back has all the information to generate the reverse pass independently of the forward pass: the method can be retrieved using meta(S) and the relevant data and input parameters from back.data.

Here is the ignored functions list (source):

function ignored(ex::Expr)
    f = ex.args[1]
    ignored_f(f)
end

ignored_f(f) = f in (
    GlobalRef(Base, :not_int),
    GlobalRef(Core.Intrinsics, :not_int),
    GlobalRef(Core, :(===)),
    GlobalRef(Core, :apply_type),
    GlobalRef(Core, :typeof),
    GlobalRef(Core, :throw),
    GlobalRef(Base, :kwerr),
    GlobalRef(Core, :kwfunc),
    GlobalRef(Core, :isdefined)
)