![\"TMSB4X\"/](\"/avcc/TMSB4X.png\")
Arc Institute recently unveiled the Virtual Cell Challenge. 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.\nFor 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.
\n\n\nGoal
\n
\nTrain 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?\nThis is the goal of the virtual cell challenge — a model (most likely a neural network) that can simulate exactly what happens\nto a cell when we change some parameter. Given that tightening your feedback loop is often the best way to speed up progress,\na model capable of doing this accurately would have significant impact.
\nTo 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 before continuing. This essay will build off of this to provide the ~minimum biology knowledge you'll need for the challenge.
\nThe training set consists of a sparse matrix and some associated metadata. More specifically, we have 220k cells, and\nfor each cell we have a transcriptome. This transcriptome is a sparse row vector, where each\nentry is the raw count of RNA molecules (transcripts) that the corresponding gene (our column) encodes for. Of the 220k cells,\n~38k are unperturbed, meaning no gene has been silenced using CRISPR. These control cells are crucial as we will see shortly.
\nTo 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\nperturbed cell.
\n
We can see that the cell with TMSB4X silenced has a greatly reduced number of transcripts compared with the control\ncells.
\nThe astute among you may be wondering why you don't just measure the count of the RNA molecules before and after\nsilencing 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.
\nThis inability to measure the cell state before and after introduces many issues, as we are forced to use a population of basal\n(a.k.a control, unperturbed) cells as a reference point. The control cells and perturbed cells are not entirely\nhomogeneous even prior to the perturbation. This means that we have to now separate out our true signal, the perturbation, from\nnoise induced by the heterogeneity.
\nMore formally, we can model observed gene expression in perturbed cells as:
\nX^p∼T^p(Dbasal)+H(Dbasal)+ε,ε∼Pε\nwhere:
\nPrior to the Virtual Cell Challenge, Arc released STATE, their own attempt to solve the challenge\nusing a pair of transformer based models. This serves as a strong baseline for participants to start with, so we will\nexplore it in detail.
\nSTATE 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.
\n![\"ST\"/](\"/avcc/ST.png\")
The State Transition Model is a relatively simple transformer with a Llama backbone that operates upon the following:
\nUsing a covariate matched set of control cells with paired target cells should assist the model in discerning the\nactual 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.\nIf working directly in gene expression space (i.e producing a full transcriptome), they pass the output through a learned\ndecoder.
\nST is trained using Maximum Mean Discrepancy. Put simply, the model learns to minimize the difference between the two probability distributions.
\n![\"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\ntake a little detour for some more biological grounding.
\n![\"Alternative](\"/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. 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.
\nWith this basic understanding, we can move on to how the SE model works. Remember, our core goal for SE is to create meaningful\ncell embeddings. To do this, we must first create meaningful gene embeddings.
\nTo produce a single gene embedding, we first obtain the amino acid sequence (e.g SDKPDMAEI... for TMSB4X) of all the different protein isoforms encoded for by the gene in question. We then feed these sequences to ESM2, 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.
\nNow 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:
\ng~j=SiLU(LayerNorm(gjWg+bg))\nWe've now obtained a gene embedding, but what we really want is a cell embedding. To do this, Arc represents each cell\nas the top 2048 genes ranked by log fold expression level.
\nWe then construct a "cell sentence" from our 2048 gene embeddings as follows:
\nc~(i)=[zcls,g~1(i),g~2(i),…,g~L(i),zds]∈R(L+2)×h\nWe add a [CLS] token and [DS] token to our sentence. The [CLS] token ends up being used as our "cell embedding" (very BERT-like)\nand the [DS] token is used to "disentangle dataset-specific effects". Although the genes are sorted by log fold\nexpression level, Arc further enforces the magnitude of each genes expression by incorporating the transcriptome in a\nfashion analogous to positional embeddings. Through an odd "soft binning" algorithm and 2 MLPs, they create some\n"expression encodings" which they then add to each gene embedding. This should modulate the magnitude of each gene\nembedding by how intensely it is expressed in the transcriptome.
\nTo train the model, they mask 1280 genes per cell, and the model is tasked with predicting them. The 1280 genes are\nselected such that they have a wide range of expression intensities. For the graphically inclined, the below\ndemonstrates the construction of the cell sentence.
\n![\"SE](\"/avcc/SE_path.png\")
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.
\n![\"Perturbation](\"/avcc/pert_disc.png\" "\\"Perturbation")
Perturbation Discrimination intends to evaluate how well your model can uncover relative differences between\nperturbations. To do this, we compute the Manhattan distances for all the measured perturbed transcriptomes in the test set (the ground\ntruth we are trying to predict, yt and all other perturbed transcriptomes, ypn) to our predicted transcriptome y^t. We then rank where the\nground truth lands with respect to all transcriptomes as follows:
\nrt=p=t∑1{d(y^t,yp)<d(y^t,yt)}\nAfter, we normalize by the total number of transcriptomes:
\nPDisct=Trt\nWhere 0 would be a perfect match. The overall score for your predictions is the mean of all PDisct. This is then normalized to:
\nPDiscNorm=1−2PDisc\nWe multiply by 2 as for a random prediction, ~half of the results would be closer and half would be further away.
\nDifferential 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. We do this for both our predicted perturbation distribution and the ground truth perturbation distribution.
\nNext, we apply the Benjamini-Hochberg 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 Gp,pred, and the ground truth set of differentially expressed genes Gp,true.
\nIf 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:
\nDEp=np,trueGp,pred∩Gp,true\nIf 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 G~p,pred), take the intersection with the ground truth set, and then divide by the true number.
\nDEp=np,trueG~p,pred∩Gp,true\nDo this for all predicted perturbations and take the mean to obtain the final score.
\nIf a virtual cell can accurately model the change in a cells state in response to perturbations, we can look forward to a\nvery interesting time in pharma. I hope this post accelerated your understanding of the challenge.
\nMay the best team win.
", "url": "https://fleetwood.dev/posts/virtual-cell-challenge", "title": "Arc Virtual Cell Challenge: A Primer", "summary": "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.", "date_modified": "2025-07-05T00:00:00.000Z", "author": { "name": "Christopher Fleetwood", "url": "https://fleetwood.dev" } }, { "id": "https://fleetwood.dev/posts/domain-specific-architectures", "content_html": "Billions of people may be continuously running AI inference for their waking hours in the near future. Satisfying this demand requires relentless focus on efficiency to reduce the required quantities of two key inputs: energy and capital. The constraints on these inputs in conjunction with the\nslowing and/or stagnation of both Moore's Law and Dennard Scaling has left hardware architects no choice but to pursue Domain Specific Architectures (DSAs) - architectures tailored to the task at hand.
\nThe current dominance of GPUs in modern deep learning is largely accidental - it was pure\nserendipity that the computational workload of graphics and deep learning were\nsimilar. Remnants of their graphical heritage still persist in GPU architectures today.\nWhat would AI inference hardware look like if it was redesigned carte blanche? By working backwards from the AI inference workload, we can determine some optimal properties these DSAs should have. Furthermore, we will attempt to\npredict the direction the inference paradigm will shift over time - a crucial exercise for\nhardware architects and engineers alike to ensure return on investment.
\n\n\n>90% of the total system energy is spent on memory in large ML models\n
\n
\n— Onur Mutlu
Often times it feels like many people have a view of computing\nsystems much like natural philosophers did prior to Copernicus — compute is like\nthe earth, at the centre of the universe. This is a bad model for understanding\nAI inference.
\nProgress in memory latency and throughput has lagged behind Moore's law since\nthe beginning 1.Reading and writing from memory is extraordinarily slow when compared to computation. Below is one of my favourite animations\nadapted from Vrushank Desai, demonstrating the\nlatency of different operations on an A100 GPU2.
\n\nLatency and throughput isn't the only issue with memory. The one true optimization objective for any accelerator is performance per dollar, and reducing power consumption is one key way of reducing the Total Cost of Ownership (TCO). The graph below compares the energy cost of different operations on different\ndatatypes for the Google TPU v4i 3.
\n![\"Energy](\"/npu/energy_consumption.png\")
Note the log y-axis, this should highlight how incredibly energy intensive reading and writing to DRAM\nis, regardless of the memory standard, when compared with actual computation.
\nAll of this together allows us to define one key property of any DSA\nfor AI inference to decrease power consumption and improve performance — minimize data movement.
\nWe will use this as an axiom for the remainder of the essay.
\nAs always, we must begin with the problem we are trying to solve, and work\nbackwards towards our solution. We will iteratively build up a more and more\ndetailed picture of inference in an attempt to derive optimal hardware\nprinciples.
\n\nThe above animation intends to give a very simple visual model of performing work\non a traditional accelerator (e.g. a GPU) with different bottlenecks - try clicking on the different bound\ntypes! This simple model affords us some intuition about performing a\nforward pass of a transformer. A simplified forward pass consists solely of moving extremely large, contiguous weight\ntensors from Global Memory (DRAM) to Shared Memory (L1 / SRAM), performing some number of operations on them\nin combination with our inputs, then storing the results back in DRAM.
\nWhen memory bound, we can see there are instances when the\nProcessing Elements (PEs) are doing nothing (a.k.a stalling). When compute bound, we see the PEs cannot\nkeep up with the data flowing through memory. As compute progress has outpaced\nmemory progress, keeping our compute units fed is one of the core optimization principles in deep learning.\nTherefore, optimizing the process of moving these weight tensors to and from memory should be the underlying principle behind many of\nour hardware design decisions.
\nThe first port of call when optimizing a piece of work is to simply try to avoid doing the work at all. One obvious way to do this is reducing the size of our weight tensors by using smaller data types.\nWe can cut our work in half, quarters or more by representing our weights with\nless bits. All modern accelerators support FP8 inference, pushing into\nFP4 territory in the latest generation. In order to maintain numerical accuracy with lower bit widths, hardware support\nfor high precision accumulation and online quantization is necessary. There is however another less obvious incentive for reducing the bit width of our data types.
\nThe number of full adders that are required for multiplying two floating point\nnumbers together is a quadratic function of the number of mantissa bits.\nFor two floating point numbers with M mantissa bits, the multiplication requires an\n(M+1)⋅(M+1) array of full adders. This translates to 576 full adders for FP32, 121 for FP16, 64 for BF16 and 16 for FP8(E4M3) 4.\nLess adders per multiplier means we can fit more compute per unit area, or trade the die area\nfor L1 cache.
\nThese two forces together are a powerful incentive to reduce the precision of our\nmodel. Unfortunately, there are limits — how many bits does a transformer really\nneed? Given that intelligence and compression are deeply intertwined5, we cannot\nexpect to continue to reduce the precision of our datatypes without compromising\nmodel performance. Scaling Laws for Precision6 explores this in depth, and it\nseems that the optimal region sits somewhere between 7 and 8 bits.
\n\n\nPrinciple 1: Hardware support for low precision data types.
\n
One simple way to reduce the opportunity cost of moving a tensor from Global\nMemory to Shared Memory is to introduce asynchronous data movement. By\nimplementing double buffering, the compute unit can work on one buffer while the other is being filled (skip ahead to the TPU animation below for a\ndemonstration). The overlap of computation, memory transfers, and inter-accelerator communication is a prerequisite for optimal hardware utilization. This has been the defacto standard since the H100,\nand we should expect any self respecting DSAs to be designed "asynchronous-first".
\n\n\nPrinciple 2: Design for asynchronous transfers from day 1.
\n
We know that moving tensors from Global Memory to Shared Memory is a frequent\noperation in our forward pass. Therefore, it would be prudent\nto create an optimized subsystem within our hardware to expedite the transfers.\nDirect Memory Access allows us to move our tensors between memory spaces (be it Global Memory to Shared Memory within one accelerator, or even to\nthe Shared Memory of another accelerator) without involving the\nhost (read: CPU) at all. Nvidia terms their subsystem the Tensor Memory Accelerator (TMA).
\nIt may sound ridiculous, but prior to the introduction of the TMA with the H100, Nvidia GPUs required all data to go via the\nregisters in the journey from Global Memory and Shared Memory, like forcing a river through a hosepipe. DMA that is aware of tensor layouts and can copy both within a node and between nodes (termed Remote Direct Memory Access or RDMA) should be a principle concern for any future DSA.
\n\n\nPrinciple 3: Dedicated hardware for tensor aware memory transfers.
\n
Typically when we move data within a computing system, it travels through a\nmulti-level cache hierarchy, with each level of the hierarchy trading off latency/throughput for size or the inverse. This is analogous to the human memory hierarchy (sensory memory, working memory and long-term memory).
\nThe number of levels in your hierarchy and their size is a function of the expected workload with respect to two dimensions:
\nIn multiprocessor architectures, privacy and data consistency are also key considerations in cache design. For example, the H100 features an 80MB shared L2 cache that is accessible by multiple Streaming Multiprocessors (SMs), requiring sophisticated synchronization and coherency checks to maintain data consistency, which is inherently slow.
\nUtilizing a memory hierarchy designed for one workload will be inefficient for\nanother. CPUs are designed for very unpredictable workloads, with random access\npatterns being the norm. GPUs on the other hand have threads frequently reusing\ndata and in larger contiguous chunks. These different access patterns lend\nthemselves to different designs, with CPUs benefiting from multiple smaller cache levels to expedite random accesses at different temporal and spatial localities. GPUs have more predictable and bulk-oriented access patterns, meaning they can effectively utilize fewer, larger cache levels. We should intuitively understand why CPUs often have a 3/4 tier cache hierarchy whereas GPUs typically have only 2 tiers.
\nLet's return to our simplified model and derive the optimal hierarchy for AI\ninference. We move weight tensors from DRAM into L1, using the whole weight tensor at ~the same time. Once we are done with that tensor, we will not use it again until the subsequent forward pass.\nTherefore, the idea of caching (storing data nearby for reuse) falls apart altogether. If we use a typical\ncache hierarchy for AI inference, we will be filling our L1/L2 with data that won't be reused any time soon.
\nWhat we really want is a high bandwidth, low latency storage space, much like L1/SRAM, but outside of the usual caching\nparadigm. This storage space can store chunks of our weight tensors whilst we\noperate on them. This is known as a scratchpad. A scratchpad negates all of the usual overheads that come with cache\nmanagement, and gives us the fast working memory we need! We have now arrived at our next principle: replace your cache hierarchy with an outsized scratchpad.
\nIn theory, if our scratchpad was large enough, we could store entire intermediate tensors between kernel invocations! This would completely remove roundtrips to\nGlobal Memory. Unfortunately, L1 is located on-chip and is implemented using SRAM which requires 6(!) transistors per bit of storage, meaning even small amounts\nof storage (read: kilobytes) takes an enormous amount of die area which could\notherwise be used for compute. This hasn't stopped many hardware startups like Cerebras, Groq and Tenstorrent from trading off some amount of compute for cache.
\n\n\nPrinciple 4: Replace your cache hierarchy with an outsized scratchpad.
\n
We will now focus on an algorithmic analysis of Transformer inference from a single accelerator perspective. This section is largely derived from Chapter 7 of How to scale your model7 and Transformer Inference Arithmetic8. I would recommend studying them both for a comprehensive understanding.
\nA modern Transformer consists of only 6 operations:
\nThe Attention and Feedforward blocks are the most relevant when discussing\nperformance. In an ideal world, the other operations can be fused into the aforementioned\noperations and thus incur ~little cost.
\nInference consists of two distinct phases with very different computational requirements:
\nEvery algorithm has a property known as the Arithmetic Intensity, which is the\nratio of operations performed to bytes accessed. The counterpart to this is\nthe ops:byte ratio of the hardware accelerator, which is the ratio of the\naccelerators "math" bandwidth and memory bandwidth. An algorithm is compute\nbound if the arithmetic intensity is higher than the accelerators ops:byte\nratio. We should endeavour to be compute bound so that our hardware is always doing useful work. Understanding the interplay between these ratios is key for performance.
Below is a table of key metrics for the most popular hardware accelerators as of\nFebruary 2025:
\n| Hardware | FP16 TFLOPs | DRAM | SRAM | ||||
|---|---|---|---|---|---|---|---|
| Size (GB) | BW (TB/s) | Ops:byte | Size (MB) | BW (TB/s) | Ops:byte | ||
| NVIDIA H100-SXM5 | 989 | 80 | 3.35 | 295.2 | 30.81 | 31 | 31.9 |
| NVIDIA A100 | 312 | 40 | 2.0 | 156.0 | 20.72 | ~19 | 16.4 |
| Tenstorrent Blackhole | 387 | 32 | 0.51 | 758.8 | 204.93 | 46.6 | 8.3 |
| Google TPU v4i | 138 | 8 | 0.61 | 226.2 | 144 | ~13.54 | 10.2 |
| 1 228KiB * 132 SMs, 2 192KB * 108 SMs, 3 1464KB * 140 Tensix Cores, 4Estimated |
As a mental reference point for these capacities, a [4096,4096] matrix (e.g. Llama2\nattention projections) in BF16 is 32MiB (~33.6MB). Determining if an operation is compute or memory bandwidth bound becomes quite complex, given the capacity constraints and throughput variation between the different memory spaces. Despite the difficulty, we will attempt to determine the typical boundedness of our two core operations in a transformer forward pass. This should inform the optimal balance between memory bandwidth and compute capacity.
\nAlmost the entirety of the FLOPs spent on our model will be in matrix\nmultiplications. Matrix multiplication has a very interesting property that bears repeating: the computational cost grows with the cube of the dimension O(n3), whereas the number of memory accesses grows with the square of the dimension O(n2). This means as the dimension grows the ratio of compute to data improves (from a hardware perspective).
\nIn a transformer, matmuls occur primarily in two places:
\nThe most common form of matmul we will perform in a forward pass is [B,Dmodel]⋅[Dmodel,F], where B is the batch size, Dmodel is our model dimension and F is the outer dimension of the weight tensor in question.\nAt what point does this operation become compute bound on a given accelerator?
\nWe are looking to determine when Tmath≥Tmem for the following:
\nTmath=Accelerator FLOPs/sTotal FLOPs=Accelerator FLOPs/s2BDF\nTmem=DRAM BandwidthTotal Bytes=DRAM Bandwidth2BD+2FD+2BF\nFor our loads, we multiply each term by a factor of two as our values are 2 bytes. In order for the matmul to become compute bound on a given accelerator, the arithmetic intensity must exceed the ops:byte ratio:
As demonstrated in How to scale your model7, we can simplify the denominator\nto 2DF, as both D and F are typically much larger than B. Using an Nvidia H100 as our reference, capable of 989 TFLOPs of BF16 compute with memory bandwidth of 3.35TB/s, we get the following:
\n2BD+2DF+BF2BDF≊2DF2BDF≥DRAM BandwidthAccelerator FLOPs/s=3.35×10129.89×1014⟹B≥295=Bcrit\nThis means our matmuls will only become compute bound when our B≥295. During\ninference, this will almost certainly occur during prefill (most system + user\nprompts exceed this (ignoring prompt caching)). However during decoding it will\nbe quite the challenge to achieve this batch size.
\nOne might think that quantizing our weights to FP8 would cut our Bcrit in\nhalf, but the H100 can also perform 2x the number of FLOPS in FP8, so we are\nback where we started. If we wanted to lower our Bcrit, what hardware changes could we make? The obvious one is\nto simply increase our memory bandwidth by stumping up for more advanced High Bandwidth\nMemory (HBM). However, this has diminishing returns with respect to our true optimization objection of performance /\ndollar. Another approach could be to trade off compute for SRAM as discussed\nearlier. In our previous table we can see that Bcrit=32 for the H100,\nmeaning we can accommodate algorithms with much lower intensity whilst still\nsaturating our compute.
\nHow to scale your model7 rigorously analyses attention to determine its Arithmetic Intensity. We will summarize it here, expand with recent advancements and what implications they have for hardware.
\nUsing FlashAttention9, the Arithmetic Intensity of a single head of (multi-headed) attention is as follows:
\nAttention Intensity=4BSD+4BTD4BSTD=S+TST\nWhere S is the (existing) sequence length, T is the number of tokens we are\ncurrently sampling, B is the batch size and D is our Dmodel. During prefill S=T, so the expression simplifies to T2/2T=T/2. This is quite favourable for hardware utilization! For our H100, this means that we are compute bound if T≳590 — which\nshould be very common during prefill (e.g. a standard system prompt). Decoding however, is a different story.
\nDuring decode T=1, and S≫T, which allows us to approximate the original expression as:
\nS+TST≈1\nArithmetic intensity this low means decode is fundamentally memory bound. This should make intuitive sense — we are loading a (potentially large) KV cache and performing a small number of operations due to the volume of input data ([1,Dmodel]). Without some architectural modifications (e.g. MQA, GQA, MLA) to reduce the size of the cache, there isn't much we can do to improve this situation.
\nThis makes decreasing the size of our KV cache one of our primary concerns,\nnot only for inference performance, but also for capacity. Earlier, our simplified model assumed only weight tensors are consuming DRAM capacity. However, the size of our cache quickly becomes considerable. Let's do some napkin math using Deepseek V310 671Bs hyperparameters to demonstrate, where nlayers=61 and dmodel=7168. V3 didn't use MHA, so we will use a plausible value of dhead=128.
\nIn the MHA case, we store both k,v∈Rdmodel for each layer using 2 bytes per value, meaning our cache size is as follows:
\nKV Cache Size=2⋅2⋅nlayers⋅nheads⋅dhead=2⋅2⋅61⋅56⋅128=1,749 KB/token\nGiven that a H100 has 80GB of DRAM, we will exceed the entire storage capacity once the\ncache hits ~46,000 tokens, just from the cache! This means naïve MHA fundamentally prohibits long context. We can see now why there\nhas been intense optimization pressure applied to reducing the size of the KV cache - it is key for both long context conversations and inference performance.\nThese optimizations included Multi Query Attention and Grouped Query Attention, but both methods trade model performance for cache size.
\nMore recently Deepseek introduced Multi Latent Attention11 (MLA), with the intention of compressing the KV cache by performing a low-rank decomposition. This works well as data and computation are two sides of the same coin, and here we trade data for compute in the form of our up and down projection matrices.
\nFor MLA, we simply store a single compressed latent dc for both\nk and v. Let's redo our math:
\nMLA KV Cache Size=2⋅nlayers⋅dc=2⋅61⋅512=62.5 KB/token\nThis means our KV cache is 28x smaller using MLA. Our context would need to be ~1.3M (4x the length of Ulysses) tokens before\nexceeding our H100 capacity. An order of magnitude difference has implications\nfor hardware. Using MHA, we could only store 17 tokens worth of cache in SRAM on the H100, whereas with MLA we can store 492 tokens worth of cache. This means we can\nstore entire pages of attention on-chip when using something like PagedAttention12. Much like reducing the precision of our data types, there is a lower limit here. However, I wouldn't be surprised to see inter-layer projections compressing the total cache size further.
\nDespite recent optimizations, we can clearly see that any single accelerator is going to be unable\nto accommodate a frontier model. Therefore, we must spill over from one accelerator to many.
\n\n\nPrinciple 5: For a single accelerator, turn the memory bandwidth up to 11.
\n
All models being served by the frontier labs far exceed the size of a single accelerator.\nThis necessitates sharding our model across multiple accelerators and\ncommunicating between them. Adding more chips creates a tradeoff: while it distributes both the compute and memory burden, it also increases communication overhead between chips, reducing the amount of computation available per chip to hide the communication cost. At a certain point, adding chips no longer increases performance.
\nGiven that we want to use computation to hide communication cost, the ratio\nbetween the number of operations we can perform to number of bytes we can send\nbetween accelerators is important. This is known as the ops:comms ratio.
Below is a table of the ops:comms ratio for the most popular accelerators\ntoday:
| Hardware | FP16 TFLOPs | Interconnect | Bidirectional Bandwidth (GB/s) | Ops:Comms Ratio |
|---|---|---|---|---|
| NVIDIA H100-SXM5 | 989 | NVLink 4.0 | 900 | 1099 |
| NVIDIA A100 | 312 | NVLink 3.0 | 600 | 520 |
| Tenstorrent Blackhole | 387 | Ethernet | 800 | 484 |
| Google TPU v5p | 459 | ICI | 600 | 765 |
Different sharding strategies incur different communication costs. Determining at which point an accelerator becomes communication bound for a\ngiven strategy is quite involved. Luckily for us, How to scale your model7 derives useful heuristics we can use to get an idea of when modern transformers may become\ncommunication bound, and use it to inform our hardware design.
\nWe will restrict our analysis to the decoding phase for brevity. During decode we can only use Model Parallelism for sharding our model. At what\npoint do we become bottlenecked by communication? To simplify the problem, we\nreduce our forward pass to a stack of MLPs. We consider 1D Model Parallelism for an MLP layer where we perform an AllGather on the input activations and a ReduceScatter on\nthe output activations as shown here. We want to understand at what point Tmath≥Tcomms, for our MLP performing\nthe following:
\nAin[B,Ds]⋅Win[D,Fs]→tmp[B,Fs]\ntmp[B,Fs]⋅Wout[Fs,D]→Aout[B,Ds]\nWhere B is our batch dimension, D is our model dimension, F is our FFN\ndimension and s denotes sharding.
\nGiven that we perform 2 matmuls, each performing 2BDF FLOPs, and we\ncommunicate 2 matrices of size [B,D], our Tmath and Tcomms are as follows:
\nTmath=Y⋅C4⋅B⋅D⋅FTcomms=W4⋅B⋅D\nWhere Y is the number of accelerators, C is the per accelerator FLOPs and W is our bandwidth. As we want to know when Tmath≥Tcomms, this simplifies down to:
\nF>Y⋅WC\nAstute readers will notice that WC is our ops:comms ratio! This formula is quite handy. Using F=18432 from Deepseek V3, we can see that we become communication bound on an H100 at ~16 GPUs. If we had a more\nfavourable ops:comms ratio, we could scale our tensor parallelism further. This analysis neglects latency, which can dominate communication time when the volume of data transferred is low, please refer to Transformer Inference Arithmetic8 for a more thorough\ntreatment.
We will now briefly cover Expert Parallelism (EP), as Mixture of Expert (MOE)\nmodels are the standard for frontier models today and it has implications\nfor hardware design. In Expert Parallelism, we dedicate an accelerator to hosting\none or more experts (a specialized FFN). When processing a batch of tokens, we\nmust send each of the tokens to a different subset of experts, depending on the\nresult of our routing mechanism. The result from our subset of activated experts\nthen must be aggregated afterwards. This warrants an AllToAll communication\n"primitive":
\nAllToAllX,JA[IX,J]→A[I,JX]\nAllToAll is also known as a resharding operation, as we have a batch\nof tokens sharded in the batch dimension (our batch entries are entirely\nindependent which is good for sharding), and we reshard them based on the result of our routing mechanism. You can imagine this as matrices flying all over to different accelerators, which necessitates a low latency, high throughput communication fabric.\nGiven that we are performing the same communication primitives repeatedly, we should explore specialized hardware subsystems to optimize their performance.
\nThe Deepseek V3 paper highlights the need\nfor hardware optimizations for communications. In lieu of vendor provided hardware optimizations, Deepseek "specialized" 20/132 SMs specifically for communications. This includes performing reduce operations for the MOE AllToAll aggregation, managing memory layout of tensors being transferred between experts etc. These SMs are designed to do computation, not network management, and by specializing them for communications, not matrix multiplication, 15% of all Tensor Cores are wasted.
\nWe should expect any performant domain specific architecture to have dedicated hardware for managing communications. It would be ideal if during data transfers the data could be transformed (e.g. arriving at the destination reduced, transposed etc). Remote Direct Memory Access (RDMA) being able to perform some primitive operations is suitably biomimetic — axons in the brain aren't just passive transmission channels, but perform "diverse functional operations" 13 beyond mere signal propagation.
\n\n\nPrinciple 6: Design for scale-out from day 1.
\n
\nPrinciple 7: Dedicated communication hardware should complement compute hardware.
By now we have analysed today's state of the art. Unfortunately, the paradigm does not stand still whilst we are\ndesigning hardware. Recent advancements in test time compute scaling have been called "the next\nscaling law", as demonstrated by the OpenAI O series. There are 2 dimensions\nin which we can scale inference compute: serially and parallel.
\nScaling inference compute serially manifests as extended multi step reasoning chains, analogous to Kahneman's System 2 thinking.\nThis is opposed to saying the first thing off the top of your head (System 1). In\ncontrast, scaling inference computation in parallel is like having n clones\napproach the same problem, with each potentially landing at different places in the solution\nspace.
\nScaling in the serial dimension (massively) increases demand for inference, in\nparticular the decoding phase. As we determined earlier, attention during decode is inherently\nmemory bound (MLA notwithstanding) without high batch sizes. This may contribute to something like "economies\nof scale" for providers that can hit high numbers of concurrent users.
\nScaling in parallel is more nuanced. As we stated earlier, we can only reach high utilization of our\nhardware during decode when our batch size is high (B≳295), which is not always easy to\nachieve even in a serving setting using Continuous Batching14. To reiterate: given that we must pay the fixed cost of moving the weights from\nmemory, low batch sizes have low Arithmetic Intensity and therefore do not\nutilize our hardware effectively. Now that scaling in parallel has been shown to\nincrease response quality, we can run our serving infrastructure and achieve\nhigher utilization by simply padding our batches with more parallel chains of thought. The\nusers get an improved experience that can be delivered for free.
\nI expect research into the optimal trade off between these 2 scaling dimensions to be a\nhot topic this year, resulting in "token optimal inference time scaling". This\nclosely mirrors the exploration-exploitation dilemma.
\n\n\nSummary
\nUsing multiple different approaches, we've derived a (non-exhaustive) list of design decisions\nthat should hold for any AI inference accelerator, namely:
\n\n
\n- Hardware support for low precision data types
\n- Design for asynchronous transfers from day 1
\n- Dedicated hardware for tensor aware memory transfers
\n- Replace your cache hierarchy with an outsized scratchpad
\n- For a single accelerator, turn the memory bandwidth up to 11
\n- Design for scale-out from day 1
\n- Dedicated communication hardware should complement compute hardware
\n
With an understanding of the inference workload, we can now explore some domain\nspecific architectures and determine if they adhere to the design principles we derived.
\nThe Google TPU is one of the most successful DSAs in\nthe history of computing. It would be derivative to explore the architecture in\ndetail, particularly as Google has been open about the design across the generations15^, ^16. In brief, the TPU adheres to all of the\nprinciples we have previously outlined, which is why it is currently the leading\naccelerator on a performance / dollar basis, our one true optimization objective.
\nInstead, we will focus on one of the most interesting components of the TPU - the Systolic array. A systolic array is a special purpose compute unit designed to accelerate certain regular computations, popularized by H.T Kung17, who rediscovered systolic arrays in 1982 (It’s important to be the last person to discover something). Its ability to accelerate matrix multiplications makes it a popular choice for AI inference accelerators.
\nThe TPUv2 onwards features one or more 128x128 weight stationary systolic arrays. They take their fitting name from the similarity to the human circulatory system. Data (blood) passes through many processing elements (cells) before returning to memory (heart).
\n\nThe above animation shows a 4x4 weight stationary systolic array, a scaled down\nversion of the unit in the TPU, performing a matrix multiplication (Y=WX). Weights are preloaded into the processing\nelements, and the output is accumulated from left to right.
\nThe principle benefit of these units is minimizing data movement to the theoretical limit. A naïve square matrix multiplication performs n3 memory reads, and has a space complexity of 3n2. This can be improved with tiling, using tiles of size M to Θ(bMn3), where b is the size of the cache line and M is the total size of the cache in bytes. A systolic array further improves this to the optimal number of memory reads of 2n2. In theory, each element of the input and weights is read only once, and is used many times as it flows through the array. In practice this is not the case, given that you have a fixed dimension for your systolic array and a dynamic workload consisting of many problem sizes. Even if the array is only partially utilised, the entire unit must do work, which can harm power efficiency. Furthermore, keeping the systolic array fed requires complex hardware and software infrastructure.
\nDo we see the theoretical efficiency gains from a systolic array play out in practice? The TPUv4 uses 1.3x-1.9x less power than its generational counterpart the A100. The systolic array is a staple for any hardware startup looking to build an AI accelerator, including OpenAI, who are reported to be taping out their new systolic array based accelerator this year.
\nTenstorrent is a startup founded in 2016 with the ambitious goal of creating an AI\naccelerator and supporting software stack from scratch capable of competing\nwith Nvidia. Their upcoming accelerator, Blackhole, offers an excellent case study in how the principles we've derived might be implemented in practice.
\n![\"Tensix](\"/npu/Tensix.png\")
Blackhole's layout should feel familiar to those who have seen that of a GPU, with a grid of independent cores. However, the resemblance is only surface level. We can see that the chip features families of specialized cores, namely D for\ndata movement, T for Tensix (compute), E for Ethernet, CPU cores for\nrunning Linux, and some additional cores for board management. This specialization embodies our principle of optimized hardware subsystems for their intended functions.
\nThe Tensix core, analogous to Nvidia's Streaming Multiprocessor, serves as the fundamental compute unit. Each Tensix core contains:
\nThe NoC connects the entire grid of cores together, allowing cores to copy data\nbetween each other and across the ethernet network to any other\naccelerator. This is the underlying thesis behind Tenstorrent: cost-effective\naccelerators with massive high bandwidth scale out.
\nSupporting this architecture is a novel software stack that represents a departure from traditional programming models. Unlike CUDA, where Shared Memory/L1 is undefined between kernel invocations, Tenstorrent's architecture allows intermediate results to remain in L1 and undergo successive transformations. This approach is almost systolic in nature: designed for repeated transformations before returning to memory — strongly adhering to our core principle of minimizing data movement.
\nThe most controversial design choice for Blackhole is the refusal to use\nHigh Bandwidth Memory. Tenstorrent has opted to use GDDR6 for Blackhole with a bandwidth of only 512GB/s — contrast this with the H100 with 3.35TB/s.\nThe CEO, Jim Keller, has stated that HBM is too expensive for the accelerators to hit commodity pricing. How accurate really is this?
\nAccording to Tae Kim of Barrons, the total bill of materials (BOM) for a H100\nwas $3320 in 2023. Combining this with the figure of $15/GB from SemiAnalysis for the HBM3 memory, the memory\ncontributes ~36% of the overall BOM of the H100. GDDR6 is far cheaper than\nHBM.
\nBy leveraging an outsized L1, asynchronous transfers, the NoC, and a more\nflexible programming model, Tenstorrent is hoping that they don't need HBM. The\nwinds of AI architecture do seem to be blowing in their favour. The introduction of MLA have shown that attention may not be forever memory bound.\nFurthermore, if they can crack low latency, high bandwidth interconnect, then\nthe burden of loading model parameters can be distributed across accelerators as\nfollows:
\nTmem=N⋅Abw2⋅P\nWhere P is our model parameters, N is our number of accelerators and\nAbw is our memory bandwidth.
\nTenstorrent still has huge challenges ahead of them, the largest of which being\nsoftware. The "CUDA moat" is usually the first topic of discussion when disparaging any hardware upstart. The moat as\nit's commonly understood is the maturity and (relative) robustness of the CUDA software stack. However, the real moat is the 1000's of engineers worldwide who have internalized the SIMT paradigm. Developing an entire software stack from zero, transparently scaling it out to 1000s of accelerators and forcing engineers to trade SIMT for MIMD may be a bridge too far.
\nFurthermore, Nvidia is now royally flush with cash, sitting with ~$38B of cash equivalents as of Q3 2024, and this industry is extremely capital intensive. This is primarily a problem for 2 reasons:
\nHowever, if Tenstorrent can pull off a solid software stack in the next ~24 months and AI\narchitecture trends continue to move in their favour, they'll be a strong\ncontender.
\nThe iteration time for designing and fabricating a new chip is typically 2-3 years. The AI paradigm has\nbeen moving much faster than this, with new architectures and datatypes being\ncreated on an increasing cadence. This is a huge problem for hardware architects, as you may end up with hardware optimized for\nthe models of yesteryear. If you're a hyperscaler spending $80B on hardware and the paradigm shifts dramatically before the 3 years\nrequired to recoup your costs3, well\nyou're out of luck! Whilst no one can concretely predict future architectural advances, is there any way we can gain some semblance of directionality of the paradigm in order to derisk our investments?
\nIf we believe that the search space of possible methods for constructing AGI is "large and\nsparse", then it would be prudent to take inspiration from our only existence proof: the\nhuman brain. In his remarkably prescient talk from 2010, Demis Hassabis\noutlines an iterative approach to developing the components of AGI via Systems Neuroscience.
\nSystems neuroscience does not mindlessly mimic the brain, but attempts to understand/extract the underlying algorithms and creatively and pragmatically reimplement them in software and hardware. He\nfurther explores this argument in a 2017 paper18, highlighting parallels between previous AI innovations and their counterparts in the brain.\nThese aren't coincidental similarities, but rather convergent evolution toward optimal solutions in the landscape of intelligence.
\nHassabis' thesis has continued to hold true. Frontier models are now all Mixture of Experts, which exhibit the same regional specialization we see in the brain. Looking forward through this lens, what other biological analogues could we see be implemented in-silico? I would hazard a guess at episodic memory systems and something similar to the hippocampal-neocortical consolidation system.
\nThe growth in demand for AI inference will not abate for the foreseeable future. Hardware architects attempting to hedge their bets when designing inference hardware in this market will almost certainly fail. Given that from specialisation comes efficiency, and flexibility and specialisation are diametrically opposed, one must be bold and bet on the way the winds of AI architecture will blow. Luckily, systems neuroscience can give us a lens through which to attempt to peer into the future.
\nThis principle of bold specialization extends beyond just startups. Every programmer is effectively an investor — their time and expertise are their capital. Being an early adopter of a technology that goes on to succeed can pay huge dividends, but this requires the same willingness to make bold, calculated bets on specific technological directions rather than hedging as a generalist.
\nThe principles we've derived: minimizing data movement, optimizing memory hierarchies, and designing for scale-out — provide a roadmap for those willing to make such bets in the inference hardware space. It remains to be seen if there exists a David to take down Goliath.
\nThanks to Erik Kaunismäki, Madeline Ephgrave, Luca Peric and Amine Dirhoussi for their insightful feedback on this post. Thanks to Felix LeClair and Martin Chang for their corrections.
\nD Patterson. (2004), "Latency lags bandwith", https://dl.acm.org/doi/10.1145/1022594.1022596 ↩
\nH Abdelkhalik et al. (2022), "Demystifying the Nvidia Ampere Architecture through Microbenchmarking and Instruction-level Analysis", https://arxiv.org/pdf/2208.11174 ↩
\nJouppi et al. (2021), "Ten Lessons From Three Generations Shaped Google’s TPUv4i", https://gwern.net/doc/ai/scaling/hardware/2021-jouppi.pdf ↩ ↩2
\nJ Dean. (2019), "The Deep Learning Revolution and Its Implications for Computer Architecture and Chip Design", https://arxiv.org/pdf/1911.05289 ↩
\nLegg, S et al. "Universal intelligence: A definition of machine intelligence.", https://arxiv.org/pdf/0712.3329 ↩
\nKumar et al. (2024), "Scaling Laws for Precision", https://arxiv.org/pdf/2411.04330 ↩
\nAustin et al. (2025), "How to Scale Your Model", https://jax-ml.github.io/scaling-book ↩ ↩2 ↩3 ↩4
\nChen, Carol. (2022), "Transformer Inference Arithmetic", https://kipp.ly/blog/transformer-inference-arithmetic/ ↩ ↩2
\nT Dao et al. (2022), "Flashattention: Fast and memory-efficient exact attention with io-awareness", https://arxiv.org/pdf/2205.14135 ↩
\nLiu, Aixin, et al. "Deepseek-v3 technical report.", https://arxiv.org/abs/2412.19437 ↩
\nLiu, Aixin, et al. "Deepseek-v2: A strong, economical, and efficient mixture-of-experts language model.", https://arxiv.org/pdf/2405.04434 ↩
\nW Kwon et al. (2023), "Efficient Memory Management for Large Language Model Serving with PagedAttention", https://arxiv.org/pdf/2309.06180 ↩
\nSasaki, T. (2013), "The axon as a unique computational unit in neurons", https://pubmed.ncbi.nlm.nih.gov/23298528/ ↩
\nYu, Gyeong-In, et al. "Orca: A distributed serving system for Transformer-Based generative models.", https://www.usenix.org/system/files/osdi22-yu.pdf ↩
\nJouppi et al. (2017), "In-Datacenter Performance Analysis of a Tensor Processing Unit", https://arxiv.org/abs/1704.04760 ↩
\nJouppi et al. (2021), "Ten Lessons From Three Generations Shaped Google’s TPUv4i", https://gwern.net/doc/ai/scaling/hardware/2021-jouppi.pdf ↩
\nH.T Kung. (1982), "Why Systolic Architectures", https://www.eecs.harvard.edu/~htk/publication/1982-kung-why-systolic-architecture.pdf ↩
\nHassabis et al. "Neuroscience-inspired artificial intelligence.", https://pubmed.ncbi.nlm.nih.gov/28728020/ ↩
\n\n\nGall's Law
\n
\nA complex system that works is invariably found to have evolved from a simple\nsystem that worked
\n— John Gall
This post walks you through the step-by-step discovery of state-of-the-art positional encoding in transformer models. We will achieve\nthis by iteratively improving our approach to encoding position, arriving at Rotary Postional Encoding (RoPE) used in the latest LLama 3.2 release and most modern transformers. This post intends to limit the mathematical knowledge required to follow along, but some basic linear algebra, trigonometry and understanding of self attention is expected.
\n\n\nYou shall know a word by the company it keeps
\n
\n— John Rupert Firth
As with all problems, it is best to first start with understanding exactly what we are trying to achieve. The self attention mechanism in transformers is utilized to understand relationships\nbetween tokens in a sequence. Self attention is a set operation, which\nmeans it is permutation equivariant. If we do not\nenrich self attention with positional information, many important relationships are\nincapable of being determined.
\nThis is best demonstrated by example.
\nConsider this sentence with the same word in different positions:
\nThe dog chased another dog\nIntuitively, "dog" refers to two different entities. Let's see what happens if we first tokenize them, map to the real token embeddings of Llama 3.2 1B and pass them through torch.nn.MultiheadAttention.
\nimport torch\nimport torch.nn as nn\nfrom transformers import AutoTokenizer, AutoModel\n\nmodel_id = "meta-llama/Llama-3.2-1B"\ntok = AutoTokenizer.from_pretrained(model_id)\nmodel = AutoModel.from_pretrained(model_id)\n\ntext = "The dog chased another dog"\ntokens = tok(text, return_tensors="pt")["input_ids"]\nembeddings = model.embed_tokens(tokens)\nhdim = embeddings.shape[-1]\n\nW_q = nn.Linear(hdim, hdim, bias=False)\nW_k = nn.Linear(hdim, hdim, bias=False)\nW_v = nn.Linear(hdim, hdim, bias=False)\nmha = nn.MultiheadAttention(embed_dim=hdim, num_heads=4, batch_first=True)\n\nwith torch.no_grad():\n for param in mha.parameters():\n nn.init.normal_(param, std=0.1) # Initialize weights to be non-negligible\n\noutput, _ = mha(W_q(embeddings), W_k(embeddings), W_v(embeddings))\n\ndog1_out = output[0, 2]\ndog2_out = output[0, 5]\nprint(f"Dog output identical?: {torch.allclose(dog1_out, dog2_out, atol=1e-6)}") #True\nAs we can see, without any positional information, the output of a (multi\nheaded) self attention operation is identical for the same token in\ndifferent positions, despite the tokens clearly representing distinct entities. Let's begin designing a method of enhancing self attention with positional information, such that it can determine relationships between words encoded by\ntheir positions.
\nTo understand and design an optimal encoding scheme, let's explore some desirable properties such a scheme should have.
\nLet's try and define some desirable properties that will make the optimization\nprocess as easy as possible.
\nEach position needs a unique encoding that remains consistent regardless of sequence length - a token at position 5 should have the same encoding whether the current sequence is of length 10 or 10,000.
\nThe relationship between positions should be mathematically simple. If we know the encoding for position p, it should be straightforward to compute the encoding for position p+k, making it easier for the model to learn positional patterns.
\nIf you think about how we represent numbers on a number line, it's easy to understand that 5 is 2 steps away from 3, or that 10 is 5 steps from 15. The same intuitive relationship should exist in our encodings.
\nTo increase our models' utility in the real world, they should generalize outside\ntheir training distribution. Therefore, our encoding scheme needs to be\nadaptable enough to handle unexpected input lengths, without\nviolating any of our other desirable properties.
\nIt would be ideal if our positional encodings could be drawn from a\ndeterministic process. This should allow the model to learn the mechanism\nbehind our encoding scheme efficiently.
\nWith multimodal models becoming the norm, it is crucial that our positional\nencoding scheme can naturally extend from 1D to nD. This will allow models to consume data like images or brain scans, which are 2D and 4D\nrespectively.
\nNow we know the ideal properties (henceforth referred to as Prn), let's start designing and iterating on our encoding scheme.
\nThe first approach that may jump to mind is simply to add the integer value of the token position to each component of the token embedding, with values ranging from 0→L where L is the\nlength of our current sequence.
\n\nIn the above animation, we create our positional encoding vector for the token chased from the index and add it to our token embedding. The embedding values here are a subset of the real values from Llama 3.2 1B. We can observe that they're clustered around 0. This\nis desirable to avoid vanishing or exploding gradients during training and therefore is something we'd like to maintain throughout the model.
\nIt's clear that our current naïve approach is going to cause problems. The magnitude of the position value\nvastly exceeds the actual values of our input. This means the signal-to-noise\nratio is very low, and it's hard for the model to separate the semantic\ninformation from the positional information.
\nWith this new knowledge, a natural follow on might be to normalize the position value by N1. This constrains the values between 0 and 1, but introduces another problem. If we choose N to be the length of the current sequence, then the position values will be completely different for each sequence of differing lengths, violating Pr1.
\nIs there a better way to ensure our numbers are between 0 and 1?\nIf we thought really hard about this for a while, we might come up with switching\nfrom decimal to binary numbers.
\nInstead of adding our (potentially normalized) integer position to each\ncomponent of the embedding, we could instead convert it into its binary\nrepresentation and s t r e t c h our value out to match our embedding dimension, as demonstrated below.
\n\nWe've converted the position of interest (252) into its binary representation\n(11111100) and added each bit to the corresponding component of the\ntoken embedding. The least significant bit (LSB) will cycle between 0 and 1 for every\nsubsequent token, whilst the most significant bit (MSB) will cycle every\n2n−1 tokens where n is the number of bits.\nYou can see the positional encoding vector for different indices in the animation below 1.
\n\nWe've solved the value range problem, and we now have unique encodings that are\nconsistent across different sequence lengths. What happens if we plot a low dimensional version of our token embedding and visualize the addition of our binary positional vector for different values.
\n\nWe can see that the result is very "jumpy" (as we might expect from the\ndiscrete nature of binary). The optimization process likes smooth, continuous and\npredictable changes. Do we know any functions with similar value ranges that are smooth and continuous?
\nIf we looked around a little, we might notice that both sin and cos fit the bill!
\nThe above animation visualizes our position embedding if each component is\nalternatively drawn from sin and cos with gradually increasing\nwavelengths. If you compare it with the previous animation, you'll notice a striking similarity!
\nWe've now arrived at Sinusoidal embeddings; originally defined in the Attention is all you need paper.\nLet's look at the equations:
\nPE(pos,2i)=sin(100002i/dpos)PE(pos,2i+1)=cos(100002i/dpos)\nwhere pos is the tokens position index, i is the component index\nin the positional encoding vector, and d is the model dimension. 10,000 is the base wavelength (henceforth referred to as\nθ), which we stretch or compress as a function of the component index. I encourage you to plug in some realistic values to get a feel for this\ngeometric progression.
\nThere's a few parts of this equation that are confusing at first glance. How did the\nauthors choose 10,000? Why are we using sin and cos for even and odd positions respectively?
\nIt seems that using 10,000 for the base wavelength was determined experimentally 2. Deciphering the usage of both sin and cos is more involved, but crucial\nfor our iterative approach to understanding. The key here is our desire for a linear relation between two encoded positions (Pr2). To understand how using sin and cos in tandem produce this linear relation, we will have to dive into some trigonometry.
\nConsider a sequence of sine and cosine pairs, each associated with a frequency ωi. Our goal is to find a linear transformation matrix M that can shift these sinusoidal functions by a fixed offset k:
\nM⋅[sin(ωip)cos(ωip)]=[sin(ωi(p+k))cos(ωi(p+k))]\nThe frequencies ωi follow a geometric progression that decreases with dimension index i, defined as:
\nωi=100002i/d1\nTo find this transformation matrix, we can express it as a general 2×2 matrix with unknown coefficients u1, v1, u2, and v2:
\n[u1u2v1v2]⋅[sin(ωip)cos(ωip)]=[sin(ωi(p+k))cos(ωi(p+k))]\nBy applying the trigonometric addition theorem to the right-hand side, we can expand this into:
\n[u1u2v1v2]⋅[sin(ωip)cos(ωip)]=[sin(ωip)cos(ωik)+cos(ωip)sin(ωik)cos(ωip)cos(ωik)−sin(ωip)sin(ωik)]\nThis expansion gives us a system of two equations by matching coefficients:
\nu1sin(ωip)+v1cos(ωip)u2sin(ωip)+v2cos(ωip)=cos(ωik)sin(ωip)+sin(ωik)cos(ωip)=−sin(ωik)sin(ωip)+cos(ωik)cos(ωip)\nBy comparing terms with sin(ωip) and cos(ωip) on both sides, we can solve for the unknown coefficients:
\nu1u2=cos(ωik)=−sin(ωik)v1v2=sin(ωik)=cos(ωik)\nThese solutions give us our final transformation matrix Mk:
\nMk=[cos(ωik)−sin(ωik)sin(ωik)cos(ωik)]\nIf you've done any game programming before, you might notice that the\nresult of our derivation is oddly familiar. That's right, it's the Rotation Matrix! 3.
\nSo the encoding scheme designed by Noam Shazeer in Attention is all you need was already encoding relative position as a rotation back in 2017! It took another 4 years to go from Sinusoidal Encoding to RoPE, despite rotations already being on the table...
\nWith the knowledge in hand that rotations are important here, let's\nreturn to our motivating example and try to discover some intuitions for our next iteration.
\n01234The dog chased another dog-2-1012The dog chased another dog\nAbove we can see the absolute positions of our tokens, and the relative\npositions from chased to every other token. With Sinusoidal Encoding, we\ngenerated a separate vector which represents the absolute position,\nand using some trigonometric trickery we were able to encode relative positions.
\nWhen we're trying to understand these sentences, does it matter that this word is the 2157th word in this blog post? Or do we care about its relationship to the words around it? The absolute position of a word rarely matters for meaning - what matters is how words relate to each other.
\nFrom this point on, it's key to consider positional encoding in the context of\nself attention. To reiterate, the self-attention mechanism enables the model to weigh the importance of different elements in an input sequence and dynamically adjust their influence on the output.
\nAttn(Q,K,V)=softmax(dkQKT)V\nIn all our previous iterations, we've generated a separate positional encoding\nvector and added it to our token embedding prior to our Q, K and V projections.\nBy adding the positional information directly to our token embedding, we are\npolluting the semantic information with the positional information. We should\nbe attempting to encode the information without modifying the norm. Shifting to multiplicative is the\nkey.
\nUsing the dictionary analogy, when looking up a word (query) in our dictionary (keys), nearby words should have more influence than distant ones. The influence of one token upon another is determined by the QKT dot product - so that's exactly where we should focus our positional encoding!
\na⋅b=∣a∣∣b∣cosθ\nThe geometric interpretation of the dot product shown above gives us a magnificent insight.\nWe can modulate the result of our dot product of two vectors purely by\nincreasing or decreasing the angle between them. Furthermore, by rotating the\nvector, we have absolutely zero impact on the norm of the vector, which encodes\nthe semantic information of our token.
\nSo now we know where to focus our attention, and have seen from another angle why a\nrotation might be a sensible "channel" in which to encode our positional\ninformation, let's put it all together!
\nRotary Postional Embedding or RoPE was defined in the\nRoFormer paper (Jianlin Su designed it independently on his blog here and here).\nWhile it may seem like voodoo if you skip to the end result, by thinking about Sinusoidal Encoding in the\ncontext of self attention (and more specifically dot products), we can see how\nit all comes together.
\nMuch like in Sinusoidal Encoding, we decompose our vectors (q or k, instead of pre-projection x) into 2D pairs/chunks. Rather than encoding absolute position directly by adding a vector we drew from sinusoidal functions of slowly decreasing frequencies, we cut to the chase and encode relative position by multiplying each pair with the rotation matrix.
\nLet q or k be our input vector at position p. We create a block diagonal matrix\nwhere Mi is the corresponding rotation matrix for that component\npairs desired rotation:
\nR(q,p)=M1M2⋱Md/2q1q2⋮qd\nMuch the same as Sinusoidal Encoding, Mi is simply:
\nMi=[cos(ωip)−sin(ωip)sin(ωip)cos(ωip)]\n\nIn practice, we don't use a matrix multiplication to compute RoPE as it would be\ncomputationally inefficient with such a sparse matrix. Instead, we can directly apply the rotations to pairs of elements independently, taking advantage of the regular pattern in the computation:
\nRΘ,pdq=q1q2q3q4⋮qd−1qd⊙cospθ1cospθ1cospθ2cospθ2⋮cospθd/2cospθd/2+−q2q1−q4q3⋮−qdqd−1⊙sinpθ1sinpθ1sinpθ2sinpθ2⋮sinpθd/2sinpθd/2\nThat's all there is to it! By artfully applying our rotations to 2D chunks of q and\nk prior to their dot product, and switching from additive to\nmultiplicative, we can gain a big performance boost in evaluations 4.
\nWe've explored the 1D case for RoPE and by this point I hope you've gained an\nintuitive understanding of an admittedly unintuitive component of transformers.\nFinally, let's explore extending it to higher dimensions, such as images.
\nA natural first intuition could be to directly use the [xy] coordinate pairs from the image. This might seem intuitive, after all, we were almost arbitrarily pairing up our components previously. However, this would be a mistake!
\nIn the 1D case, we encode the relative position m−n through a rotation of pairs\nof values from our input vector. For 2D data, we need to encode both horizontal and vertical relative positions, say m−n and i−j independently. RoPE's brilliance lies in how it handles multiple dimensions. Instead of trying to encode all positional information in a single rotation, we pair components within the same dimension and rotate those, otherwise we would be intermixing the x and y offset information. By handling each dimension independently, we maintain the natural structure of the space. This can generalize to as many dimensions as required!
\nIs RoPE the final incarnation of positional encoding? This recent paper from DeepMind deeply analyses RoPE and highlights some fundamental problems.
\nI anticipate some future breakthroughs, perhaps taking inspiration from\nsignal processing with ideas like wavelets or hierarchical implementations. As models\nare increasingly quantized for deployment, I'd also expect to see some\ninnovation in encoding schemes that remain robust under low-precision arithmetic.
\nPositional encoding has and continues to be treated as an after thought in\ntransformers. I believe we should view it differently - self attention has an\nAchilles heel that has been repeatedly patched.
\nI hope this blog post showed you that you too could have discovered state of the\nart positional encoding, despite it being unintuitive at first. In a follow up\npost I'd love to explore practical implementation details for RoPE in order to\nmaximise performance.
\nThanks to Madeline Ephgrave for proof reading this, and thanks to James\nMcDermott for his insightful corrections.
\nBinary and Sinusoidal animations are reproductions of animations contained\nin this video. ↩
\nUsing θ=10000 gives us ⌊2π⋅10000⌋ unique positions, or a\ntheoretical upper bound on the context length at ~63,000. ↩
\nPieces of this post are based on this fantastic\npost\nby Amirhossein Kazemnejad. ↩
\nFor empirical evidence, see this great post by EleutherAI. ↩
\nIn the novel Accelerando by Charles Stross, the protagonist Manfred Macx is a power user of smartglasses. These glasses dramatically enhance his ability to interact with the world. Over the course of the book, we learn about the glasses' incredible capabilities, from real-time language translation to dispatching agents to perform complex tasks across the internet. At one point, Manfred has the glasses stolen from him, and it becomes clear that they form a significant part of his cognition.
\nOffloading parts of your cognition to a device is already commonplace today. Many of us already rely on our smartphones for upcoming events, directions, reminders etc. If you take the idea of having a "silicon lobe" to its logical conclusion, you end up at... a computer being physically integrated into your brain aka Neuralink. This is true in the novel too, with the glasses being "jacked in" to the brain directly, allowing for a higher bandwidth interface and direct sensory manipulation.
\nIn the interim between the iPhone and Neuralink, I predict we will see some kind of Augmented Reality (AR) glasses hit the mainstream for consumers, similar to those in the book (minus the direct neural interface). This device may become ~as integral to your life as Manfreds, where losing them is equivalent to a lobotomy!
\nThis post is a deep dive into the following:
\nAR glasses are a form of wearable technology that superimposes digital information onto the real world. While all AR glasses are smartglasses, not all smartglasses are AR glasses.
\nBelow you can see the 2024 Ray-Ban + Meta smartglasses, one of the first forays into the consumer market. I'll be using them as a reference point throughout this post.
\n\nNotice how the arms are quite a bit thicker than standard sunglasses from Ray-Ban. These glasses pack in a huge amount of electronics to provide users with a rich "AI + Audio" experience. Using the integrated camera and spatial audio, you can ask questions, listen to music and use AI to perform other tasks. However, they're lacking the crucial component of AR glasses: a display. Meta isn't stopping there though, and are expected to announce Orion and Hypernova, their true AR glasses, at Meta Connect on September 25 - 26, 2024 1.
\nIn order to usurp the smartphone and become the primary interface with the digital world, AR glasses must provide a significant boost in utility. The use case that everyone is talking about today is AI integration, and for good reason. In my opinion, every human could get a staggering amount of utility from having insert SOTA LLM deeply integrated into their lives. These LLMs are already proving incredibly useful in chatbot form, however they've quickly become bottlenecked - not by IQ, but by context and bandwidth.
\nPerhaps the latest wave of AI innovation will be the catalyst for AR glasses to hit the mainstream. To see what the glasses might be capable of, let's analyse some use cases.
\nThe obvious use case for AR glasses is providing updates and real-time information in a timely and contextual manner. Microsoft has achieved a modicum of success in the enterprise AR market with their HoloLens 2, which is primarily targeting manufacturing and military applications. The glasses provide real-time, hands-free instructions to workers, allowing them to carry out routine tasks more efficiently. Despite being relatively absent from the public consciousness, the HoloLens line has sold over half a million units since its release 4.
\nWhilst we won't be able to directly interface with our auditory and visual systems for at least the next 5 years, that doesn't mean to say that AR glasses can't give users a new found control over their senses. One of the most liberating things about the proliferation of high quality Active Noise Cancelling (ANC) headphones is the ability for an individual to filter their audio input. The market has responded well to this ability, and in 2020, Apple's revenue from AirPods sales was larger than the entirety of Twitter, Spotify and Square combined 5.
\nWhat would this look like for visual data? Perhaps something reminiscent of "Times Square with Adblock on".
\n![\"Times](\"/ar-glasses/ts-adblock.gif\")
The ability to modify your visual input is a powerful one, and it could be used for more than just blocking ads. For example, during a conversation with a person speaking a different language, you could use a diffusion model to modify their lip movements to match your language. This, plus ANC and a speech-to-speech model, could be an excellent start to the dissolution of language barriers. In order to achieve both of these use cases, the AR would need to be extremely high-resolution and low-latency, which is a significant challenge we'll discuss later.
\n\n\nA person's name is to that person, the sweetest, most important sound in any language.
\n— Dale Carnegie
\n
Whilst this is perhaps a subset of "Contextual Overlay", it's worth highlighting, as it is a compelling use case far beyond what smartphones can provide today. By providing pieces of data throughout the course of a day, everyone could become more charismatic and thoughtful. You'll never have to forget a name or a face ever again.
\nYou can take this further too. According to Carnegie 6, Abraham Lincoln would dedicate considerable time to learning about the interests and backgrounds of people he was scheduled to meet. This allowed him to engage them in conversations about topics they were passionate about, greatly boosting his likeability. With AR glasses, everyone could do this effortlessly (much to the chagrin of those of us who have expended considerable effort practicing this).\nThis is on top of the more obvious personal assistant utilities! However, for this to be totally effective, the glasses would need to serialize everything you ever see or hear, which is quite the social and regulatory challenge.
\nThe outlined use cases above are just the tip of the iceberg. In Accelerando, the glasses also perform a variety of futuristic tasks, such as:
\nEven without these capabilities, by simply combining contextual overlay, sensory control and personal assistant functions, we can already see the immense value that AR glasses could provide to users within the next ~3 years. Let's dive into the technical challenges that must be overcome to make this vision a reality.
\nIn order to better understand the technical challenges, let's first establish some axioms based on the physical constraints of the glasses.
\nThere is surprisingly little literature on the acceptable weight of consumer glasses, so let's look at some current products on the market.
\n![\"AR](\"/ar-glasses/weight_plot.png\")
Anecdotally, it seems like the Meta smartglasses are light enough for all day wear, with the XREAL Air 2 Ultra not being too far off. Given the potential utility of the glasses, I think using 75g as an upper bound is reasonable. This also aligns with Meta's rumoured Hypernova glasses we mentioned earlier, which reportedly sit at ~70g 7.
\nAnother concrete axiom we can build off of is volume. The shape of the human head is going to stay consistent for the foreseeable future, and whilst there is some leeway to be found in the flexibilty of social dynamics, we can expect the form factor (and therefore volume) of the glasses to be approximately similar to glasses today. To calculate a baseline volume, I went to the closest demo location and measured the Meta smartglasses with calipers.
\nFrom my measurements, the volume of each arm is approximately 10cm3, so 20cm3 for both. If we naïvely used the entire volume of a single arm as a battery (and allocate the rest to electronics), and assuming we use a SOTA battery with 800Wh/L at a voltage of 3.7V, we can calculate the maximum possible battery life of any smartglasses as follows:
\n800Wh/L×0.01L=8Wh8Wh÷3.7V=2162mAh\n2162mAh is approximately half of the battery life of a modern smartphone, but is an order of magnitude larger than the battery that currently ships in the Meta smartglasses highlighted below (154mAh). To put it simply, if we want 8 hours of battery life, the glasses are limited to 3600J/h (1Wh==3600J).
\n![\"Glasses](\"/ar-glasses/mrb-battery.png\")
Maximising battery life is absolutely critical for AR glasses, and it's part of what makes them such an enticing challenge. It requires a full stack solution, from efficient software to efficient hardware to efficient power management. From the above teardown, we can see that Meta hasn't exactly maximised volume usage, so I expect that later generations of glasses could have significantly longer battery life. Perhaps like in aeroplanes and electric cars, we will see power sources being used as structural components in the glasses to increase volume utilisation.
\nOn top of the physical constraints, we have to deal with the following technical challenges:
\nWe won't explore all of these challenges in this post, but progress is being made everywhere! Let's take a look at the most significant challenge: near eye optics.
\n\n\nSuccess of smartglasses in a consumer acceptable form factor begins & ends with near eye optics.
\n— Christopher Grayson, 2016
\n
The above quote seems to be the underlying truth behind the lack of true consumer AR glasses. Therefore, we will devote significant time and energy to understanding the two key components of near eye optical systems and the advances required.
\n![\"VR](\"/ar-glasses/compare-vr-ar.png\")
The above diagram from Radiant Vision Systems 8 demonstrates why AR optics are much more challenging than VR. On the left hand side, you can see the simplicity of a VR optical system - the light from a display is fed directly into the eye with lenses for magnification. On the right hand side, you can see the AR optical system, with a compact display offset from the eye. This requires a complex optical combiner to overlay the digital image onto the real world.
\nKarl Guttag (A legend in the AR space), frequently comments that AR is many orders of magnitude more difficult than VR. The two major bottlenecks in AR optics that require significant advances are the underlying display technology for image generation, and the optical combiner for merging real-world and digital imagery into a single view.
\nTo start to get a sense of the challenges, let's break down the requirements for the display:
\nThese stringent requirements are compounded by the fact the optical combiner will inherently degrade the light from both the real world and the display!
\nHow do we actually combine the light from the two sources? Below you can see a map from IDTechEx of all optical combiner methods:
\n![\"Optical](\"/ar-glasses/optical-combiners.png\")
As you can see, optical combiners are divided into waveguides and non-waveguides. I'm not going to discuss non-waveguides in this post, as from my research it seems like waveguides are going to be the technology of choice for the first generation of consumer AR glasses.
\nThe function of a waveguide is fundamentally quite simple - guide the light from the display into the viewer's eye. Waveguides are usually defined by their underlying mechanism: diffraction or reflection. Whilst diffractive waveguides have been the dominant technology up until now, I believe reflective waveguides are the future of AR glasses due to their superior efficiency and image quality. For more information on waveguides, I highly recommend reading this post by Optofidelity 10.
\n![\"Lumus](\"/ar-glasses/lumus.png\")
\nAbove you can see a diagram of a reflective waveguide from Lumus Vision. Lumus Vision is an extremely interesting company that has pioneered reflective waveguides. The light is generated by the micro-display, and expanded and directed into the eye by the waveguide through a series of transflective mirrors. Lumus has been working on reflective waveguides for many years, but they have proved to be extremely difficult to manufacture at scale. However, Lumus has recently renewed its partnership with the glass manufacturing giant Schott to manufacture their waveguides at a new facility in Malaysia. This is a significant step forward for the technology, and could pave the way for the first generation of consumer AR glasses.
\nThere are a plethora of display technologies available, however only a limited subset are applicable to AR glasses. This is because the losses in a waveguide based system are astronomical, as demonstrated in the diagram below from Karl Guttag 11.
\n![\"Diffractive](\"/ar-glasses/waveguide-efficiency.webp\")
\nIn order for a suitable amount of nits to reach the eye, the display must be capable of generating a huge amount of light. For this reason, Liquid Crystal on Silicon (LCOS) and MicroLED are the most promising display technologies for AR glasses. Many industry experts see MicroLED as the future of AR displays, with companies like Apple and Meta investing heavily in the technology. Read more about MicroLED in this excellent article by Karl Guttag 11.
\nAR glasses are a full stack problem, and the compute requirements are variable and intense. Let's break down the compute workload into its core components:
\nOffloading as much compute as possible to auxiliary devices is key to the success of AR glasses, for both power and performance reasons. Most humans carry around a supercomputer in their pocket, and we can push even more intensive workloads to the cloud, much like Meta is already doing for their current generation of smartglasses.
\n![\"Compute](\"/ar-glasses/compute-move.png\")
\nEach task has to be carefully placed somewhere on this unconjoined triangle of success, trading off latency, power usage and compute requirements. For the neural network use case, using a Router LLM 12 to evaluate query difficulty and selecting the most appropriate device for the workload seems like an interesting avenue.
\nIn order to leverage the two external compute sources, you will need to have a high bandwidth communication channel. In the cloud case, this will obviously be done via WiFi or cellular networks. However, for the glasses-to-phone connection, more care must be taken in selecting the communication protocol. Bluetooth 5.0, whilst ubiquitous and low power, has an extremely limited bandwidth of 2Mbps, which is far from sufficient for AR glasses. Ultra Wideband (UWB) 13 could be the answer, and has been shipping in the iPhone since 2019, albeit purely for real-time location not data transfer. With a theoretical transfer rate of 1Gbps, it may be the solution that AR and other high bandwidth peripherals are looking for. Meta has been filing numerous patents relating to UWB (e.g 20240235606), so we may see this technology in their glasses soon.
\n![\"Tomorrow's](\"/ar-glasses/tomorrows-glasses.png\")
\nThe AR glasses we get in the next ~3 years will not be the ultimate form. In order to rival Manfred's glasses in Accelerando, we will need better specs. The above schematic is taken from one of Meta's most recent patents 14, and demonstrates some interesting innovations such as transducers in both the nose bridge and arms for more immersive spatial audio.
\nIn my opinion, the most salient difference between the first and final generation of AR glasses will be the display technology, particularly the resolution and FOV. As the visual fidelity of the display approaches that of the human eye, the line between the real and the virtual will become blurred, then vanish entirely. This may arrive sooner than you think, with commerical micro displays already hitting 4K4K resolution @ ~3µm pixel pitch 15.
\nIf we were to assemble a team to build the ultimate AR glasses, it would look something like this:
\nFrom all my research, it seems like Meta is well positioned in this market. I would not be surprised to see exponential adoption if they can deliver on their 2028 roadmap. With the recent advancements in AI, the value proposition for AR glasses has gone from exciting to essential. With developments from Lumus and Meta's recent patent filings, it seems like this technology may finally be approaching its "iPhone moment" - provided the near eye optics problem can be solved.
\nOnce the "iPhone moment" is reached, there will be a plethora of regulatory and social challenges to overcome, as highlighted in this excellent SMBC comic. In spite of these challenges, I'm looking forward to the future of AR glasses, and the startups and innovations that will come with them.
\nTLDR: Long $META.
\nI'm grateful to Mithun Hunsur, Madeline Ephgrave & Benjamin Perkins for their feedback on this post.
\nhttps://www.roadtovr.com/report-meta-ar-glasses-orion-connect-2024/ ↩
\nhttps://www.buildwagon.com/What-happened-to-the-Hololens.html ↩
\nhttps://www.acumenfinancial.co.uk/advice/airpods-revenue-vs-top-tech-companies/ ↩
\nhttps://www.rfpmm.org/pdf/how-to-win-friends-and-influence-people.pdf ↩
\nhttps://www.uploadvr.com/meta-hud-glasses-wont-be-ray-bans/ ↩
\nhttps://www.radiantvisionsystems.com/blog/ride-wave-augmented-reality-devices-rely-waveguides ↩
\nhttps://www.optofidelity.com/insights/blogs/comparing-and-contrasting-different-waveguide-technologies-diffractive-reflective-and-holographic-waveguides ↩
\nhttps://kguttag.com/2023/03/12/microleds-with-waveguides-ces-ar-vr-mr-2023-pt-7/ ↩ ↩2
\nhttps://image-ppubs.uspto.gov/dirsearch-public/print/downloadPdf/20240118423 ↩
\nhttps://kguttag.com/2024/04/20/mixed-reality-at-ces-ar-vr-mr-2024-part-3-display-devices/ ↩
\nIn order to perform classification on a functional brain scan it first\nundergoes many preprocessing steps. One of these steps is the transformation from\nthe timeseries output of the fMRI scan, transforming a m×n (where m is\nthe number of timepoints recorded and n is the number of brain regions used)\nto an n×n matrix of similiarity values (called a connectome). This\nsimilarity value is a measure of the neural synchronization between the 2 regions.
\nSo how do we quantify the similarity of 2 different timeseries? This blog post will\nexplore the common ways of quantifying time series similarity in a neuroscientific\nsetting. Before we get into the actual methods used to calculate time series similarity,\nwe need to cover the corner stone of almost all of the methods we are about to\nexplore - covariance.
\nCovariance is simply a measure of how two random variables change together.\nBelow is the formula for calculating covariance:
\nΣ=E[(X−E[X])(X−E[X])T]\nIn the case of fMRI, we have a multivariate random variable, allowing us to use\nMaximum Likelihood Estimation to estimate the covariance matrix. Below is a toy\nexample of our estimated covariance matrix.
\nvar(X)cov(X,Y)cov(X,Z)cov(X,Y)var(Y)cov(X,Y)cov(X,Z)cov(X,Y)var(Z)\nThis estimated covariance matrix is called the empirical covariance matrix. In\npractice we don't use the covariance matrix for our fMRI analyis.\nThis for a number of reasons:
\nMany of the following approaches will aim to address some or all of the above\ndrawbacks with empirical covariance.
\nTo address the excessive estimation error, it is common to perform a\ntransformation to the covariance coefficients, known as "shrinkage". In\ntheir seminal paper, "Honey, I shrunk the Sample Covariance Matrix" [1], Ledoit & Wolf\nproposed using shrinkage to regularize the sample covariance matrix.\n"Shrinkage" as the name implies, pulls the most extreme covariance coefficients\ntowards more central values. This not only resolves our excessive estimation\nerror, it can also make the matrix easily invertable by encouraging numerical\nstability. (The interested reader should consult the scikit-learn docs\n[2] )
\nNow that we have a well conditioned covariance matrix, we can attempt to address\nsome of the other identified drawbacks.
\nPearson's correlation coefficient (Pearson's R) or simply correlation, is\nthe most commonly used method to quantify similarity between 2 fMRI timeseries.\nCorrelation is a linear metric computed from the covariance of the 2 timeseries.\nBelow is the mathematical formula to compute correlation for a pair of random\nvariables:
\nρX,Y=σXσYcov(X,Y)\nCorrelation is widely used in neuroscience as it has long statistical history\nand is bound between -1 and 1. However, correlation does have some disadvantages.\nThe below figure should demonstrate one clearly:
\n![\"Correlation](\"/corr.gif\" "\\"Correlation")
Due to correlations linear nature, the same timeseries being slightly out of\nphase causes a huge decrease in the correlation value. Additionally, correlation\nprovides no distinction between whether 2 regions are directly\nconnected or indirectly connected via another region. To account for this, we\ncan use partial correlation!
\nPartial correlation is a variant of PCC that attempts to address\ndistinguishing between direct and indirect connections. This is done by\ncomputing correlation between regions after regressing all other timeseries.\nPartial correlation is computed from the inverse of the covariance matrix\n(this is where the shrinkage comes in handy), also known as the precision matrix.\nBelow is the mathematical formula for computing partial correlation for a pair\nof random variables:
\nρ=−piipjjpij\nwhere p is the respective precision matrix (obtained through inversion of our\nwell conditioned covariance matrix).
\nTangent space parameterization is a recently proposed solution to the problematic\nnature of SPD matricies, in that they don't naturally form Euclidean spaces.\nVaroquax et al. [3] proposed a solution to this\nsticky problem. Whilst I won't profess to be an expert on differential\nmanifolds or Riemmenian geometry, here is how I understand it.
\nAny covariance matrix is a symmetric positive semi-definite matrix. In order to\npreserve the geometry during our transformations, we need to\nperform a projection into a space where our geometry is preserved.\nAs SPD matricies form a diffentiable manifold, we can use this fact and\nRiemannian geometry to keep our geometric intutions in tact.
\n![\"Manifold\"](\"/manifold.jpg\" "\\"Manifold\\"/")
Image source
\nThe above image makes it clear that if we perform Euclidean operations like\nsubtraction in our manifold, the distances and therefore the geometry will break\ndown. By defining a homomorphic (ie structure preserving) tangent space, we can\napproximate the distance on the manifold using Euclidean distance in the tangent\nspace!
\nBut we aren't quite out of the woods yet, since we still need to define the\nreference point G on the manifold. As all of our covariance matricies need\nto be projected into the same tangent space, we need to choose an appropriate\nreference point. To do this, Varoquax et al. compute the geometric mean of the\nmatricies to define a reference space.
\nIn their paper, Dadi et al. [4] showed\nthat across many fMRI classification tasks, tangent space parameterization\nshould be preferred over both correlation and partial correlation.
\nDynamic Time Warping (DTW) is not like the previously discussed methods, as it does\nnot rely on covariance matricies. DTW is one of the canonical algorithms used for measuring\nsimilarity between timeseries. It is commonly used in scenarios where timeseries\nmay have a different length or be temporally offest. DTW is a non linear metric,\nand calculates the optimal match between two given sequences with certain\nrestrictions.
\nDTW was first proposed for usage with fMRI timeseries by Meszlényi et al. [5].\nThey propose its value for fMRI timeseries analysis as it has been noted that\n"dynamic switch of brain states" can cause non-stationary time lags. If there\nare known or unknown mechanisms in the brain that cause regional time differences\nin hemodynamic response - DTW could be a great way to account for this.
\nBelow is a Rust implementation of windowed DTW.\nThe window refers to the maximal warping distance.
\nlet m = s.len() + 1;\nlet n = t.len() + 1;\nlet mut dtw = Array::from_elem((m, n), f64::MAX);\n\ndtw[[0, 0]] = 0.;\n\nlet max_window = i32::max(*window, i32::abs((n - m) as i32));\nfor si in 1..n {\n let lower_bound = i32::max(1, si as i32 - max_window);\n let upper_bound = i32::min(m as i32, si as i32 + max_window);\n for ti in lower_bound as usize..upper_bound as usize {\n let cost = distance_fn(&s[si - 1], &t[ti - 1]);\n dtw[[si, ti]] = cost + \n f64::min(\n f64::min(dtw[[si - 1, ti]], dtw[[si, ti - 1]]), \n dtw[[si - 1, ti - 1]]\n );\n }\n}\nDynamic Time Warping does have some disadvantages. One particularly\ntroublesome one being the computational complexity - O(n2). The paper from\nMeszlényi et al. does come with a DTW implementation for fMRI data on GitHub,\nhowever it is extremely cumbersome and difficult to get running (particularly\non HPC systems where you don't have root access.)
\nTo address this I wrote RustDTW, a Python module backed by Rust code that works\nnatively within Python on numpy matricies. Check out an example of using DTW\nfor fMRI classification here.\nMy basic testing shows that it should be nearly 10x faster than a mutlithreaded\nPython implementation.
\nWhilst not meant to be an exhaustive exploration of the different connectivity\nmetrics explored above (and itself is not an exhaustive exploration of available\nconnectivity metrics, for more please see Pervaiz et al. [6]\nI thought I would demonstrate a comparison between them using a large open source dataset.\nYou can check out the inner workings and run it for yourself on Colab here!
\nFor those just interested in the results, the below graph summarizes the\nperformance of each of the similarity metrics.
\n![\"Accuracy\"](\"/accuracy.png\" "\\"Accuracy\\"/")
The results are pretty interesting! We can see that both non-linear matching and\ntransforming the covariance matricies into a natural space can provide improvements\nover standard correlation. Perhaps a future metric will incorporate\ncharacteristics from both tangent space parameterization and DTW to improve\nperformance.
", "url": "https://fleetwood.dev/posts/fmri-timeseries-similarity", "title": "An exploration of fMRI timeseries similarity metrics", "summary": "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", "date_modified": "2021-07-10T00:00:00.000Z", "author": { "name": "Christopher Fleetwood", "url": "https://fleetwood.dev" } } ] }