
[Arc Institute](https://arcinstitute.org/) recently unveiled the [Virtual Cell Challenge](https://virtualcellchallenge.org/). Participants are required to train a model capable of predicting the effect of silencing a gene in a (partially) unseen cell type, a task they term _context generalization_.
For ML engineers with little to no biology background, the jargon and required context can seem quite daunting. To encourage participation, we recapitulate the challenge in a form better suited to engineers from other disciplines. 

# Background

> **Goal** \
> Train a model to predict the effect on a cell of silencing a gene using CRISPR.

Doing things in the world of atoms is expensive, laborious and error prone. What if we could test thousands of drug candidates without ever touching a petri dish?
This is the goal of the virtual cell challenge — a model (most likely a neural network) that can simulate exactly what happens
to a cell when we change some parameter. Given that tightening your feedback loop is often the best way to speed up progress,
a model capable of doing this accurately would have significant impact.

To train this neural network, we will need data. For the challenge, Arc has curated a dataset of ~300k single-cell RNA sequencing profiles. It may be worthwhile to revisit the [Central Dogma](https://www.khanacademy.org/science/biology/gene-expression-central-dogma/central-dogma-transcription/v/rna-transcription-and-translation) before continuing. This essay will build off of this to provide the ~minimum biology knowledge you'll need for the challenge.

## Training data

The training set consists of a sparse matrix and some associated metadata. More specifically, we have 220k cells, and
for each cell we have a [transcriptome](https://en.wikipedia.org/wiki/Transcriptome). This transcriptome is a sparse row vector, where each
entry is the **raw count of RNA molecules** (transcripts) that the corresponding gene (our column) encodes for. Of the 220k cells,
~38k are _unperturbed_, meaning no gene has been silenced using CRISPR. These control cells are crucial as we will see shortly. 

To understand the dataset more concretely, let's select a gene, TMSB4X (the most frequently silenced gene in the dataset) and compare the number of RNA molecules detected for a control cell and a
perturbed cell. 

![TMSB4X](/avcc/TMSB4X.png)

We can see that the cell with TMSB4X silenced has a greatly reduced number of transcripts compared with the control
cells.

### Modelling the challenge 

The astute among you may be wondering why you don't just measure the count of the RNA molecules before and after
silencing the gene — why do we need the control cells at all? Unfortunately, **reading the transcriptome destroys the cell**, which is a problem reminiscent of the [observer effect](https://en.wikipedia.org/wiki/Observer_effect_(physics)). 

This inability to measure the cell state before and after introduces many issues, as we are forced to use a population of **basal**
(a.k.a control, unperturbed) cells as a reference point. The control cells and perturbed cells are not entirely
homogeneous even prior to the perturbation. This means that we have to now separate out our true signal, the perturbation, from
noise induced by the heterogeneity.

More formally, we can model observed gene expression in perturbed cells as:

$$
\hat{X}_p \sim \hat{T}_p(\mathcal{D}_{\text{basal}}) + H(\mathcal{D}_{\text{basal}}) + \varepsilon, \quad \varepsilon \sim P_\varepsilon 
$$

where:
- $\hat{X}_p$: The observed gene expression measurements in cells with perturbation $p$
- $\mathcal{D}_{\text{basal}}$: The distribution of the unperturbed, baseline cell population.
- $\hat{T}_p(\mathcal{D}_{\text{basal}})$: True effect caused by perturbation $p$ on the population.
- $H(\mathcal{D}_{\text{basal}})$: Biological heterogeneity of the baseline population.
- $\varepsilon$: Experiment-specific technical noise, assumed independent of the unperturbed cell state and $\mathcal{D}_{\text{basal}}$.

# STATE: The baseline from Arc

Prior to the Virtual Cell Challenge, Arc released [STATE](https://arcinstitute.org/manuscripts/State), their own attempt to solve the challenge
using a pair of transformer based models. This serves as a strong baseline for participants to start with, so we will
explore it in detail.

STATE consists of two models, the **State Transition Model** (ST) and the **State Embedding Model** (SE). SE is designed to produce rich semantic embeddings of cells in an effort to improve cross cell type generalization. ST is the "cell simulator", that takes in either a transcriptome of a control cell, or an embedding of a cell produced by SE, along with a one hot encoded vector representing the perturbation of interest, and outputs the perturbed transcriptome. 

## State Transition Model (ST) 

![ST](/avcc/ST.png)

The State Transition Model is a relatively simple transformer with a Llama backbone that operates upon the following: 
1. A set of transcriptomes (or SE embeddings) for covariate matched basal cells.
2. A set of one hot vectors representing our gene perturbation for each cell.

Using a covariate matched set of control cells with paired target cells should assist the model in discerning the
actual effect of our intended perturbation. Both the control set tensor and the perturbation tensor are fed through independent encoders, which are simply 4 layer MLPs with GELU activations.
If working directly in gene expression space (i.e producing a full transcriptome), they pass the output through a learned
decoder.

ST is trained using [Maximum Mean Discrepancy](https://en.wikipedia.org/wiki/Kernel_embedding_of_distributions). Put simply, the model learns to minimize the difference between the two probability distributions.

## State Embedding Model (SE) 

![SE](/avcc/SE.png)

The State Embedding Model is a BERT-like model trained using a masked prediction task. To understand this more deeply, first we have to
take a little detour for some more biological grounding. 

### A little biological detour 

![Alternative Splicing](/avcc/alt_splicing.gif)

A gene consists of _exons_ (protein coding sections) and _introns_ (non-protein coding sections). DNA is first _transcribed_ into pre-mRNA, as shown above. The cell then performs [Alternative Splicing](https://en.wikipedia.org/wiki/Alternative_splicing). This is basically "pick and choose exons", cut out all introns. You can think of the gene as an IKEA manual for making a table. One could also construct a 3 legged table, perhaps an odd bookshelf with some effort, by leaving out some parts. These different objects are analogous to **protein isoforms**, proteins coded for by the same gene.

### Back to the model

With this basic understanding, we can move on to how the SE model works. Remember, our core goal for SE is to create **meaningful
cell embeddings**. To do this, we must first create meaningful gene embeddings.

To produce a single gene embedding, we first obtain the amino acid sequence (e.g $\texttt{SDKPDMAEI}$... for TMSB4X) of all the different protein isoforms encoded for by the gene in question. We then feed these sequences to [ESM2](https://huggingface.co/facebook/esm2_t48_15B_UR50D), a 15B parameter Protein Language Model from FAIR. ESM produces an embedding _per amino acid_, and we mean pool them together to obtain a "transcript" (a.k.a protein isoform) embedding. 

Now we have all of these protein isoform embeddings, we then just mean pool those to get the gene embedding. Next, we project these gene embeddings to our model dimension using a learned encoder as follows:

$$
\tilde{g}_j = \text{SiLU}(\text{LayerNorm}(g_j \mathbf{W}_g + \mathbf{b}_g))
$$

We've now obtained a gene embedding, but what we really want is a _cell embedding_. To do this, Arc represents each cell
as the top 2048 genes ranked by [log fold expression level](https://en.wikipedia.org/wiki/Fold_change#Fold_changes_in_genomics_and_bioinformatics).

We then construct a "cell sentence" from our 2048 gene embeddings as follows:
$$
\tilde{\mathbf{c}}^{(i)} = \left[\mathbf{z}_{\text{cls}}, \tilde{\mathbf{g}}_1^{(i)}, \tilde{\mathbf{g}}_2^{(i)}, \ldots, \tilde{\mathbf{g}}_L^{(i)}, \mathbf{z}_{\text{ds}}\right] \in \mathbb{R}^{(L+2) \times h}
$$

We add a $\texttt{[CLS]}$ token and $\texttt{[DS]}$ token to our sentence. The $\texttt{[CLS]}$ token ends up being used as our "cell embedding" (very BERT-like)
and the $\texttt{[DS]}$ token is used to "disentangle dataset-specific effects". Although the genes are sorted by log fold
expression level, Arc further enforces the magnitude of each genes expression by incorporating the transcriptome in a
fashion analogous to positional embeddings. Through an odd ["soft binning" algorithm](https://github.com/ArcInstitute/state/blob/main/src/state/emb/nn/model.py#L374) and 2 MLPs, they create some
"expression encodings" which they then add to each gene embedding. This should modulate the magnitude of each gene
embedding by how intensely it is expressed in the transcriptome.

To train the model, they mask 1280 genes per cell, and the model is tasked with predicting them. The 1280 genes are
selected such that they have a wide range of expression intensities. For the graphically inclined, the below
demonstrates the construction of the cell sentence. 

![SE PATH](/avcc/SE_path.png)

# Evaluations 

Understanding how your submission will be evaluated is key to success. The 3 evaluation metrics chosen by Arc are **Perturbation Discrimination**, **Differential Expression** and **Mean Average Error**. Given that Mean Average Error is simple and exactly as it sounds, we will omit it from our analysis.

## Perturbation Discrimination

![Perturbation Discrimination](/avcc/pert_disc.png "Perturbation Discrimination")

Perturbation Discrimination intends to evaluate how well your model can uncover _relative differences_ between
perturbations. To do this, we compute the Manhattan distances for all the measured perturbed transcriptomes in the test set (the ground
truth we are trying to predict, $y_t$ and all other perturbed transcriptomes, $y_p^n$) to our predicted transcriptome $\hat{y}_t$. We then rank where the
ground truth lands with respect to all transcriptomes as follows:

$$
r_t = \sum_{p \neq t} \mathbf{1}\{d(\hat{y}_t, y_p) < d(\hat{y}_t, y_t)\}
$$

After, we normalize by the total number of transcriptomes:

$$
\text{PDisc}_t = \frac{r_t}{T}
$$

Where $0$ would be a perfect match. The overall score for your predictions is the mean of all $$\text{PDisc}_t$$. This is then normalized to: 

$$
\text{PDiscNorm} = 1 - 2\text{PDisc}
$$ 

We multiply by 2 as for a random prediction, ~half of the results would be closer and half would be further away.

## Differential Expression

Differential Expression intends to evaluate what fraction of the truly affected genes did you correctly identify as significantly affected. Firstly, for each gene compute a $p$-value using a [Wilcoxon rank-sum test with tie correction](https://en.wikipedia.org/wiki/Mann%E2%80%93Whitney_U_test). We do this for both our predicted perturbation distribution and the ground truth perturbation distribution.

Next, we apply the [Benjamini-Hochberg procedure](https://en.wikipedia.org/wiki/False_discovery_rate#Benjamini%E2%80%93Hochberg_procedure), basically some stats to modulate the $p$-values, as with $20,000$ genes and a $p$-value threshold of $0.05$, you'd expect $1,000$ false positives. We denote our set of predicted differentially expressed genes $G_{p,pred}$, and the ground truth set of differentially expressed genes $G_{p,true}$. 

If the size of our set is less than the ground truth set size, take the intersection of the sets, and divide by the true number of differentially expressed genes as follows:

$$
DE_p = \frac{G_{p,pred} \cap G_{p,true}}{n_{p,true}}
$$

If the size of our set is greater than the ground truth set size, select the subset we predict are most differentially expressed (our "most confident" predictions, denoted $\tilde{G}_{p,pred}$), take the intersection with the ground truth set, and then divide by the true number.

$$
DE_p = \frac{\tilde{G}_{p,pred} \cap G_{p,true}}{n_{p,true}}
$$

Do this for all predicted perturbations and take the mean to obtain the final score.

# Conclusion

If a virtual cell can accurately model the change in a cells state in response to perturbations, we can look forward to a
very interesting time in pharma. I hope this post accelerated your understanding of the challenge. 

May the best team win.


Arc recently unveiled the Virtual Cell Challenge - we explore the challenge from a non-biologist's perspective in the hope of encouraging participation from a range of disciplines.

Arc Virtual Cell Challenge: A Primer

With exploding demand for AI inference, many hardware startups are designing
Domain Specific Architectures. Working backwards from the Transformer workload, we will identify optimal design
choices, promising hardware and how to predict the future of inference.


Domain specific architectures for AI inference

Rotary Position Encoding (RoPE) can seem intimidating at first. This post endeavours to rediscover RoPE by iteratively improving upon proposed positional encoding schemes.

You could have designed state of the art Positional Encoding

Smartglasses are a form of Augmented Reality (AR) providing significant compute power. What are the current fundamental limitations prohibiting the consumer adoption of smartglasses, and when can we expect their widespread adoption?

A first principles analysis of Smartglasses


LayerNorm (and its close sibling RMSNorm) have superseded batch normalization as the go-to
normalization technique for deep learning. Some kind of normalization is
essential in stabilizing inputs to each layer ensuring the model can learn
efficiently. LLaMA, Whisper and other recent transformer architectures all use (Layer|RMS)Norm. As a result,
if you want to run these **big** models in the browser on the GPU ([like me](https://github.com/FL33TW00D/whisper-turbo)), you need to implement
a fast LayerNorm kernel yourself.

Of all the kernels I wrote, I think LayerNorm is the most insightful for less seasoned GPU programmers. In this post we will explore how we can iteratively arrive at
quite a competitive WebGPU LayerNorm kernel, whilst covering some of the cornerstones of GPU programming.

## The problem

On the surface, LayerNorm seems like a simple operation - however it is not naively parallelizable. Let's see how PyTorch defines LayerNorm in [their documentation](https://pytorch.org/docs/stable/generated/torch.nn.LayerNorm.html):

$$
\begin{align}
y = \frac{x - E[x]}{\sqrt{Var[x] + \epsilon}} * \gamma + \beta
\end{align}
$$

$x$ is the input tensor, $\gamma$ and $\beta$ are learnable parameters, and $\epsilon$ is a small constant to avoid division by zero. Looking at this formula, the first thing to note as a GPU programmer is that it requires 2 group statistics: mean and variance.
This means that we can't immediately parallelize the
computation of each output element. First, we need to compute the mean and variance along the
dimension we intend to normalize, and then use those to compute each output element.
Computing the mean is straightforward, but the variance is a bit more complicated.

## Reductions

Whenever we are required to compute some kind of total or group statistic, your first thought
should be to use a reduction. Reductions are a core primitive of GPU programming, and it is
**essential** you build an understanding of them. Fortunately, Mark Harris of
NVIDIA has written a [timeless resource](https://developer.download.nvidia.com/assets/cuda/files/reduction.pdf)
on the topic. I highly recommend you read it before proceeding. Without the use
of reductions, you'll find yourself unable to parallelize many common operations.

However, if you are short on time, here is a quick visual that
should convey the core idea of a reduction:

![Reduction](/reduction.png)

Above you can see a summation of 16 elements using 8 threads. Each thread
accesses 2 elements, offset by the number of active invocations in that
iteration.

## One pass algorithm

Now that we have a basic understanding of reductions, let's see how we can use them!
The most obvious way to compute the mean and variance for our LayerNorm is to use a one-pass algorithm.
In the most common case, we will be normalizing the last dimension of a tensor,
where each workgroup computes the mean and variance of a single row.

This method requires the use of shared memory, where all threads in a block share data and synchronize their execution. Each thread will accumulate 2 partial results: the sum and the squared sum. These can then be stored in shared memory, and reduced to a single value. The mean and variance can then be computed from the reduced values.

The kernel looks something like this ([Full Code](https://github.com/FL33TW00D/wgpu-bench/blob/master/kernels/layernorm_scalar_onepass.wgsl)):

```wgsl 
const BLOCK_SIZE: u32 = 128u;
var<workgroup> sum: array<f32, BLOCK_SIZE>;
var<workgroup> sq_sum: array<f32, BLOCK_SIZE>;
var<workgroup> mean: f32;
var<workgroup> sigma: f32;

fn block_reduce(index: u32, stride: u32) {
    if index < stride {
        sum[index] += sum[index + stride];
        sq_sum[index] += sq_sum[index + stride];
    }
    workgroupBarrier();
}

fn main() {
    //Each thread computes the partial sum and squared sum
    for (var i = local_id.x; i < metadata.N; i += BLOCK_SIZE) {
       let val = X[anchor + i];
       sum[local_id.x] += val;
       sq_sum[local_id.x] += val * val;
    }
    workgroupBarrier(); //Wait for all threads to finish

    block_reduce(local_id.x, 64u);
    //... manual unroll ...
    block_reduce(local_id.x, 1u);

    if local_id.x == 0u {
        mean = sum[0] / f32(metadata.N);
        //We use inverseSqrt for performance, so technically this isn't sigma
        sigma = inverseSqrt((sq_sum[0] - (mean * mean) / f32(metadata.N)) + metadata.eps);
    }
    workgroupBarrier();
    //...compute LN...
}
```

We process the data in blocks of 128 elements, and use a manually unrolled reduction
to reduce down to our 2 final values. Let's see how it performs:

```markdown
One-Pass
time: [124979.2722 ns 126546.2184 ns 128406.2134 ns]
thrpt: [30.4210 GiB/s 30.8682 GiB/s 31.2552 GiB/s]

2 samples not close - AVGE=0.0000016181261 MAE=0.000022888184 at [0, 137, 615]
```

Whilst the performance is good, we have a problem. The output isn't matching PyTorch!
What's going on here?

Unfortunately, we have bumped into a common problem in GPU programming: [Catastrophic cancellation](https://en.wikipedia.org/wiki/Catastrophic_cancellation).
On the highlighted line (line 31) in the code sample above, we do the following (we use `inverseSqrt` for peformance,
but it's the same idea):

$$
\begin{align}
\sigma &= \sqrt{\sum{x^2} - \frac{\mu^2}{N} + \epsilon} \\
\end{align}
$$

If the squared sum is close to the mean squared over $N$, computing the difference will lead to a significant loss of precision.
Whenever we have 2 good approximations of values of a similar magnitude, subtracting them can lead to a very poor approximation of the true value.
For a formal analysis of this phenomenon, I recommend exploring the Wikipedia article linked above.

Whilst catastrophic cancellation can be mitigated by using higher precision types
or something like [Kahan summation](https://en.wikipedia.org/wiki/Kahan_summation_algorithm), these both have performance implications.
Is there a way we can avoid this problem without sacrificing performance?

## Two-pass algorithm

Next, let's try a two-pass algorithm to avoid the troublesome
subtraction in the one-pass algorithm. By first computing the mean, and
subtracting it from each element as we compute the variance, we can avoid
catastrophic cancellation. The kernel looks something like this ([Full Code](https://github.com/FL33TW00D/wgpu-bench/blob/master/kernels/layernorm_scalar.wgsl)):

```wgsl
fn mu(local_id: vec3<u32>, anchor: u32) -> f32 {
    var threadSum = 0f;
    for (var i: u32 = local_id.x; i < metadata.N; i += BLOCK_SIZE) {
        threadSum += X[anchor + i];
    }
    smem[local_id.x] = threadSum;
    workgroupBarrier();

    block_reduce(local_id.x, 64u);
    ///... manual unroll ...
    block_reduce(local_id.x, 1u);

    return smem[0] / f32(metadata.N);
}

fn sigma(local_id: vec3<u32>, anchor: u32, mu: f32) -> f32 {
    var threadSum = 0f;
    for (var i: u32 = local_id.x; i < metadata.N; i += BLOCK_SIZE) {
        let val = X[anchor + i] - mu;
        threadSum += (val * val);
    }
    smem[local_id.x] = threadSum;
    workgroupBarrier();

    block_reduce(local_id.x, 64u);
    ///... manual unroll ...
    block_reduce(local_id.x, 1u);

    return smem[0] / f32(metadata.N);
}
```

This is very similar to the one-pass algorithm, but will obviously be much slower
as we need to perform two passes over the data. Let's see how it performs:

```markdown
Two-Pass
time: [158307.3426 ns 159382.1927 ns 160646.1034 ns]
thrpt: [24.3159 GiB/s 24.5087 GiB/s 24.6751 GiB/s]

All close - AVGE=0.0000005285047 MAE=0.0000076293945 at [0, 18, 370]
```

Great, we match PyTorch's output! But we've taken a significant performance hit.
Is there a way we can get the best of both worlds?

## Welford's algorithm

Fortunately, there is a well known one-pass algorithm that gives us the best of
both worlds, with high performance and numerical precision. This algorithm is known as [Welford's algorithm](https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Welford's_online_algorithm).
Welford's algorithm is an [online algorithm](https://en.wikipedia.org/wiki/Online_algorithm), allowing us to build up the complete solution from partial solutions. To do this, we need to form a [recurrence relation](https://en.wikipedia.org/wiki/Recurrence_relation) for each component.

Let's look at the relation for the mean:

$$
\begin{align}
\mu_{k} &= \mu_{k-1} + \frac{x_k - \mu_{k-1}}{k} \\
\end{align}
$$

We can update our mean by adding the difference between the
current value and the previous mean, divided by the new number of elements (our
new values "contribution" to the mean).

Let's do the same for variance:

$$
\begin{align}
\sigma_{k}^2 &= \sigma_{k-1}^2 + \frac{(x_k - \mu_{k-1})(x_k - \mu_{k}) - \sigma_{k-1}^2}{k} \\
\end{align}
$$

This is a bit more complicated, but it's the same idea. We can update our variance
using only the previous variance, mean and the new value. In practice, the above formula for variance is not used, and the algorithm
originally stated by Welford uses the sum of squares of differences from the current mean as the quantity for updating.
This is known as $M_{2,k}$, and is defined as:

$$
\begin{align}
M_{2,k} &= M_{2,k-1} + (x_k - \mu_{k-1})(x_k - \mu_{k}) \\
\sigma_{k}^2 &= \frac{M_{2,k}}{k} \\
\end{align}
$$

Now that we have these recurrence relations, threads can now compute "partial
means" and "partial variances", and using a reduction we can construct
the full solution.

## Implementing Welford's Algorithm

To achieve peak performance with Welford's algorithm, we are going to make use of the
bleeding edge [WebGPU Subgroups Proposal](https://github.com/gpuweb/gpuweb/pull/4368).

For those unfamiliar with the term, a subgroup is
a generalization of the instrisic thread grouping of the underlying GPU hardware. For CUDA
it's known as a warp, wave for DirectX and simdgroup for Metal. The group
of threads executes the same instruction in parallel, **and** can share data
between their registers. This circumvents the need for shared memory, and
can lead to significant speedups. Designing your algorithms with the thread grouping
in mind is essential for achieving peak performance.

Let's take a look at the core components of Welford's algorithm:

```wgsl
fn welford_combine(val: f32, mean: ptr<function, f32>, m2: ptr<function, f32>, count: ptr<function, f32>) {
    *count += 1.0;
    let delta1 = val - *mean;
    *mean += delta1 / *count;
    let delta2 = val - *mean;
    *m2 += delta1 * delta2;
}
```

This is the heart of the algorithm, which updates the running mean and $M_{2}$ defined
above. Once each thread has computed its partial mean and $M_{2}$, we can combine them
together to form the full solution. This is where the subgroup extension comes in:

```wgsl
fn welford_warp_reduce(thread_mean: f32, thread_m2: f32, thread_count: f32, mean: ptr<function, f32>, m2: ptr<function, f32>, count: ptr<function, f32>) {
    *mean = thread_mean;
    *m2 = thread_m2;
    *count = thread_count;
    //Subgroup reduce into first lane
    for (var offset = subgrp_size >> 1u; offset > 0u; offset >>= 1u) {
        let b_mean = subgroupShuffleDown(*mean, offset);
        let b_m2 = subgroupShuffleDown(*m2, offset);
        let b_count = subgroupShuffleDown(*count, offset);
        block_welford_combine(b_mean, b_m2, b_count, mean, m2, count);
    }

    //Broadcast result to all threads
    *mean = subgroupBroadcast(*mean, 0u);
    *m2 = subgroupBroadcast(*m2, 0u);
    *count = subgroupBroadcast(*count, 0u);
}
```

This uses the new `subgroupShuffleDown` function, which passes the value from the active invocation at `lane + offset` to the invocation at `lane`.
We then use a version of `welford_combine` that handles the case where `count > 1` to merge them. Once the subgroup has reduced the value to the first lane, we broadcast it to all threads using `subgroupBroadcast`.

```markdown
Welford Vectorized
time: [111251.7095 ns 113935.1043 ns 117158.9148 ns]
thrpt: [33.3415 GiB/s 34.2849 GiB/s 35.1118 GiB/s]

All close - AVGE=0.0000014362298 MAE=0.000017166138 at [0, 1688, 413]
```

As we can see, our performance is now on par with the one-pass algorithm, but
no longer suffers from numerical instability!

## Conclusion

We've seen how we can arrive at a fast, numerically stable LayerNorm kernel by using
the proposed subgroup extension, which will be a welcome addition to the WebGPU ecosystem. There is still more to
explore here, for more information I recommend reading [this excellent post](https://archive.is/wip/rB7II) by OneFlow
on which this post is based.

Check out [the benchmarking repo](https://github.com/FL33TW00D/wgpu-bencher) for more, and thanks to [sekstini](https://x.com/fluffykittnmeow) for their insightful comments!

## References

[[1]](https://archive.is/wip/rB7II)
[[2]](https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance)


A pedagogical exploration of writing a GPU kernel to compute Layer Normalization as fast as possible.

Layer Normalization as fast as possible


With WebGPU landing on stable Chrome today, billions of internet users will have unimpeded access to [GPGPU compute](https://en.wikipedia.org/wiki/General-purpose_computing_on_graphics_processing_units) in the browser.
This is **very** exciting and paves the way for new possibilities in web based applications.
After assessing the feasibility in my previous [post](https://fleetwood.dev/posts/a-case-for-client-side-machine-learning), I set out to leverage WebGPU to run large machine learning models locally.
This effort has culminated in my _pre-pre-alpha_ library [`laserbeak`](https://github.com/FL33TW00D/laserbeak), which allows users to run a 250M parameter large(ish?) language model (LLM) on the web.

To showcase some of the capabilities, I've built a simple document editor on top of the excellent [`slate`](https://github.com/ianstormtaylor/slate) which you can see in the demo above. If you've already updated Chrome or are from the future, you can [try out the demo here](https://summize.fleetwood.dev) - and if you're interested to learn more, read on!

## Applications

![RunwayML Green Screen](/runwayml-gs.apng "RunwayML Green Screen")

This post by [RunwayML](https://runwayml.com/blog/building-a-web-based-real-time-video-editing-tool-with-machine-learning/) first got the gears turning about the applications of models in the browser. In the post, they break down their AI-powered [Green Screening tool](https://runwayml.com/green-screen/), which is built using a chain of models. They expound on the difficulty of setting up a streaming inference system, rapidly shuttling frames to many server side V100s to reach the level of interactivity required for good UX.

Meta AI took this a step further with their Segment Anything Model (SAM). SAM cleverly shifts almost all the model parameters to the encoder, leaving an extremely lightweight decoder which can be conditioned on user selected points to refine the mask (check out the [demo](https://segment-anything.com/demo)). Unfortunately, with the encoder on the server, you still need your GPU farm and an efficient streaming system. For a truly interactive experience, having the entire model locally will prove essential for real-time applications like audio and video.

With more and more applications requiring chains of models, it's no surprise that [`LangChain`](https://github.com/hwchase17/langchain) has exploded in popularity. `LangChain` enables developers to harness complex emergent behaviour by using LLMs as a primitive. Unfortunately, repeatedly calling into ~~the Oracle~~ GPT-4 is going to run up a huge bill and organisations that naïvely
build on top of OpenAIs APIs without proper optimization are going to end up out of business. I found this stat from Chip Huyen in her detailed [LLMOps
post](https://huyenchip.com/2023/04/11/llm-engineering.html) mindblowing:

> GPT-4 with 10k tokens in input and 200 tokens in output, it’ll be **$0.624 / prediction**.

This cost is untenable for a production service, even before we take multiple calls into consideration. Fortunately, not every call requires the mental horsepower
of GPT-4. Offloading some parts of the call chain to finetuned local models could dramatically reduce costs while offering additional benefits such as privacy and personalization.

## WebGPU

To effectively run these models and achieve the desired interactivity and distribution, a powerful cross-platform API is needed. Enter WebGPU - a _native_ graphics API landing in Chrome after six years of development. With its predecessor WebGL gone, compute shaders are now a first class citizen on the web, and the performance has proved **quite** impressive. From my [benchmarks](https://github.com/FL33TW00D/wgpu-mm), you can expect nearly 1TFLOP/s on an M1! [This post](https://pixelscommander.com/javascript/webgpu-computations-performance-in-comparison-to-webgl/) compares WebGPU against WebGL, demonstrating a huge 3.5x improvement for matrix muliplications.
With WebGPU also comes a brand new shading language, [`WGSL`](https://www.w3.org/TR/WGSL/), and despite [some critisism](https://github.com/gpuweb/gpuweb/issues/566), I've found the Rust inspired syntax very approachable. Fellow Rust programmers will enjoy WebGPU due to the excellent Rust support, with the [`wgpu`](https://github.com/gfx-rs/wgpu) crate implementing the API and [`naga`](https://github.com/gfx-rs/naga) handling shader compiliation (these two crates power the Firefox WebGPU implementation). If you're interested in delving deep into WebGPU, I highly recommend [this](https://surma.dev/things/webgpu/index.html) seminal blog post by Surma.

## Laserbeak

Just having access to GPGPU compute is a far cry from plug-and-play web based models. For this, I've created [`Laserbeak`](https://github.com/FL33TW00D/laserbeak) - the frontend for my WebGPU runtime, `Rumble`. Named after the smallest transformers, `Laserbeak` & `Rumble` empower you to integrate GPU accelerated transformers directly into your web or Electron applications. You can install `Laserbeak` [here](https://www.npmjs.com/package/laserbeak)!

`Laserbeak` is more than just a thin wrapper over the runtime; it offers some creative solutions to problems with running models in the browser. For example, we want to keep models as small as possible, both over the wire and in memory, so weight sharing for encoder-decoder models is a must. Through some hacking of the ONNX specification and [IndexedDB](https://developer.mozilla.org/en-US/docs/Web/API/IndexedDB_API/Using_IndexedDB), I've implemented weight sharing by storing tensors individually and reading them from the DB for each model. Storing each of the tensors individually enables some pretty interesting use cases that I will be exploring in the near future.

## The first model

If you're going to write a runtime, it's a good idea to have a model in mind
that you want to run. [`FLAN-T5`](https://huggingface.co/docs/transformers/model_doc/flan-t5) is an encoder-decoder model trained on text-to-text tasks and released by Google in Oct. 2022 in their [Scaling Instruction-Finetuned Language Models](https://arxiv.org/abs/2210.11416) paper. I chose FLAN as the inaugural model for my runtime for the following reasons:

-   It's available in a [wide range of sizes](https://s3.amazonaws.com/moonup/production/uploads/1668072995230-62441d1d9fdefb55a0b7d12c.png), with the 250M and 780M being the most attractive initially.
-   The 780M parameter model boasts impressive performance, achieving a 45.1 [MMLU score](https://paperswithcode.com/sota/multi-task-language-understanding-on-mmlu), which is comparable to LLaMA 13B (46.9).
-   It can be implemented in only ~32 operations (Explore the model with [this diagram](https://fleetwood.dev/flan_base.svg)).
-   RLHF tuning drastically improves performance and usability.
-   250M parameters seemed a feasible starting point even in FP32.

Although the model isn't finetuned for any particular task at the moment, its potential is promising, especially when considering the substantial performance increase demonstrated by the 780M parameter variant. Getting to this point will require some more work to reduce memory usage, but with a finetuned and larger model in the pipeline, we should be a strong competitor to OpenAI for narrow use cases.

## Implementation

The primary bottleneck in executing large models locally is not processing power, but memory capacity. For instance,
the smallest model from Meta's recent [`LLaMA`](https://ai.facebook.com/blog/large-language-model-llama-meta-ai/) release is 7B
parameters, which would require **14GB** of memory if stored in FP16. This is why projects like [`llama.cpp`](https://github.com/ggerganov/llama.cpp) have caused such a frenzy. Using techniques like [4-bit quantisation](https://arxiv.org/abs/2210.17323), the same 7B weights now take up only 3.5GB (at the cost of higher perplexity), which is much more feasible for consumer hardware. Unfortunately, the `GGML` family primarly targets the CPU and whilst this gives much more flexibility, it does have a major downside - it will eat [all of your cores](https://twitter.com/ggerganov/status/1640022482307502085). Since most modern devices are shipping with capable GPU/NPUs, we should capitalize on that silicon! The `GGML` community is very aware of this, and they've already moved [`whisper.cpp`](https://github.com/ggerganov/whisper.cpp) onto the Apple ANE.

Whilst the raw performance of WebGPU is pretty impressive, outperforming hyper optimised CPU implementations written by extremely talented engineers is still a daunting task. One of the most significant challenges stems from most modern models having dynamic input shapes. In the good old days, most models would have static input and output shapes. This makes running something like a ResNet on the GPU trivial - you just precompile all the shaders and execute them in sequence. Unfortunately, this is not the case for LLMs or diffusion models, where the input dimensions change at every time step.

To handle this, you need to just-in-time (JIT) compile the shaders, which makes beating CPU performance _slightly_ more challenging than anticipated. You need to compile all the shaders, move all the data to the GPU, run the model, and read the results back all in the same timespan as the CPU which... simply runs the model. However, all is not lost. Given that models are typically constructed of repeating blocks, you only need to JIT compile the shaders for the first block, then you can just repeatedly fetch them from the cache.

## Limitations

In its current state, my runtime can run a 250M parameter model at ~20 tokens/s without any optimization, except for caching. With further optimization, it is possible to significantly outpace CPU implementations. Moreover, the advantage over CPU will grow in tandem with the model size, making it an increasingly attractive option for running large models in the browser.

The main drawback of `Rumble` is that it currently only supports FP32(!), as this is what's supported via `naga` and `wgpu`, which makes our models _quite_ large. Luckily, the [`SHADER-F16`](https://gpuweb.github.io/gpuweb/wgsl/#extension-f16) extension is in the spec and will soon be available in Chrome, which should allow us to easily half our model size. We can take our size reduction even further, with clever tricks to implement INT8 and INT4 in the future. This has recently been done to great effect by the MLC-AI team for their project, [`web-llm`](https://github.com/mlc-ai/web-llm).

For models that can fit within the available memory, optimizing GPU utilization is essential for achieving high performance. In my [benchmarking](https://github.com/FL33TW00D/wgpu-mm), WebGPU is only able to reach ~40% of theoretical maximum FLOPs on native for matrix multiplies. This drops down to ~30% when running in the browser, due to the strict bounds checking implemented to avoid potential abuse. This is a shame, as it is quite easy to write a shader directly in [`Metal`](https://developer.apple.com/metal/) that reaches ~85% of maximum FLOPs, as demonstrated by Geohot in [`tinygrad`](https://github.com/geohot/tinygrad/blob/master/extra/gemm/metal_matmul.py). To do this, he leverages lots of intrinsics that aren't currently available through WebGPU, primarily SIMD groups. However, there have been several proposals (e.g [here](https://www.w3.org/2020/06/machine-learning-workshop/talks/simd_operations_in_webgpu_for_ml.html)) to introduce SIMD groups and I am confident it will eventually make its way into the browser.

## Conclusions

I am really excited to see what happens in this space as the Rust + WebGPU ecosystem matures. I anticipate an explosion of web apps and Chrome extensions leveraging sub-1B parameter models locally.

With more optimizations, operator support, and quantization tricks, `LaserBeak` &
`Rumble` should be able to run more useful models - **and soon**. Delivering this solo is
extremely ambitious, and taking it to the next level will not
be easy. If you find this interesting, check out [the demo](https://summize.fleetwood.dev) and if you want to get involved - [reach out to me](mailto:chris@fleetwood.dev).


Explore the implementation of an ML runtime to run large language models (LLMs) in the browser using Rust and WebGPU.

Running LLMs in the Browser with Rust + WebGPU


In order to perform classification on a functional brain scan it first
undergoes many preprocessing steps. One of these steps is the transformation from
the timeseries output of the fMRI scan, transforming a $m{\times}n$ (where $ m $ is
the number of timepoints recorded and $ n $ is the number of brain regions used)
to an $n{\times}n$ matrix of similiarity values (called a connectome). This
similarity value is a measure of the neural synchronization between the 2 regions.

So how do we quantify the similarity of 2 different timeseries? This blog post will
explore the common ways of quantifying time series similarity in a neuroscientific
setting. Before we get into the actual methods used to calculate time series similarity,
we need to cover the corner stone of almost all of the methods we are about to
explore - covariance.

## Covariance

**Covariance** is simply a measure of how two random variables change together.
Below is the formula for calculating covariance:

$$
\operatorname{\Sigma} = \operatorname{E}\left[\left(X-\operatorname{E}[X]\right ) \left (X-\operatorname{E}[X]\right)^\mathrm{T}\right]
$$

In the case of fMRI, we have a multivariate random variable, allowing us to use
Maximum Likelihood Estimation to estimate the covariance matrix. Below is a toy
example of our estimated covariance matrix.

$$
\begin{bmatrix}
var(X) & cov(X,Y) & cov(X,Z) \\
cov(X,Y) & var(Y) & cov(X,Y) \\
cov(X,Z) & cov(X,Y) & var(Z)
\end{bmatrix}
$$

This estimated covariance matrix is called the empirical covariance matrix. In
practice we don't use the covariance matrix for our fMRI analyis.
This for a number of reasons:

-   Covariance is bound between $ -\infty $ and $ \infty $ making it less
    suitable for downstream classifiers.
-   Covariance coefficients are not standardized and cannot be used to quantify the
    strength of the relationship.
-   Symmetric Positive semi-definite (SPD) matricies do not naturally form a Euclidean space (this will be
    important later).
-   When the number of features is large relative to the number of observations,
    the sample/empirical covariance matrix has excessive estimation error.

Many of the following approaches will aim to address some or all of the above
drawbacks with empirical covariance.

To address the excessive estimation error, it is common to perform a
transformation to the covariance coefficients, known as **"shrinkage"**. In
their seminal paper, "Honey, I shrunk the Sample Covariance Matrix" \[[1](http://www.ledoit.net/honey.pdf)\], Ledoit \& Wolf
proposed using shrinkage to regularize the sample covariance matrix.
"Shrinkage" as the name implies, pulls the most extreme covariance coefficients
towards more central values. This not only resolves our excessive estimation
error, it can also make the matrix easily invertable by encouraging numerical
stability. (The interested reader should consult the scikit-learn docs
\[[2](https://scikit-learn.org/stable/auto_examples/covariance/plot_covariance_estimation.html)\] )

Now that we have a well conditioned covariance matrix, we can attempt to address
some of the other identified drawbacks.

## Canonical Approaches

**Pearson's correlation coefficient** (Pearson's $ R $) or simply correlation, is
the most commonly used method to quantify similarity between 2 fMRI timeseries.
Correlation is a linear metric computed from the covariance of the 2 timeseries.
Below is the mathematical formula to compute correlation for a pair of random
variables:

$$
{\displaystyle \rho _{X,Y}={\frac {\operatorname {cov} (X,Y)}{\sigma _{X}\sigma _{Y}}}}
$$

Correlation is widely used in neuroscience as it has long statistical history
and is bound between -1 and 1. However, correlation does have some disadvantages.
The below figure should demonstrate one clearly:

![Correlation Decay](/corr.gif "Correlation Decay")

Due to correlations linear nature, the same timeseries being slightly out of
phase causes a huge decrease in the correlation value. Additionally, correlation
provides no distinction between whether 2 regions are directly
connected or indirectly connected via another region. To account for this, we
can use partial correlation!

**Partial correlation** is a variant of PCC that attempts to address
distinguishing between direct and indirect connections. This is done by
computing correlation between regions after regressing all other timeseries.
Partial correlation is computed from the inverse of the covariance matrix
(this is where the shrinkage comes in handy), also known as the precision matrix.
Below is the mathematical formula for computing partial correlation for a pair
of random variables:

$$
\rho = -{\frac {p_{ij}}{\sqrt {p_{ii} p_{jj}}}}
$$

where $p$ is the respective precision matrix (obtained through inversion of our
well conditioned covariance matrix).

## Novel Approaches

**Tangent space parameterization** is a recently proposed solution to the problematic
nature of SPD matricies, in that they don't naturally form Euclidean spaces.
Varoquax et al. \[[3](https://arxiv.org/abs/1008.5070)\] proposed a solution to this
sticky problem. Whilst I won't profess to be an expert on differential
manifolds or Riemmenian geometry, here is how I understand it.

Any covariance matrix is a symmetric positive semi-definite matrix. In order to
preserve the geometry during our transformations, we need to
perform a projection into a space where our geometry is preserved.
As SPD matricies form a diffentiable manifold, we can use this fact and
Riemannian geometry to keep our geometric intutions in tact.

![Manifold](/manifold.jpg "Manifold")

<a href="https://juliahub.com/docs/PosDefManifold/vb3YZ/0.4.8/introToRiemannianGeometry/#Intro-to-Riemannian-Geometry-1" styles="text-align: center;">
Image source
</a>

The above image makes it clear that if we perform Euclidean operations like
subtraction in our manifold, the distances and therefore the geometry will break
down. By defining a homomorphic (ie structure preserving) tangent space, we can
approximate the distance on the manifold using Euclidean distance in the tangent
space!

But we aren't quite out of the woods yet, since we still need to define the
reference point $ G $ on the manifold. As all of our covariance matricies need
to be projected into the same tangent space, we need to choose an appropriate
reference point. To do this, Varoquax et al. compute the geometric mean of the
matricies to define a reference space.

In their paper, Dadi et al. \[[4](https://hal.inria.fr/hal-01824205v3)\] showed
that across many fMRI classification tasks, tangent space parameterization
should be preferred over both correlation and partial correlation.

### Dynamic Time Warping

**Dynamic Time Warping** (DTW) is not like the previously discussed methods, as it does
not rely on covariance matricies. DTW is one of the canonical algorithms used for measuring
similarity between timeseries. It is commonly used in scenarios where timeseries
may have a different length or be temporally offest. DTW is a non linear metric,
and calculates the optimal match between two given sequences with certain
restrictions.

DTW was first proposed for usage with fMRI timeseries by Meszlényi et al. \[[5](https://www.frontiersin.org/articles/10.3389/fnins.2017.00075/full)\].
They propose its value for fMRI timeseries analysis as it has been noted that
"dynamic switch of brain states" can cause non-stationary time lags. If there
are known or unknown mechanisms in the brain that cause regional time differences
in hemodynamic response - DTW could be a great way to account for this.

Below is a Rust implementation of windowed DTW.
The window refers to the maximal warping distance.

```rust
let m = s.len() + 1;
let n = t.len() + 1;
let mut dtw = Array::from_elem((m, n), f64::MAX);

dtw[[0, 0]] = 0.;

let max_window = i32::max(*window, i32::abs((n - m) as i32));
for si in 1..n {
    let lower_bound = i32::max(1, si as i32 - max_window);
    let upper_bound = i32::min(m as i32, si as i32 + max_window);
    for ti in lower_bound as usize..upper_bound as usize {
        let cost = distance_fn(&s[si - 1], &t[ti - 1]);
        dtw[[si, ti]] = cost + 
            f64::min(
                f64::min(dtw[[si - 1, ti]], dtw[[si, ti - 1]]), 
                dtw[[si - 1, ti - 1]]
            );
    }
}
```

Dynamic Time Warping does have some disadvantages. One particularly
troublesome one being the computational complexity - $ O(n^2) $. The paper from
Meszlényi et al. does come with a DTW implementation for fMRI data on GitHub,
however it is extremely cumbersome and difficult to get running (particularly
on HPC systems where you don't have root access.)

To address this I wrote RustDTW, a Python module backed by Rust code that works
natively within Python on numpy matricies. Check out an example of using DTW
for fMRI classification [here.](https://github.com/FL33TW00D/rustDTW/blob/master/examples/classification/ABIDE_classification.ipynb)
My basic testing shows that it should be nearly 10x faster than a mutlithreaded
Python implementation.

## Classification Performance

Whilst not meant to be an exhaustive exploration of the different connectivity
metrics explored above (and itself is not an exhaustive exploration of available
connectivity metrics, for more please see Pervaiz et al. \[[6](https://www.biorxiv.org/content/10.1101/741595v1.full.pdf)\]
I thought I would demonstrate a comparison between them using a large open source dataset.
You can check out the inner workings and run it for yourself on Colab [here](https://colab.research.google.com/github/FL33TW00D/rustDTW/blob/master/examples/classification/ABIDE_classification.ipynb)!

For those just interested in the results, the below graph summarizes the
performance of each of the similarity metrics.


![Accuracy](/accuracy.png "Accuracy")

The results are pretty interesting! We can see that both non-linear matching and
transforming the covariance matricies into a natural space can provide improvements
over standard correlation. Perhaps a future metric will incorporate
characteristics from both tangent space parameterization and DTW to improve
performance.


A key step in diagnosing pathology using functional magnetic resonance imaging (fMRI) data is computing the similarity of timeseries, as this is what determines if 2 brain regions are "connected" or not. This post explores different methods of computing timeseries similarity

Hi, I'm Chris.

Essays