
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 <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

What's the future of machine learning going to look like? How can we attempt to predict what's going to come next given the rate of innovation?
We can't make any predictions with certainty, but in this post I will examine
past trends and ongoing work in an attempt to make predictions about the future. 

The current paradigm of machine learning is "static" - users requests are shipped off to GPUs running in the cloud, inference is performed and results are returned. It seems unlikely that this paradigm will remain unchanged going forward. I believe the next paradigm of ML is "dynamic", with client side encoding of user specific information in model weights, models becoming more personalized, and your tertiary self adapting to you. 

But this isn't some Jetsons fantasy, this future will arrive sooner than you
think. Let's apply some first principles thinking in an attempt to understand
when and how this will take shape.

## The Past 
‎|‎
:-------------------------:|:-------------------------:
![Gates&Allen](/client-side/Gates.png) |  ![GE-400](/client-side/GE-400.png)

Here's Bill Gates & Paul Allen sitting at a [Teletype Model 33](https://en.wikipedia.org/wiki/Teletype_Model_33 "Teletype Model 33") ASR Terminal in 1968. This Model 33 would be connected to a General Electric computer designed for **time-sharing**, where requests from multiple users would be multiplexed in order to maximise compute utilization. Machine learning is much like this today, with a huge warehouse of A100s located in ~~North Virginia~~ **undisclosed location** that your device shares with 100,000s of others via API requests.

Much like the time-sharing of the 60's and 70's, this approach is destined to change as more compute is put into the hands of the users. Given todays computing capabilties, is it feasible to do machine learning on the client side?

## First Principles Thinking 

Let's do some back of the envelope calculations to evaluate if client side machine learning is currently feasible.
There are 3 factors to consider:

- Memory - can we store the model parameters in (V)RAM on consumer hardware?
- Network - what size model can we reasonably ship to a user?
- Compute - can we perform inference (and possibly training) in a
   reasonable time?

Memory is not the bottleneck, since integrated GPUs share their memory with the
CPU, and even mobile devices are commonly shipping with 4GB or more of RAM. 4GB equates to
2B FP16 parameters, or large enough to comfortably fit
[InstructGPT](https://openai.com/blog/instruction-following/) with enough to spare
for an OS.

Network is more constrained, with the average internet speed in the
United States currently sitting at [~150Mbit/s](https://en.wikipedia.org/wiki/List_of_countries_by_Internet_connection_speeds).
This means that it would currently take ~140 seconds to download InstructGPT's 2.6GB FP16 weights.
The download however, only needs to be done once, and can be cached client side.

Now onto compute, the focus of this essay. I'll be using an Apple M1 as our baseline (engineers bias). Disregarding the locked down [Apple Neural Engine](https://github.com/hollance/neural-engine) (ANE) for now, the M1 has an 8 core GPU with 128 execution units, capable of [~2.6 TFLOP/s](https://web.archive.org/web/20201110184757/https://www.apple.com/mac/m1/) of FP32 compute. 

Unfortunately, the M1 GPU has 1:1 FP16/FP32 ratio [[1]](https://www.realworldtech.com/forum/?threadid=197759&curpostid=197759), so we can't just double our
FP32 FLOPs when we convert our models to FP16. But FP16 gives us a number of
other advantages, like reducing memory usage and allowing easy
offload to the ANE in future, so we will use this going forward. 

We can say that a running a model client side is feasible when the following holds:
$$
	 \text{Client Side Inference} < \text{RTT} + \text{GPU Inference}
$$

The median round trip time to all AWS servers from London is ~150ms (Check yours [here](https://cloudping.info)). And yes, whilst this may be high for an E-corp with deployments in every AWS availability zone, or you who lives in **undisclosed location**, it's a good figure to start with. The A100 supposedly does **312 TFLOP/s** of FP16 compute, or 120x faster than the M1 GPU. As the A100 is so fast, we will ignore the GPU inference time for now. Let's evaluate how much work we could do on the client side before the first packet even arrives at the data center. 

From these figures, we can calculate that:
$$
(2.6\text{TFLOP/s} * 150\text{ms}) = 390\text{GFLOPs}
$$

Keep this figure in mind as we explore where 390GFLOPs would be suitable given
current model requirements. 

## Software 2.0

Karpathy talks alot about the [Software 2.0](https://karpathy.medium.com/software-2-0-a64152b37c35) paradigm, where at Tesla, they replaced huge chunks of painstakingly written code with models who've learned the same functions. We can view applied machine learning from a functional programming perspective - a series of function applications, either defined by humans or learned by networks, to create incredibly complex behaviour. [RunwayML](https://runwayml.com/) talks about the success of **model chaining** for their AI powered video editor in this [post](https://runwayml.com/blog/building-a-web-based-real-time-video-editing-tool-with-machine-learning/).

Not all of the components in a model chain will have huge compute requirements. [Stable Diffusion](https://github.com/CompVis/stable-diffusion) is an example of model chaining with varying compute needs, composed of 3 networks for different purposes. You can see the structure in the following diagram.

![Stable Diffusion Architecture](/client-side/SD_arch.png)

The main part of the network is the [UNet](https://arxiv.org/abs/1505.04597), which consists of 860M parameters, or ~6.9GB
of FP16 weights. This would have pretty huge memory, network & compute requirements, so let's
focus on the lightest part of the network, the text encoder. Would it be possible to run the text encoder on the client side, and just ship
the output vector to the cloud for inference?

To find out we need to profile it. Unlike my own toy profiler [Steelix](https://github.com/FL33TW00D/steelix), 
which only supports a limited set of operations, [Tract](https://github.com/sonos/tract) has no problem profiling the text encoder, telling us it takes around 
$6.7\mathrm{x}{10}^9$ Fused Multiply Adds (FMA) or ~13.4GFLOPs for one
forward pass (N.B We will use 1 FMA == 2 FLOPs [[2]](https://hanlab.mit.edu/files/course/slides/MIT-TinyML-Lec02-Basics-of-Neural-Networks.pdf), but this is the source of some debate).

Tract also tells us the text encoder uses $8.5\mathrm{x}{10}^7$ parameters. This means that the FP16 parameters take up ~170MB of space. No problem for typical client side RAM, but still a sizeable file to ship over the network for a website. However, with techniques like compression, quantisation, distilation and pruning, we could certainly get the model size down to a managable download. 

Given how many GFLOPs are required to do an inference pass, and how fast the M1
GPU is, we can calculate that the M1 can perform **29 inference passes** in our 390GFLOPs budget. Before you protest, yes, since you have to go to the server
anyway for the UNet and decoder, why not run the whole model there? This may
result in the truly fastest experience, but from a pragmatic perspective, I am sure
Stability would like to offload some % of their FLOPs onto the client with
negligble impact on user experience.

## Future Capabilities 

To get a good look at how client side compute capabilities will grow, let's
analyze the Apple [A-series](https://en.wikipedia.org/wiki/Apple_silicon#A_series) of mobile chips as an example (given their longer 
history than the M-series).

![Apple Series A GPUs](/client-side/AChips22.png)

As we can see by the exponential growth, raw GPU compute is going to grow enormously
on the client side ([TSMC isn't out of ideas yet](https://www.anandtech.com/show/17356/tsmc-roadmap-update-n3e-in-2024-n2-in-2026-major-changes-incoming)). This doesn't even begin to explore the custom silicon being
shipped with ML in mind such as the aforementioned ANE. 

Even if we *just* had the above exponential improvement in raw compute, we would
certainly see client side machine learning capabilties grow. But that's not all.
Improvements at the silicon level will only be complemented by improvements at
the software level. Techniques such as [Pruning](https://hanlab.mit.edu/files/course/slides/MIT-TinyML-Lec03-Pruning-I.pdf), [Distillation](https://hanlab.mit.edu/files/course/slides/MIT-TinyML-Lec10-Knowledge-Distillation.pdf) and [RLHF](https://openai.com/blog/deep-reinforcement-learning-from-human-preferences/) all enable
impressive capabilities in a small package. InstructGPT outputs were
preferred by assessors whilst only having ~1/130 of the
parameters of GPT-3!

The combination of these 2 factors means that we will have incredible ML
capabilities within the next 24 months without a single packet ever needing to
be sent to the cloud.

## The next paradigm of machine learning

Now that we know that within the next few generations of hardware, machine
learning capability on the client side will *explode*, what do these new
capabilties enable?

### Dynamic ML

When we have the capability of training networks client side, we will enter the
era of **dynamic** machine learning. Models won't just be the same for everyone, 
and the weights will adapt dynamically as you use them.

A backwards pass (training) of a neural network takes approximately 2x the FLOPs of a
foward pass (inference) [[3]](https://www.lesswrong.com/posts/fnjKpBoWJXcSDwhZk/what-s-the-backward-forward-flop-ratio-for-neural-networks). So training larger models client side will arrive later than inference. But this isn't to say the idea is far fetched - the EfficientML lab [has already proved](https://tinytraining.mit.edu/) it's possible to train networks on extremely resource constrained systems, such as a microcontroller with **256Kb** of
memory.

Lucky for us, we don't need to train an entire model client side, only fine tune
a subset of the weights to get the dynamic behaviour we are looking for. This all sounds great, until some blob of weights on a server somewhere encodes every aspect of your life. It would be optimal if the subset of weights that get updated to learn about the user **remained on their device**. 

It turns out using techniques like [Adapters](https://github.com/Adapter-Hub/adapter-transformers) and [LORA](https://arxiv.org/pdf/2106.09685.pdf), you can simply pluck out different chunks of your Transformer, fine tune it on some new data, and plug it back in. Whilst those papers use it for finetuning a model for a new task, it should be feasible to fine tune a blob of weights locally using a user specific dataset. I expect it to be common place for users to be communicating with the data centre purely via seemingly meaningless vectors in the future.

## How?
So, how can we deliver on the futuristic vision of the ["The Intelligent Internet"](https://twitter.com/EMostaque/status/1562762975500374017)? 

It should be obvious that the web browser is the best software distribution
platform ever created; enabling machine learning in the browser is going to
be key. There have been rumblings in this direction for years, like [TensorFlow.JS](https://github.com/tensorflow/tfjs),
[ONNX Runtime Web](https://github.com/microsoft/onnxruntime/tree/main/js/web) and others, but none of them have ever made a huge impact.

Enter [WebGPU](https://www.w3.org/TR/webgpu/). Taking up the *torch* from WebGL after a decade, WebGPU promises near native performance on the GPU. Unlike WebGL, WebGPU is designed with compute workloads in mind. This isn't some nuclear fusion style promise. WebGPU **is** coming, with the Chrome origin trial ending in January 2023.

W3 is taking it even further, with a dedicated [Web Neural Network API](https://www.w3.org/TR/webnn) arriving... *eventually*. But I'm not interested in waiting that long. To me, it looks possible to build yourself once you have access to the GPU - [WONNX](https://github.com/webonnx/wonnx) is a great attempt at this.

I think that building a consise ML framework designed from the ground up for WebGPU,
capable of both inference and training, would enable a world of new possibilities. 
Rust seems well placed to do this, given that the whole WebGPU standard is [implemented in Rust](https://github.com/gfx-rs/wgpu) and powering FireFox already. Unfortunately, the first incarnation of WebGPU won't support INT8, so FP16 will have to do for now.

If you find this interesting - [reach out to me](mailto:chris@fleetwood.dev).


Is client-side machine learning feasible or desirable with todays compute?

A case for client-side machine learning


I've recently been reexploring the world of Rust & Machine Learning and what
the future of it is going to look like. Whilst doing so, I built a simple ONNX
inference framework in Rust. [ONNX](https://onnx.ai/) is an open standard for serializing
machine learning models, with seemingly strong industry backing. I
declared my framework finished when it could perform a forward pass of an [EfficientNet Lite](https://github.com/tensorflow/tpu/tree/master/models/official/efficientnet/lite). 

You can see the demo here: 
![Infer](/infer_demo.gif "Infer")

To achieve the above, I had to implement all of the necessary operations such as 
Convolutions, Transposes, Pooling etc. The most frequent cause of bugs while writing
the framework: **mismatched shapes**. 

It would have been great to know exactly what shape I needed at each edge in my network graph. This is where a model visualizer comes in.
The most popular visualizer is currently Lutz Roeder's [Netron](https://netron.app/).
However, when I opened the EfficientNet model in Netron - I got the following:

![Netron](/netron.png "Netron")

...no shapes after the input!

This post isn't to dunk on Netron - it's an amazing tool that has incredibly
wide format support and is very robust. However, seeing as I had already written a
parser for ONNX files, I decided to scratch my own itch and write a visualizer - with a few
additional features along the way - culminating in [Steelix](https://github.com/FL33TW00D/steelix). 
Steelix is ONNX only (greatly reducing the scope of this work unlike Netron) and
supports a very limited set of operators. However, despite these
limitations it can do some very useful things!

## Plotting

Steelix can produce a DOT file for any ONNX model, allowing you to explore 
the structure. You can see a demo below on our EfficientNet:

![Plotting](https://raw.githubusercontent.com/FL33TW00D/steelix/master/.github/images/steelix_plot.gif "Plotting")

Steelix also has the added advantage of being composable. You can embed a call to
Steelix in your Jupyter Notebook and automatically generate updated model graphs
for your documentation. If Steelix doesn't support an operator, you can still visualize your model, just
without any shapes as it is unable to do the mock inference pass.

## Summarization 

Often when playing with models, it would be good to know some summary statistics
about them. Steelix provides a handy summary table with some interesting metrics - check out the demo below on AlexNet:

![Summary](https://raw.githubusercontent.com/FL33TW00D/steelix/master/.github/images/steelix_summary.gif "Summary")

The main metrics of interest to me are the FLOP count, parameter count and hardware 
performance. Still to be added are peak
activiations and other useful metrics (you can learn more in these great
[slides](https://efficientml.ai/schedule/)). These metrics are particularly interesting if you want to run models on microcontrollers - expect to see me exploring that space soon.
Thanks to [Marius Hobbhahn](https://www.mariushobbhahn.com/) for some useful
insights on FLOP counting.

## Takeaways

Writing Steelix was a good learning experience, but there's a huge number of
improvements to be made. Currently, Steelix creates heaps of
unnecessary zeroed tensors as some operators
require runtime information (i.e the outputs are computed from the data of a
previous operation, not solely from metadata). It should be possible to identify
which subgraphs require runtime information and which don't, and only load
tensors for those subgraphs. Perhaps I will take this further in the future.




Steelix is a Rust CLI tool that allows you to analyze your ONNX models.

A Rust based ML model analyzer


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!

&nbsp;

**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" style="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.

&nbsp;

**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 fMRI data is computing the similarity of timeseries, as this is what determines if 2 functional regions are "connected" or not. How is this done? and where can we improve?

Hi, I'm Chris.

Posts