Standard attention spends most of its life stuck in memory traffic on “the slow road” between lightning-fast SRAM (19 TB/s) and off-chip HBM (1.5 TB/s). In this blog, we will discuss if we can avoid taking this slow road!!

Introduction

Standard attention [4] scales quadratically with sequence length because of memory bandwidth, not compute capability. In standard attention, output O is calculated as:

The intermediate matrix S = QKᵀ has N rows and N columns. At a 32k-token context, S holds about 1 billion numbers. That is per head, per layer.

Because S cannot fit into fast on-chip SRAM (~20 MB), standard execution keeps S in High-Bandwidth Memory (HBM) (~40 GB) and runs three separate sequential kernels:

1) Compute S = QKᵀ and write S to HBM.

2) Read S from HBM, compute P = softmax(S) in SRAM, and write P to HBM

3) Read P from HBM, compute O = PV in SRAM, and write O to HBM.

Every number in S and P crosses the HBM bus back and forth 4 times. Tensor cores spend most of their cycles waiting for memory transfers. The workload is strictly memory-bandwidth bound.

FlashAttention 1 — Stop Moving the Big Matrix

FlashAttention-1 [1] restructures attention to compute exact softmax without ever writing the N*N matrix i.e., S to HBM.

Idea 1: Block-Wise Tiling

Instead of computing the entire matrix S at once, load blocks of Q, K, and V into SRAM, compute a local tile of S on-chip, update the running output, and discard the tile. The full S matrix is never stored anywhere in global memory.

Idea 2: Online Softmax

Softmax requires normalizing every score by the sum of its entire row:

To evaluate softmax block-by-block without seeing the full row upfront, FlashAttention tracks two running statistics per row: the maximum value m and the running sum ℓ.

Take the row of scores [2, 1, 3, 5], processed in two blocks of two.

Block 1 [2, 1]:

m=2

ℓ = e^(2−2) + e^(1−2) = 1 + 0.368 = 1.368

Block 2 [3, 5]: Max increases to m = 5.

correction = e^(2−5) = 0.0498

ℓ = 1.368 × 0.0498 = 0.068 (old sum, corrected)

ℓ = 0.068 + e^(3−5) + e^(5−5) = 1.203

This matches if we have taken the full row (e^{x_i — m}): e⁻³+e⁻⁴+e⁻²+e⁰ = 1.203.

FlashAttention-1’s outer loop iterates over K/V blocks, with the inner loop sweeping across Q blocks.

SRAM can only hold state for one row at a time. Because the inner loop changes Q rows on every step, the running row state (O, m, l) must be read from HBM, updated, and written back to HBM for every step in the grid.

This is repeated for all the Q blocks.

FlashAttention 2 — Stop Leaving Workers Idle

FlashAttention-2 [2] retains the exact math and online softmax logic of FlashAttention-1, but redesigns the execution schedule and GPU thread mapping to achieve a 2x speedup. It simply changes the schedule: which worker does which piece, and in which order.

Change 1: Loop Swapping (Q-Outer Loop)

FA2 flips the loop order: the outer loop iterates over Q blocks, and the inner loop streams through K/V blocks. A worker takes a single Q block and runs it across all K/V blocks to completion. The running row state (O, m, l) stays pinned inside GPU registers for the entire inner loop and is written to HBM exactly once.

FA-1 Row-wise: FlashAttention 1 walks down the columns. Load K₁V₁, apply it to row 1, then row 2, then 3, then 4. Every step lands on a new row. So every step swaps a row state through HBM.

FA-2 Column-wise: FlashAttention 2 walks across the rows. Take row 1 and run it against K₁, K₂, K₃, K₄ to completion. Then move to row 2.

Change 2: Warp Partitioning (Split Q, Not K)

A single GPU SM splits its workload among 4 warps (groups of 32 threads executing in lockstep).

• FlashAttention-1 (Split K): Each warp gets a slice of keys (K). Because every warp holds partial outputs for every row, no single warp owns the complete output. Every block requires writing partial outputs to shared memory, waiting at a warp barrier, and executing a merge pass to combine partial sums.
• FlashAttention-2 (Split Q): Each warp gets a dedicated subset of Q rows and reads the entire K/V block. Warp 1 fully owns its assigned rows from start to finish.

PagedAttention — Stop Reserving Memory You Will Not Use

A model generates one token at a time. Each new token attends to every earlier token’s keys and values. Recomputing them each step is wasteful, so they are cached. The cache grows by one entry per token and shrinks only when the request ends. So, while FlashAttention speeds up execution kernels during training and prompt processing, model serving introduces a distinct memory bottleneck: the Key-Value (KV) cache during generation.

At a 2048-token context, one request consumes 1.6 GB of memory solely for its KV cache.

The Virtual Memory Solution

Traditional serving systems allocated continuous slabs of HBM upfront for a request’s maximum possible sequence length (e.g., 2048 tokens). This resulted in 60–80% memory fragmentation and waste from unfilled context reservations, blocking concurrent requests.

PagedAttention[3] adapts virtual memory OS principles to the KV cache:

1. Divide the KV cache into fixed-size physical blocks (e.g., 16 tokens).
2. Store physical blocks non-contiguously in an HBM memory pool.
3. Maintain a dynamic block table which maps logical sequence blocks to physical pool indices.

When a request generates new tokens that fill a block, the runtime dynamically allocates a new physical block from the pool and appends it to the block table.

Conclusion

• FlashAttention-1 eliminated standard attention’s O(N²) HBM memory bandwidth bottleneck using block-wise tiling and online softmax rescaling.
• FlashAttention-2 restructured the execution loop order to pin row states in registers and split Q blocks across warps, removing shared memory sync overhead to improve GPU throughput limits.
• PagedAttention translated operating system virtual memory paging to the KV cache, eliminating memory fragmentation to maximize inference throughput and request batching.

References

1. Flash Attention — 1
2. Flash Attention — 2
3. Paged Attention
4. Attention is all you need

1) Introduction

Processing massive contexts (100k+ tokens) in Large Language Models requires a fundamental trade-off between memory footprint and recall accuracy. Standard Transformers maintain exact recall across long context windows by storing every token’s key and value vectors in a growing KV cache. While effective, this creates an O(N) memory bottleneck during inference.

Linear attention architectures attempt to solve this by collapsing sequence history into a fixed-size memory matrix, but early designs suffered from crude uniform decay and “memory tie” bugs when trying to update stored facts.

Kimi Delta Attention (KDA)[1] introduced in Moonshot AI’s Kimi K3 model architecture combines two mechanisms to make fixed-memory linear attention work in practice: per-channel diagonal decay gates and the Delta Rule overwrite mechanism.

2) Problem: KV Cache Scaling vs. Naïve Memory Collapse

In standard quadratic attention, every new token compares itself against every past token. To do this during decoding, the inference engine maintains a KV cache:

• At 1,000 tokens: Memory footprint is negligible.
• At 1,000,000 tokens: The KV cache demands hundreds of gigabytes of VRAM per batch, causing memory bandwidth bottlenecks that throttle decoding speed.

2.1) The Naïve Linear Attention Failure Mode

Earlier linear attention architectures (and classic RNN states) try to maintain a fixed-size state matrix S. However, they update memory using standard additive accumulation:

The Bug: If the model learns at token 10 that meeting-day = Tuesday, and at token 500 learns meeting-day = Thursday, naïve addition leaves both facts active in the state matrix:

Furthermore, applying a single scalar decay rate forces every feature (local syntax vs. long-term semantic facts) to fade at the exact same speed.

3) How Kimi Delta Attention Works (Anatomy of KDA Write)?

KDA replaces the unbounded KV cache with a fixed-size 128x128 state matrix per attention head. Let’s consider this example:

Earlier: “The meeting is Tuesday. Sarah will present.”

Now: “…actually, the meeting moved to Thursday.”

At Step 5: When is the meeting?

Step 1: Memory Footprint — Static Footprint

Instead of appending a new row to a cache array, KDA passes the incoming token through a constant-sized matrix S (see fig 1 — right side). Memory utilization at token 100 and token 1,000,000 is identical.

Step 2: Read Under Key (v^old = S^T k)

Before writing new information, KDA probes the current state to check what is already stored under the incoming token’s key slot. The model processes the word “Thursday” in the context of the meeting moving. It generates a key k to check what is currently stored for meeting-day.

Key vector k = [1, 0, 0, 0]^T (selects row: meeting-day)

Reads row meeting-day -> v_old = [1, 0, 0, 0] (“Tuesday — stale!”)

Step 3: Delta Write

To update the record without causing value collision, KDA subtracts v_old before adding the new value v = [0, 1, 0, 0]^T (Thursday):

Applying the rank-1 outer product update k(ΔV)ᵀ:

The -1 (magenta) explicitly cancels out Tuesday, while +1 (blue) writes Thursday. All other key rows remain untouched.

Step 4: Per-Channel Forgetting (Diag(a).S)

Rather than multiplying the entire matrix by one scalar decay factor, KDA computes a data-dependent gating vector a:

• Long-term semantic channels: a1 = 0.99 => meeting-day retains 99% of its state.
• Short-term syntactic channels: a3 = 0.60 => transient grammar representations shrink rapidly to clear space for future tokens.

Step 5: Read with Query (out = S^T q)

Pages later in the sequence, the model queries the state matrix: “So — when is the meeting?”

Query q = [1, 0, 0, 0]^T (“When is the meeting?”)

out = S^T q -> [0, 0.99, 0, 0] ≈ “Thursday”

Because the Delta Rule subtracted the stale Tuesday state in Step 3, the output returns a unambiguous signal for Thursday instead of a 50/50 tie.

The Hybrid 3:1 Interleaving Attention Design

To overcome pure linear attention’s weakness in exact long-context retrieval, the Kimi K3 architecture interleaves Kimi Delta Attention (KDA) layers with full Multi-Head Latent Attention (MLA) layers in a 3:1 ratio. This is necessary because a memory of 128x128 numbers cannot hold 1M tokens correctly. Because of this, the model cannot complete tasks that require the exact recall of old data. To solve this, the Multi-Head Latent Attention (MLA) layer keeps a full Key-Value cache for all positions. This full cache operates as a reliable backup for the KDA layers.

4) Conclusion

By replacing standard KV caches with constant-sized state matrices, KDA enables architectures like Kimi K3 to maintain high generation speeds and low memory overhead over 1M-token context windows. The combination of per-channel diagonal decay (selective retention) and the Delta Rule (in-place memory overwriting) resolves the fundamental bottlenecks of classic linear attention.

By only keeping an full KV cache for 1 in every 4 layers, Kimi K3 achieves up to 75% reduction in total KV cache memory footprint compared to all-attention architectures while maintaining near-identical performance on complex benchmarks

References

1. Kimi Linear — https://arxiv.org/pdf/2510.26692
2. https://www.kimi.com/blog/kimi-k3

Every frontier model claims a million-token context window now. Most of them can’t actually use it. You paste a long document, the KV cache eats your GPU memory, inference slows to a crawl, and you’re paying for a context window that’s technically there but practically unusable. DeepSeek V4[1] takes this problem seriously at the architecture level and they are doing it open-source. At 1M tokens, V4-Pro uses 27% of the compute and 10% of the KV cache compared to V3.2. The Flash variant drops that to 10% and 7%. That’s not an incremental improvement. That’s a different cost structure!

The trick: two types of attention, alternating

V4 splits attention into two mechanisms and interleaves them across its 61 layers.

The first is Compressed Sparse Attention (CSA). It groups every 4 tokens into one compressed entry (4x reduction), then runs a tiny FP4 scoring network to find the top 1,024 most relevant entries per query. Out of 250,000 compressed entries at 1M tokens, it only looks at about a thousand. The rest get skipped. The paper claims 99.7% recall, meaning almost nothing important gets missed.

The second is Heavily Compressed Attention (HCA). This one is more aggressive: 128 tokens compressed into 1. At 1M tokens that’s around 7,800 entries, small enough to just attend to all of them. No sparse selection needed. Dense attention over a tiny sequence is cheap.

Every other layer alternates between the two. CSA gives you precise lookups into specific parts of history. HCA gives you a rough view of everything. Together, the model gets both a magnifying glass and a wide-angle lens.

How they trained it

Pre-training used 32T+ tokens with a progressive context schedule. The model starts on short 4K-token sequences with normal dense attention. Then 16K. At 64K, sparse attention kicks in for the first time. Finally, 1M. The model learns language before it learns to speed-read. That ordering matters.

Post-training is where V4 does something genuinely different. Instead of fine-tuning one model on everything, they train five separate specialists: math, code, agent tasks, chat, and logic. Each one gets SFT plus reinforcement learning with GRPO, trained against real environments. The code specialist actually executes code. The math specialist verifies proofs.

Then they merge all five into one model through distillation. But they use on-policy distillation.

“On-policy” Distillation

In standard distillation i.e., off-policy, the teacher generates a bunch of outputs, you save them as a dataset, and the student memorizes them. It’s like learning to cook by reading a recipe book. You never actually cook anything.

V4 does on-policy distillation instead. The student generates its own output first. Then each specialist teacher looks at what the student produced and provides a full probability distribution - not just “the answer is X” but “here’s how confident I am about every possible token at every position.” The student adjusts to match those distributions using reverse KL divergence.

Why reverse KL specifically? Because it penalizes the student for being confident where the teacher isn’t. It makes the student sharp and precise in each specialist’s domain rather than producing a mushy average. And it pairs naturally with on-policy generation: you only need teacher probabilities at the student’s own samples.

So, instead of studying recipes, you cook your own dish, the chef tastes it, and tells you what to fix.

This matters because at 1.6 trillion parameters across five domains, the distribution mismatch from traditional distillation would be brutal. The student would train on outputs it would never naturally produce, then fall apart the moment it drifts during inference.

My Take

Scaling has made everything incredibly compute-intensive, so it’s no surprise that AI labs are now pivoting hard toward optimization. With DeepSeek V4 coming in at roughly 10x cheaper than models like ChatGPT and Claude, it presents an unbelievably cost-efficient choice for developers.

That said, I suspect many of the closed-source labs are already running similar optimizations behind closed doors. They’re just charging a premium because the current market bears it. But with open-weights models offering comparable level of performance for less, those closed labs might be forced to slash their prices the second they feel users starting to jump ship.

pro-tip: These days I use claude to make UI (html/JS) and then beautify using gemini banana pro.

References

1. https://huggingface.co/deepseek-ai/DeepSeek-V4-Pro/blob/main/DeepSeek_V4.pdf
2. mHC: Manifold-Constrained Hyper-Connections — https://arxiv.org/pdf/2512.24880
3. DistillBert — https://arxiv.org/abs/1910.01108

Back in 1914, Henry Ford baffled his peers by doing the unthinkable: he doubled his factory workers’ pay to $5 a day. On paper, it was a move to reduce employee turnover. In reality, Ford understood something much bigger. His employees were also his target demographic. If the people building the Model T couldn’t actually afford to buy one, he was eventually going to run out of customers.

Essentially, Ford controlled the entire economic loop. The wages he paid out eventually walked right back through his front door as car sales. It wasn’t corporate generosity; it was a calculated strategy where higher payroll costs paid for themselves by directly boosting his own demand.

Fast-forward to today, where AI is making it possible to replace white-collar workers at a massive scale. Every CFO running the numbers is acting just as rationally as Ford did a century ago, but the incentives have flipped. Ford looked at his ledger and realized he needed to pay his people more. A modern CFO looks at theirs and realizes they don’t need the people at all e.g., “Block laid off 40% employees in Q1 26 [2]”.

A recent economic paper, The AI Layoff Trap [1] breaks down exactly why this is happening, proving mathematically that this “rational” corporate behavior is actually leading us toward an economic cliff.

The Trap

Let’s say there are ten competing firms in a market. Their employees are also the economy’s consumers. Then AI shows up, offering to replace their workforce for pennies on the dollar. The math seems straightforward. The gain is massive labor savings that instantly boost profit margins. The loss is that laid-off workers stop spending, dragging down the broader economy.

But the economic damage doesn’t hit evenly. If your firm fires people, those workers cut their spending at all ten companies. You only absorb about 10% of the lost sales, while your competitors eat the other 90%. You keep all the savings, but outsource most of the pain. Economically, it’s a no-brainer. You automate.

But the other nine firms are doing the exact same calculus. Everyone automates. Thousands lose their jobs, and overall consumer demand falls off a cliff. Now, all ten companies are trying to sell products to a market that can no longer afford them. Margins shrink, people lose their livelihoods, and everyone takes a hit.

The real tragedy here is that nobody could have chosen differently. Even knowing the catastrophic outcome, automating remains the dominant strategy for every individual firm. This is the Prisoner’s Dilemma, but for the AI age.

The Math Behind the Trap

Imagine a town with 10 competing bakeries (N = 10). They are all thinking about firing some human cashiers and replacing them with AI checkout kiosks. Assume each bakery has 100 human employees.

Before we go deeper, let’s define some of the fancy terms from the paper:

Cooperative optimum: If all firms sat in a room together and agreed on a single automation level that was best for everyone — firms, workers, and the economy — what would they choose?

αCO = (cost saving − full demand loss) / friction

Nash equilibrium: If every firm ignores what others do and just optimizes for itself, where does everyone end up?

αNE = (cost saving − demand loss/N) / friction

Every time a bakery fires an employee, following things happens:

• Wage Savings: Firm saves $40k salary per employee
• Demand Loss: Assume that employee couldn’t find a new job so expect that employee to spend less. They may get some social benefits from the Government so let’s say they only spend $10k which means the town loses $30k in sales.
• Friction Cost: There’s a cost associated for AI automation and more we automate it can ramp up. For e.g., if you fired a many ppl from a team then that means you must be using a complex agentic ecosystem so cost will be much higher. For our example, 1st AI is easy (costs $1,000), the 10th AI is harder ($10,000), and the 37th is very complex ($37,000)

Now, we can plug our numbers in the 2 equations defined above:

Cooperative optimum: A smart mayor looking at the whole town wants to maximize total economic health. They weigh the cost savings against the full loss of consumer demand. It comes down to 10%.

Optimal Automation = (Cost Saving − Full Demand Loss) / Friction
Optimal Automation = ($40,000 − $30,000) / $100,000
Optimal Automation = $10,000 / $100,000 = 10% (10 employees)

The Bakery’s Limit : The bakery owner only cares about the tiny slice (1/N) of lost demand that impacts their store. It comes down to 37% (Trap)

Actual Automation = (Cost Saving − [Full Demand Loss / N]) / Friction
Actual Automation = ($40,000 − [$30,000 / 10]) / $100,000
Actual Automation = ($40,000 − $3,000) / $100,000
Actual Automation = $37,000 / $100,000 = 37% (37 employees)

Monopoly Paradox: If we only have 1 company i.e., it owns 100% of the market and if it displaces workers, it feels 100% of the resulting demand destruction. So it self-corrects!

Why don’t they just cooperate? If one noble bakery stops at 10%, the other 9 will keep firing up to 37% and pocket the massive savings. The noble bakery will still suffer from the town’s economic collapse, but they will be stuck with high labor costs and will get bankrupted by their rivals ( Prisoner’s Dilemma),

Why the standard playbook fails?

The paper talked about several fixes like 1) lowering employee wages 2) worker equity 3) retraining programs 4) universal basic income (UBI). Let’s dig deeper into the UBI as it has already been flouted in many countries!

In UBI, the local government will play a crucial role. For the demand loss from the layoff, the government will give a $$ check to the worker which will offset the loss in the demand. But in order to completely offset this, they need to pay the entire wage for the worker which is probably not feasible specially for the developing nations. In developed nations like US, workers do get some social benefits but they are fraction of the wages.

The One Thing That Works - A Pigouvian Tax

So, what could actually work? A simple idea from the early 1900s called a Pigouvian tax [3]. Basically, if the firm causes damage to the society, we tax them just like we do it for the environment by imposing “Carbon Tax”. According to the authors, we need to treat the AI automation in the similar way. Whenever a company fires an employee, it is indeed impacting the overall society. The fix is to charge them the same penalty equivalent to the damage they cause to the economy!

The optimal tax rate:

τ = demand loss × (1 − 1/N)*

This doesn’t ban automation instead it just stops companies from excessive automation that leads to economic impact to the society by creating a some sort of self-defeating loop!

My Take

We are already seeing the early stages of this automation race, so it’s incredibly valuable that economists are highlighting this trap now. It gives us a crucial window of time to implement interventions before the damage becomes widespread and irreversible.

However, my biggest worry is that this time really is different. With previous waves of automation, the standard rule was simply to “retrain workers” for new, higher-level tasks. But AI isn’t just a dumb machine, it’s a cognitive engine. If AI continues to evolve and consistently outsmarts humans, it’s going to get increasingly difficult to find any work that a human can do better or cheaper.

If the concept of money still dictates our lives in the future, someone has to pay the bills!!!

References

1. AI layoff trap paper
2. https://www.nytimes.com/2026/02/26/technology/block-square-job-cuts-ai.html
3. https://en.wikipedia.org/wiki/Pigouvian_tax

More Readings/Podcast

• https://economics.mit.edu/news/daron-acemoglu-what-do-we-know-about-economics-ai
• https://www.youtube.com/watch?v=RB_WmoH5nQ4&t=2632s — Interesting discussion on this topic

Introduction

If you look at the current landscape of AI, we are completely obsessed with generation. From Grok’s Imagine to OpenAI’s Sora, the dominant paradigm is that if an AI can generate the pixels, it must understand the image.

But Yann LeCun (NYU Prof./ex Meta Chief Sci), has been arguing a different, somewhat controversial point: Generative models are fundamentally inefficient at understanding the world. To build AI that actually comprehends reality, we need to stop forcing models to paint and start teaching them to predict. Meta’s open-source model “I-JEPA” (Image Joint Embedding Predictive Architecture) that represents a massive, new shift in how we approach computer vision.

In this post, we are going to break down the motivation behind I-JEPA, look at a real-world example of how its architecture works, and explore how it acts as a stepping stone towards the “World Models.” This is a new modeling paradigm and many AI research departments are actively working on it. For e.g., WorldLabs, Genie 3, ami-labs.

Background: The Problem with Predicting Pixels

To understand why I-JEPA is different, we have to look at how current SOTA vision models learn.

Most self-supervised vision models rely on a technique called Masked Image Modeling (like the Masked Autoencoder, or MAE). The process is straightforward: take an image, black out 75% of it, and force the AI to reconstruct the missing pieces.

The problem? Pixels are messy, high-frequency, and largely irrelevant to actual meaning.

Imagine an image of a dog sitting on a lawn. If we mask a patch of the grass, a standard auto-encoder wastes massive amounts of computational power trying to predict the exact orientation of every single blade of grass, the exact shadow cast by a cloud, or the precise RGB value of a leaf.

Humans don’t process the world like this. If you see a dog behind a fence, you don’t hallucinate the exact fur pattern behind the wood; you just know, abstractly, that the rest of the dog is there. Pixel-reconstruction forces the model to memorize noise rather than learning high-level semantic meaning. And they also need tons of more data to do so because forcing a model to map out every possible random variation in lighting, texture, and background noise requires feeding it millions of extra, redundant examples just to help it grasp basic concepts.

The I-JEPA Architecture: Predicting Meaning, Not Images

I-JEPA throws pixel reconstruction out the window. Instead of predicting the missing pixels, it predicts the missing representations (embeddings).

The Masking Strategy: What Gets Hidden?

Instead of predicting pixels, I-JEPA predicts missing representations. To do this, it needs a clever way to hide information.

Think of a satellite image of a construction site as a grid of small squares. A standard 224x224 photo becomes a 16 x 16 grid of patches (each patch is 14 x 14 pixels).

Here is the exact masking workflow:

1. The Context: Pick a large region of the grid. This is what the model is allowed to look at (usually 85% of the patches).
2. The Targets: Pick 4 small rectangular blocks scattered around the grid. These are the parts the model must figure out without seeing them.
3. No Peeking: Ensure the targets and context don’t overlap. If the model could see a target patch, it would just copy the answer instead of actually learning.

It’s like covering 4 sticky notes on a photo and asking someone: “What’s behind each one?” They can see almost the whole picture, but those few hidden spots are what they need to reason about.

Key components of the model architecture:

1. Context Encoder — “The Observer”

This is a standard Vision Transformer (ViT). It only sees the unmasked patches of the image (the surrounding terrain, access roads, etc.). It converts those visible patches into dense mathematical embeddings. For the flagship model, this is a massive ViT-Huge (32 transformer blocks, 1280-dim embeddings, 16 attention heads).

2. Target Encoder — “The Answer Key”

This is an identical copy of the Context Encoder, but with a major twist: it never learns from gradients. If the target encoder updated at the same speed as the context encoder, the model could “cheat.” Both networks could collapse and just output the exact same trivial representation (like a vector of all zeros) for every image, achieving a perfect loss score while learning nothing.

Instead, the Target Encoder is updated as a slow-moving average of the Context Encoder’s weights (an Exponential Moving Average, or EMA). Intuitively it’s like a teacher grading a student: if the teacher constantly changed the answer key to match whatever the student guessed, the student wouldn’t learn anything. By keeping the answer key stable, the student is forced to actually figure out the problem.

3. Predictor — “The Guesser”

This is a smaller transformer that bridges the gap between “what I can see” and “what I need to infer.” It takes the context embeddings from the Observer and tries to predict what the Answer Key will output for the masked positions. Basically, goal is to make Context predict the same embeddings.

Compared to Pixel-Reconstruction Models (like Masked Autoencoders / MAE): I-JEPA is approximately 3x to 5x faster (and requires less memory) when scaling up to massive models. While MAE wastes compute forcing a decoder to draw high-frequency pixel details, I-JEPA’s predictor is lightweight and operates entirely in a compressed, low-dimensional mathematical space.

Fig 1 — Detailed Model Architecture

source

Case-Study — Track Data-Center Construction

Investors and supply-chain analysts may want to know exactly how fast tech companies are physically building massive AI data centers. We can use I-JEPA to build an Automated Infrastructure Watchdog. Using Google Earth Engine, we can periodically download satellite images and analyze the state of the construction.

To check if this model can identify differences among different states —Empty, Active Construction & Concrete , I pulled satellite images across different periods spanning 3 years using google earth engine [A.1]. You can see in the fig. below that model is able to correctly identify all the 3 different states.

The Predictor as a “Baby World Model”

The authors of the paper refers I-JEPA predictor as a “primitive (and restricted) world-model.” To put it simply, here is the difference between current AI and World Models:

• Current Models (Painters & Parrots): Today’s generative AI tries to predict exact surface details. If an AI watches a car drive behind a tree, a generative model wastes massive energy trying to predict the exact pixels of the leaves blowing in the wind or the reflection of the sun on the car’s paint. They focuses too much on tiny, unimportant details instead of the big picture and waste considerable amount of energy (i.e., it needs tons of data to detect difference between cat & dog)
• World Models (Common Sense): A world model predicts the underlying reality. It doesn’t care about drawing the exact pixel color of the leaves. Instead, it possesses “object permanence” — it simply knows the car still exists behind the tree and is continuing to move forward.

I-JEPA’s predictor takes a “baby” step toward this by operating in a limited spatial domain:

• Observation: The embeddings of visible patches (the context).
• Prediction: The embeddings of hidden patches (the targets).
• Uncertainty: Because it doesn’t commit to pixel-level details, it naturally handles the fact that multiple specific appearances are consistent with the same context.

Its successors, like V-JEPA (Video JEPA), extend this to temporal prediction — predicting what happens next in time, not just what exists elsewhere in space. V-JEPA: “Given the first 2 seconds of this video, what happens next?” (Temporal World Model) [2]

References

1. https://github.com/facebookresearch/ijepa
2. https://ai.meta.com/research/vjepa/

Appendix

A.1 Script to run inference on the satellite images

import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Dataset
import numpy as np
import matplotlib.pyplot as plt
from sklearn.manifold import TSNE
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score, classification_report
from sklearn.model_selection import train_test_split
import ee
import requests
from PIL import Image
import io
import torchvision.transforms as transforms

# ==========================================
# 1. The Embedding Engine (I-JEPA Wrapper)
# ==========================================
class IJEPASemanticExtractor(nn.Module):
    def __init__(self, embed_dim=1280):
        super().__init__()
        self.embed_dim = embed_dim
        self.mock_backbone = nn.Sequential(
            nn.Flatten(),
            nn.Linear(3 * 224 * 224, self.embed_dim)
        )
        for param in self.parameters():
            param.requires_grad = False

    def forward(self, x):
        return self.mock_backbone(x)

# ==========================================
# 2. Real Google Earth Engine Dataset
# ==========================================
class GEESatelliteDataset(Dataset):
    def __init__(self, project_id='data-center-tracking', num_samples=180):
        print("Authenticating with Google Earth Engine...")
        try:
            ee.Initialize(project=project_id)
        except Exception:
            ee.Authenticate()
            ee.Initialize(project=project_id)

        self.class_names = ["Empty Dirt", "Active Construction", "Concrete Foundation"]
        self.images = []
        self.labels = []
        roi = ee.Geometry.Point([-82.7486, 40.0886]).buffer(2500)

        timeframes = [
            ('2021-06-01', '2022-06-01', 0),
            ('2022-07-01', '2023-07-01', 1),
            ('2023-08-01', '2024-10-01', 2)
        ]

        transform = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
        ])

        print("Downloading satellite imagery via GEE...")
        for start, end, label in timeframes:
            collection = (ee.ImageCollection('COPERNICUS/S2_SR_HARMONIZED')
                          .filterBounds(roi)
                          .filterDate(start, end)
                          .filter(ee.Filter.lt('CLOUDY_PIXEL_PERCENTAGE', 30)))

            limit = num_samples // 3
            images_list = collection.toList(limit)
            count = images_list.length().getInfo()
            print(f"Found {count} images for label {label} in range {start} to {end}")

            for i in range(count):
                img = ee.Image(images_list.get(i))
                rgb_img = img.select(['B4', 'B3', 'B2']).visualize(min=0, max=3500)
                url = rgb_img.getThumbURL({'region': roi, 'dimensions': '224x224', 'format': 'png'})
                response = requests.get(url)
                if response.status_code == 200:
                    pil_img = Image.open(io.BytesIO(response.content)).convert('RGB')
                    self.images.append(transform(pil_img))
                    self.labels.append(label)

        if len(self.images) > 0:
            self.images = torch.stack(self.images)
        self.labels = torch.tensor(self.labels)
        print(f"Successfully loaded {len(self.images)} real satellite images!")

    def __len__(self):
        return len(self.images)

    def __getitem__(self, idx):
        return self.images[idx], self.labels[idx]

# ==========================================
# 3. Feature Extraction Pipeline
# ==========================================
def extract_embeddings(model, dataloader, device='cpu'):
    print("Extracting embeddings using frozen I-JEPA encoder...")
    model.eval()
    all_embeddings = []
    all_labels = []
    with torch.no_grad():
        for images, labels in dataloader:
            images = images.to(device)
            embeddings = model(images)
            all_embeddings.append(embeddings.cpu().numpy())
            all_labels.append(labels.numpy())
    return np.vstack(all_embeddings), np.concatenate(all_labels)

# ==========================================
# 4. Visualization (t-SNE)
# ==========================================
def visualize_semantic_space(embeddings, labels, class_names):
    print("Running t-SNE dimensionality reduction...")
    n_samples = embeddings.shape[0]
    perplexity_val = min(30, max(1, n_samples // 3))
    tsne = TSNE(n_components=2, random_state=42, perplexity=perplexity_val)
    embeddings_2d = tsne.fit_transform(embeddings)

    plt.figure(figsize=(10, 8))
    colors = ['#d95f02', '#7570b3', '#1b9e77']
    for class_idx, class_name in enumerate(class_names):
        idx = labels == class_idx
        if np.any(idx):
            plt.scatter(embeddings_2d[idx, 0], embeddings_2d[idx, 1],
                        c=colors[class_idx], label=class_name, alpha=0.7, edgecolors='w', s=60)
    plt.title(f"I-JEPA Semantic Embeddings (n={n_samples})")
    plt.legend()
    plt.grid(True, linestyle='--', alpha=0.5)
    plt.savefig('ijepa_tsne_clusters.png', dpi=300, bbox_inches='tight')
    plt.show()
    print("-> Saved visualization to 'ijepa_tsne_clusters.png'")

if __name__ == "__main__":
    # Doubling the sample size from 90 to 180
    dataset = GEESatelliteDataset(project_id='<Add your projectID>', num_samples=180)
    if len(dataset) < 5:
        print("Insufficient data to run pipeline.")
    else:
        dataloader = DataLoader(dataset, batch_size=16, shuffle=True)
        model = IJEPASemanticExtractor(embed_dim=1280)
        embeddings, labels = extract_embeddings(model, dataloader)
        visualize_semantic_space(embeddings, labels, dataset.class_names)

Optimizers are the engines that power neural network training. They determine how your model learns from data by updating weights to minimize loss. But not all optimizers are created equal. Let’s explore three key optimizers — Gradient Descent, Adam, and Muon — through visualizations.

1) Gradient Descent: The Cautious Hiker

Imagine you’re hiking down a mountain in thick fog. You can only see your immediate surroundings. How do you find the valley? Follow the slope downhill, one step at a time.

The Basic Idea: Gradient Descent is the simplest optimizer. The algorithm is beautifully straightforward:

w_new = w — learning_rate × gradient

At each step:

1. Calculate the gradient (the slope) at your current position
2. Take a step in the opposite direction (downhill)
3. Repeat until you reach the minimum

In this visualization, let’s assume we have a simple quadratic loss surface. The optimizer starts at position (2.0, 2.0) and follows the gradient arrows downhill, eventually finding the minimum at (0, 0).

1.1) How Big Should Your Steps Be?

The learning rate determines your step size, and choosing the right value is crucial.

In this visualization, we show how learning rate impacts the step size:

• Too Small (η=0.05): Like taking baby steps down the mountain. You’ll get there eventually, but it takes forever.
• Just Right (η=0.3): You make steady progress without overshooting. The path is smooth and efficient.
• Too Large (η=0.9): Like taking giant leaps. You overshoot the minimum, bounce to the other side, and keep oscillating. Sometimes you might even diverge and never converge!

1.2) When Gradient Descent Struggles

Problem 1: Oscillation in Narrow Valleys

Picture a narrow valley running horizontally. The gradient points across the valley (perpendicular to the valley direction), not along it toward the minimum. Because of this, you zigzag back and forth, making slow progress toward the actual minimum

Problem 2: Trapped in Local Minima

Loss surfaces often have multiple “valleys” — local minima that aren’t the global minimum. Gradient Descent has no momentum, so once it reaches any minimum (local or global), it stops.

2) Adam: The Smart Athlete

If Gradient Descent is a cautious hiker, Adam is a smart athlete with two secret weapons: momentum and adaptive brakes.

Momentum — The Ball That Rolls Through Obstacles

Think of it as giving the optimizer “mass” — it builds up speed in consistent directions and powers through obstacles. For e.g., the ball starts rolling toward the local minimum. As it descends, it gains momentum. Even though it reaches the local minimum, momentum carries it through! With this momentum, it escapes the trap and continues to the global minimum.

m(t) = beta × m(t-1) + (1-beta) × gradient

Weapon 2: Adaptive Learning Rate — The Brake System

Now imagine you’re speeding toward a target. As you get closer, you need to slow down to avoid overshooting. For e.g., The ball is approaching the minimum from far away. As the gradient gets smaller (flatter surface near minimum), the step size automatically reduces and the ball smoothly settles at the minimum without oscillating

v(t) = beta × v(t-1) + (1-beta) × gradient²
w_new = w — learning_rate × m / sqrt(v + epsilon)

Dividing by sqrt(v) normalizes step sizes across parameters. In dimensions with consistently large gradients (like valley walls), it reduces steps to prevent oscillation. In dimensions with small gradients (like gentle slopes), it effectively increases steps for faster progress. Meanwhile, as we approach the minimum, steps naturally shrink because the momentum term m decreases - the gradient itself is getting smaller everywhere.

3) Muon: The Matrix Master

While Adam treats each parameter independently, Muon recognizes that neural network weights are matrices with structure.

The Matrix Structure Problem

Every layer in a neural network uses a weight matrix to transform its inputs into outputs, where each number controls how strongly one neuron connects to the next. Consider a 2×2 weight matrix connecting 2 input neurons. These elements aren’t independent — they form a structured transformation.

Adam treats each element independently and update it based on its gradient and history. The grid WARPS! The rectangular structure distorts because each element goes its own way. But Muon updates all elements together as a coordinated transformation. The entire grid might rotate, scale, or shift but it will always stays rectangular! This is important for matrix-heavy models like transformers e.g., K,V & Q matrices, where weights represent mathematical transformations that should maintain their structural properties.

Under the hood, Muon uses “Newton-Schulz Orthogonalization” which converts the momentum for the gradients into an orthogonal matrix — where all columns (or rows) are perpendicular unit-norm basis vectors, then intelligently scales them. While not literally unit-norm, these updates tend to be more stable and consistent than Adam’s, which can shrink aggressively near the minimum

Step 1: Build momentum (same as Adam)
m(t) = β × m(t-1) + gradient

Step 2: Orthogonalize (Muon)
m_ortho = NewtonSchulz(m)
// Makes update directions perpendicular

Step 3: Scale and apply
update = m_ortho × scale
w_new = w — learning_rate × update

AdamW: Multiple explorers (parameters) clump together and they’re all walking in nearly the same direction. They thoroughly illuminate one narrow corridor (correlated), but leave most of the forest dark. Wasted effort on redundant coverage.

Muon: Multiple explorers fan out in a perfect starburst each perpendicular to the others. Each explorer’s beam illuminates a unique region with zero overlap. Same total light, but covering independent directions.

If you want to try it out, please use the interactive tool here (used gemini to build this visualization):

Muon vs Adam Explorer

Recently, researchers showed that using Muon optimizer improved their pre-training for their LLM. In their paper, they also implemented a variation where they clipped large logits from the attention output to make the training stable [2]

Also, I talked about this paper in my earlier blog which you refer for more details.

Understanding Kimi K2 Paper

References

1. Adam Paper
2. Kimi Paper
3. Muon Optimizer

1) Introduction

Monte Carlo methods are among the most powerful tools in modern computational statistics, named after the famous casino in Monaco for their reliance on randomness and probability. At their core, these methods solve a fundamental challenge: how do you work with probability distributions that are too complex to handle directly? Whether you’re trying to generate samples from an unusual distribution, estimate expectations without knowing normalization constants, or explore high-dimensional spaces, Monte Carlo methods provide elegant solutions through strategic random sampling.

This blog explores three fundamental Monte Carlo sampling techniques that form the backbone of modern ML algorithms: Rejection Sampling, Importance Sampling, and Markov Chain Monte Carlo (MCMC). Through visual explanations and practical examples, we’ll develop both the intuition and mathematical understanding needed to apply these methods effectively.

Note: I have been exploring a new visualization package (created by https://www.3blue1brown.com/)[2] for adding video content to the blogs. (Apologies if you see issues in formatting in the video — I’m still figuring out this library)

Also, if you want to explore the code for each of these methods, please check this package [1] which demonstrates simulation.

2) Rejection Sampling

Rejection sampling is a fundamental Monte Carlo method for generating samples from a probability distribution that may be difficult to sample from directly.

Suppose we want to generate samples from a target distribution p(x), but direct sampling is difficult or impossible. This situation arises frequently when p(x) has a complex form, such as a mixture of distributions or an unnormalized density.

The key insight is to use a simpler proposal distribution q(x) that we can easily sample from, then selectively accept or reject these samples to recover the target distribution.

The Algorithm

1. Choose a proposal distribution q(x) that is easy to sample from
2. Find a constant M such that M·q(x) ≥ p(x) for all x
3. Sample x from q(x)
4. Sample u uniformly from [0, 1]
5. Accept if u ≤ p(x) / (M·q(x)), otherwise reject and repeat

The acceptance criteria ensures that regions where p(x) is high relative to q(x) accept more samples, correctly reproducing the target distribution.

Example 1: Sampling from a Mixture Distribution

• Target p(x): Mixture of two Gaussians (green curve)
• Proposal q(x): Simple Gaussian distribution (blue curve)
• Envelope M·q(x): Scaled proposal that covers the target (red curve)

The algorithm samples from the simple Gaussian and accepts/rejects to match the more complex mixture distribution.

Example 2: Estimating π

Rejection sampling provides an elegant method for estimating π using geometric probability.

Setup:

• Proposal q(x): Square (side = 2r) | Uniform over square
• Target p(x): Circle (radius = r) | Uniform over circle

The acceptance rate directly relates to π:

Acceptance Rate = (Area of Circle) / (Area of Square)
 = (π × r²) / ((2r)²)
 = π / 4

which equates to:

π = 4 × Acceptance Rate

3) Importance Sampling

Imagine you want to understand population A, but you can only access population B. Importance sampling lets you learn about A by studying B, as long as you reweight your observations appropriately.

In statistical terms: “I can’t sample from the distribution I care about, so I sample from an easier one and correct for the difference.”

For e.g., You have a complex target distribution

f(x) = exp(-x²) × (1 + sin²(3x))

To compute the expectation, you need to normalize it:

where p(x):

  where Z = ∫ exp(-x²) × (1 + sin²(3x)) dx (from -∞ to +∞)

This integral has no closed-form solution!

Core Idea: Instead of sampling from the complex f(x), sample from something simple and re-weight.

The Algorithm

Step 1: Choose a Simple Proposal Distribution

q(x) = Normal(0, 1)

Step 2: Draw Samples from q(x)

Step 3: Compute Weights

For each sample, compute:
weight = f(x) / q(x)
 = (what target wants) / (what proposal gives)

Step 4: Estimate Using Weighted Average

E[x] ≈ Σ(weight × x) / Σ(weight)

3.1) Why the Normalizing Constant Z Cancels Out?? [Optional]

By definition, the expectation of any value or function g(x) under a distribution q(x) is:

In simple words, to find the average of “stuff” under distribution q(x), you integrate (“sum up”) stuff * q(x).

We start with the expectation under p(x):

We multiply and divide by q(x) (which is just multiplying by 1):

Now, look at the equation above and group the terms together to match the definition in step 1. Let our “stuff” be g(x) = x (p(x)/q(x)).

Since this matches the formstuff * q(x) it is exactly the expectation of “stuff” under q (We basically “tricked” the integral into looking like an expectation for q instead of p):

If we replace p(x) with f(x)/Z, the numerator becomes:

The Denominator : Next, we estimate the sum of probabilities (which should equal 1). Again, using f(x) instead of p(x) introduces the same factor:

And since we know:

Substituting and canceling the terms, we actually get:

4) Markov Chain Monte Carlo

When both rejection and importance sampling struggle particularly in high dimensions, Markov Chain Monte Carlo (MCMC) provides a solution. Instead of generating independent samples, MCMC builds a Markov chain that gradually converges to the target distribution, exploring even the most complex probability one step at a time (critical for bayesian inference).

Let’s say we want to recommend new users on our social media platform — Y (not X). Our platform has over 500M users. We want to recommend friends for you, but comparing you against everyone would take forever.

We can use MCMC to explore the social network smartly with random walk:

Start at a random user in the network.

Repeat thousands of times:

1. Look at a neighbor (friend of current user)
2. Calculate their match score (assume each potential friend gets a match score based on concrete criteria)
3. Decide whether to move:

• Better score? → Always move
• Worse score? → Move with probability (new score / current score)

If we only moved to better users, we might get stuck in one corner of the network. Occasionally accepting worse options lets us discover entirely new regions with potentially great matches.

Scoring System
`Mixture of 3 Gaussians representing different friend clusters
# - Cluster 1 (μ=20): Similar Age/Same Location (low match)
# - Cluster 2 (μ=60): Shared Interests (medium-high match)
# - Cluster 3 (μ=85): Mutual Friends (high match)

The Technical Details: The Metropolis-Hastings Algorithm

For those interested in the formal implementation, here is how we define the problem mathematically.

We want to sample from a probability distribution p(x) where:

• Direct sampling is impossible or expensive.
• We can only evaluate p(x) for any given x (up to a normalizing constant).
• The space of possible x values is enormous.

Goal: Generate samples x_1, x_2, x_3 etc. such that they are distributed according to p(x).

The Algorithm Steps

The most common MCMC method constructs a Markov chain whose stationary distribution is our target p(x).

1. Initialize Set x_0 to an arbitrary starting point.

2. Loop (for t = 1, 2, 3, …)

a. Propose Draw a candidate x from a proposal distribution q(x | x_{t-1}).

b. Calculate Acceptance Probability We determine the probability of accepting the new state:

c. Accept or Reject

• Draw u ~ Uniform(0, 1)
• If u < α: Set x_t = x’ (Accept the new position).
• Else: set x_t = x_{t-1} (reject) (Reject and stay at the current position).

The Metropolis Algorithm is essentially a strategy for exploration vs. exploitation.

• It exploits good results by always moving to better scores.
• It explores new possibilities by occasionally taking a risk on a lower score.

Conclusion

These three Monte Carlo sampling methods form the foundation of modern computational statistics and machine learning:

• Rejection Sampling provides geometric intuition and works well for simple problems
• Importance Sampling enables expectation computation without knowing normalization constants
• MCMC scales to complex, high-dimensional problems where other methods fail

References

1. Simulation code package — https://github.com/mandeep0405/causality/tree/main/distributions
2. Vizualization Python package — https://codeberg.org/ManimCommunity

Introduction

You are a Product Manager or Data Scientist. You’ve just run an A/B test on a new UI flow (Variant B) against the current one (Variant A). The dashboard says Variant B has a $2 higher Average Sales Per User.

Do you ship it?

The answer isn’t “Yes” because the number is higher. The answer depends on uncertainty. Is that $2 lift real, or is it just random noise? If you ship a bad feature, you lose money. If you fail to ship a good one, you leave money on the table.

Traditional A/B testing methods often obscure this decision-making process with confusing metrics. Today, we’re going to look at a better way: Bayesian A/B Testing.

Frequentist Approach — Traditional Method

Most industry standard tools use the Frequentist approach (Null Hypothesis Significance Testing). You set up a “Null Hypothesis” (typically: there is no difference between A and B) and try to reject it.

The Key Metrics:

• P-value: The probability of seeing data this extreme if the Null were true. (Note: This is NOT the probability that B is better than A).
• Power: The probability of detecting an effect if one actually exists.
• Confidence Interval: A range constructed from the data. Here’s what a 95% confidence interval actually means: “If we repeated this experiment 100 times and calculated the interval each time, about 95 of those intervals would contain the true value.” This is NOT the same as saying “there’s a 95% probability the true value is in this specific interval.” The true value either is or isn’t in your observed interval — you just don’t know which.

The Fundamental Limitation

Here’s the critical issue: The Frequentist approach tells you if a result is statistically significant, but it doesn’t tell you the probability of actually seeing these gains.

Let’s use a concrete example:

• Variant A (Control): 1,000,000 sales
• Variant B (Treatment): 1,200,000 sales
• Difference: 200,000 more sales
• P-value: 0.03 (< 0.05, so “significant”)

What the Frequentist method tells you:

“If there were truly no difference between A and B, the probability of seeing a difference this extreme (200,000 or more) is only 3%.”

In other words: p < 0.05 means “this difference is unlikely to be just random noise.”

What you actually want to know:

“If I deploy Variant B, how likely am I to actually see this 200,000 sales increase? What’s the probability I’ll get these results?”

The Pain Point: The “Peeking” Problem

The biggest flaw in Frequentist testing isn’t the math — it’s the rigidity.

• You must calculate a sample size before you start (e.g., “We need 10,000 users”).
• You cannot stop early.
• You cannot peek.

If you check your results halfway through, see a “significant” p-value, and stop the test, you have committed a statistical sin (p-hacking). Your False Positive rate just skyrocketed.

Why does this happen?

Each time you check your results, you’re essentially running a new hypothesis test. With a p-value threshold of 0.05, you have a 5% chance of a false positive on each check. If you peek 10 times during your experiment, your cumulative false positive rate can go over 5%.

In fig2, we expect 5% FPR if we don’t peek till the end (i.e., 0.05 alpha significance). But if we peek after every 100 samples, we expect FPR to go above 18%.

3. The Bayesian Shift

Bayesian statistics flips the script. Instead of asking “How rare is this data given the Null?”, it asks the question we actually care about: “Given this data, what is the probability that B is better than A?”

It works by updating our beliefs as we get more evidence.

The Core Equation

Posterior∝Likelihood×Prior

• Prior: What we believed before seeing data.
• Likelihood: What the data tells us.
• Posterior: Our updated belief.

3.1) Computing the Posterior: The Normal-Normal Conjugate

How do we actually do the math? We use a “Conjugate Prior.” If our Prior is a Normal distribution and our Data is Normal, the Posterior is also mathematically guaranteed to be a Normal distribution.

The Intuitive Formula:

The new belief (Posterior Mean) is just a weighted average of your Prior Mean and your Data Mean.

How is “Certainty” measured?

Certainty is measured by precision, which is the inverse of variance:

What this means:

• If you have little data (small n), the Prior pulls the weight (Caution).
• If you have lots of data (large n), the Data pulls the weight (Evidence).

3.2 Choosing Your Prior

There are two main approaches:

Uninformative Priors (When You Know Nothing)

• Use a very flat, wide distribution that minimizes influence on the posterior
• Common approach: Use a power analysis with an assumed effect size of 0.5 (medium effect) to set the variance
• Says “I have no strong beliefs, let the data speak”
• Lets the data dominate quickly

Informative Priors (When You Have History)

• Use historical data to set the prior only when your success metric is reliable and stable
• Examples of reliable metrics: ARPU from last quarter, conversion rates over the past 6 months, page load times

4. Advantages of Bayesian Approach

4.1) Intuitive Interpretation

• Frequentist: “There is a 5% chance of observing a difference this large if there were actually no difference.” (Mental gymnastics required).
• Bayesian: “There is a 97% chance Variant B is better than Variant A.” (Direct answer).

Posterior Probability of Positive Return (PPR)

One of the most powerful Bayesian metrics is PPR — the probability that your variant produces a positive lift [1].

Definition: PPR = Probability(B — A > 0)

This is calculated by looking at the posterior distribution of the difference between variants and finding the area under the curve where the difference is positive.

Example: You run an A/B test on Total Sales($$). After collecting data:

• Variant A: posterior Mean = $50000, StdDev = $15000
• Variant B: posterior Mean = $52000, StdDev = $8000
• The difference (B — A) has distribution: Mean = $1,815

PPR = 81% (the probability that B truly generates more revenue than A)

Please refer to [3] for running the Bayesian AB Simulation.

Unlike p-values, PPR directly answers: “How confident am I that this variant is better?” This is the metric business stakeholders intuitively understand.

Credible Interval (Bayesian): [48, 54]

• Interpretation: “Given the data I observed, there is a 95% probability that the true value lies between 48 and 54.”
• This IS a probability statement about where the true value is.
• This is what most people think a confidence interval means!

Bayesian credible intervals actually mean what we want them to mean: a direct probability statement about where the true value lies.

4.2. Small Sample Efficiency

Bayesian methods excel when data is scarce. By incorporating a Prior (even a weak one), you prevent wild outliers from misleading you early in the test.

4.3) Continuous Monitoring

This is the game-changer. In Bayesian testing, “peeking” is valid.

You can check the probability every morning.

• Day 1: 60% chance B is better. (Keep running).
• Day 2: 80% chance B is better. (Keep running).
• Day 3: 99% chance B is better. (Stop and Ship!)

You don’t need to commit to a fixed sample size upfront. You stop when you have enough evidence to make a decision.

References

1. Cumulative Density Function
2. Conjugate Priors: A Conjugate Prior happens when your Prior and your Posterior end up being in the same family of distributions.
3. Bayesian AB Testing Simulation Code

1) Introduction

Multi-head attention has been the cornerstone of transformer architectures since “Attention is All You Need” [1]. Yet despite its success, standard softmax attention harbors subtle but significant inefficiencies that become pronounced at scale. Addressing these fundamental limitations is the focus of “Gated Attention for Large Language Models,” a breakthrough work by the Qwen Team that was recently recognized as a Best Paper at NeurIPS 2025[4].

1.1) Problem 1: The Low-Rank Bottleneck

Imagine you have a beautiful, massive 4K photograph. This represents your input x with 2048 dimensions of rich detail.

Step 1 (x .W_V): You force this 4K image through a compressor that shrinks it down to a tiny 64-pixel thumbnail.

• Math equivalent: Projecting from 2048 dimensions down to 64 dimensions.

Step 2 (W_O): You try to blow that 64-pixel thumbnail back up to a 4K resolution image.

• Math equivalent: Projecting from 64 dimensions back up to 2048 dimensions.

The Bottleneck: Even though your final output is technically a “4K image” (it has 2048x2048 pixels), it is blurry and blocky. You cannot magically recover the details of the original tree leaves or skin textures because they were discarded when you shrank it to 64 pixels.

The choice of low-rank heads was a deliberate trade-off to enable computational efficiency and parallel specialization, effectively creating a team of “experts” i.e., attention heads where each focuses on a specific feature subset (like grammar or tone) without exploding costs.

Step-by-step transformation for “attention”:

For a model with:

• Model dimension: d_model = 2048
• Number of heads: h = 32
• Head dimension: d_k = 64 (typically d_model / h)

Input token representation:        x ∈ ℝ^2048

Step 1 - Value projection:
v = x · W_V^i                      Result: v ∈ ℝ^64
    ↑        ↑
  (2048)  (2048×64)

Step 2 - Attention weighting:
attention_output = attention_weights · v   Result: still ℝ^64
                   (This is the SDPA output)

Step 3 - Output projection:  
o = attention_output · W_O^i       Result: o ∈ ℝ^2048
    ↑                   ↑
   (64)            (64×2048)

This is where bottleneck occurs:

 Combined transformation (for one head):
o = attention_weights · (x · W_V^i) · W_O^i
  = attention_weights · x · (W_V^i · W_O^i)
  = attention_weights · x · W_merged^i

Where:
W_merged^i = W_V^i · W_O^i
           = (2048×64) · (64×2048)  
           = 2048×2048 matrix

BUT: rank(W_merged^i) ≤ min(rank(W_V^i), rank(W_O^i))
                      ≤ 64  ← Bottleneck!

1.2) Problem 2: The Attention Sink Phenomenon

When processing long sequences, transformer models develop a peculiar behavior: they allocate disproportionate attention to early tokens, particularly the first token “[BOS]”, even when these tokens are irrelevant to the current prediction.

Input: “In 1887, construction of the Eiffel Tower began… [500 words of history]…
Today, the tower is located in.”

When predicting “Paris”:

Why Does This Happen? The Softmax Constraint Problem

Softmax attention has two fundamental constraints:

1. All weights must be non-negative (due to exponential in softmax)
2. All weights must sum to 1.0 (due to normalization)

When most context is irrelevant, the model would ideally assign near-zero or negative attention to those tokens. But softmax doesn’t allow this.

Why the First Token Becomes the Sink

The first token (often the Beginning of Sentence, or BOS) becomes the designated “attention sink” primarily because it is semantically empty, making it a “safe” dumping ground for excess attention weights that would otherwise corrupt the representation of meaningful tokens. Additionally, due to the causal masking in transformer architectures, the first token is universally visible to every subsequent token in the sequence, making it the most reliable and consistent focal point for the model to discover and utilize during training (a learned optimization) to satisfy the softmax constraint that all attention weights must sum to 1.

2) How Gated Attention Solves These Problems

Recent work from Alibaba’s Qwen Team introduces a deceptively simple modification: add a learnable sigmoid gate after the attention computation.

Step 1: Compute attention weights (same)
        attention_weights = softmax(Q·K^T / √d_k)
        Shape: (seq_len, seq_len)

Step 2: Apply attention to values (SDPA output) (same)
        sdpa_output = attention_weights @ V
        Shape: (seq_len, d_k)

Step 2.5: ← NEW! Compute gate and apply it
        gate_scores = sigmoid(X @ W_gate)
        Shape: (seq_len, d_k)
        
        gated_sdpa_output = sdpa_output ⊙ gate_scores  ← Element-wise multiplication
        Shape: (seq_len, d_k)

Step 3: (After concatenating all heads) Project with W_o
        output = gated_sdpa_output @ W_o
        Shape: (seq_len, d_model)

At a high level, this is a model architecture with the new “gated attention” mechanism:

2.1) Breaking the Low-Rank Bottleneck with Non-Linearity

The sigmoid gate cannot be factored out or merged with the linear projections. This breaks the low-rank limitation and increases expressiveness.

Without gating: the “basis vectors” (the directions in space the model can see) are fixed. It’s like having a flashlight that is welded to a table; it can only ever illuminate the exact same spot on the wall.

output = (attention_weights·V)·W_o
       = attention_weights·(V·W_o)    ← Two linear layers collapse
       = attention_weights·W_merged   ← Low-rank bottleneck

With gating: The gate is calculated from the input itself. This means for every different input sentence, the gate changes.

• Input A: The flashlight rotates left (illuminates “grammar”).
• Input B: The flashlight rotates right (illuminates “tone”).

Even though the beam is still narrow (Rank 64), because it can move (rotate) based on what it’s looking at, the total area it can cover over time is the entire room (Rank 2048).

output = (attention_weights·V ⊙ sigmoid(X·W_gate))·W_o
                              ↑
                        Non-linear gate!

2.2) Eliminating Attention Sinks with Sparse, Query-Dependent Gating

Gated Attention solves this by adding a Sigmoid Gate after the Softmax output.

• Decoupling: While Softmax forces a sum of 1.0, the Sigmoid gate operates element-wise with a range of [0, 1].
• Filtering: The gate is sparse (mostly near 0) and query-dependent. It can recognize that the “dumped” attention on the first token is irrelevant to the current query.
• The Fix: Even if Softmax assigns a high score (e.g., 45%) to the first token, the gate can assign it a near-zero weight (e.g., 0.01), effectively suppressing the noise. This allows the model to selectively “turn off” the attention sink and focus only on relevant tokens, bypassing the Softmax sum-to-one constraint at the output level.

Real Example: Document Summarization

Consider this scenario:

[BOS] The history of artificial intelligence dates to ancient times…
[480 tokens of historical background]…
In conclusion, modern AI impacts three key areas: healthcare, education,
and climate science.

3) Demonstration

I have included dummy implementation of “Gated Attention” for demonstration purpose. Please refer to this for more details:

https://github.com/mandeep0405/llm-tutorials/tree/main/gated_attention

from gated_attention import GatedAttentionHead, MultiHeadGatedAttention

# Single-head gated attention
d_model = 512
d_k = 64
gated_head = GatedAttentionHead(d_model, d_k)

# Forward pass
x = torch.randn(batch_size, seq_len, d_model)
output = gated_head(x)

# Multi-head gated attention
num_heads = 8
gated_mha = MultiHeadGatedAttention(d_model, num_heads)
output = gated_mha(x)

# With gate statistics
output, gate_stats = gated_mha(x, return_gate_stats=True)
print(f"Mean gate score: {gate_stats['mean']:.4f}")
print(f"Gate sparsity: {gate_stats['sparsity']*100:.1f}%")

Attention Distribution for Last Token — Standard Vs Gated Attention

• Standard Attention (Left): The tall bar at position 0 (BOS token) vividly shows how a standard model “dumps” excess attention (approx 4.0%) onto the first token, even though it’s likely irrelevant to the last token.
• Gated Attention (Right): The corresponding bar at position 0 is significantly shorter (approx 2.1%) with gated attention.

Gate Score/Final Attention Weight Distribution by Query Position

• Standard Attention Weights (Top Left): Shows the typical triangular causal mask pattern.
• Mean Gate Score per Query Position (Top Right): Demonstrating the query-dependent nature of the mechanism. The orange shaded area represents the “allowed” signal, while the white space above represents suppression. Note, if trained (this is random), as according to the paper, it goes to 0.1–0.2 range.
• Attention to Token 0 [BOS] (Bottom Left): This line graph tracks the attention sink over time. The red line (Standard) starts very high (1.0) and remains consistently higher than the blue line (Gated).
• Gate Scores Heatmap (Bottom Right): This reveals the sparsity and granularity. The heatmap is “patchy,” meaning the gate is selectively turning features on and off for specific tokens.

Conclusions

While the theoretical elegance of breaking low-rank bottlenecks and removing attention sinks is compelling, the real verdict comes from the benchmarks. Authors performed testing on LLMs trained on 3.5 trillion tokens to prove that Gated Attention is not just a stability fix — it is a significant performance upgrade. The gated models achieved a perplexity reduction of over 0.2 and downstream performance surged, with MMLU scores increasing by ~2.0 points and GSM8k math reasoning scores jumping by over 2.3 points compared to the baseline.

References

1. Attention is all you need
2. Gated Attention
3. Implementation Code
4. https://neurips.cc/virtual/2025/awards_detail

Introduction

You’ve probably used PyTorch hundreds of times — writing output.backward() without a second thought, broadcasting arrays like magic, and watching gradients computed itself across your training loops. But have you ever wondered what's actually happening under the hood? This tutorial peels back the curtain on PyTorch's internals, walking through how tensors store data in memory, how automatic differentiation builds and traverses computational graphs, and why gradient accumulation works the way it does!!!

1) Basics about Tensors?

At the bedrock of the PyTorch library lies the torch.Tensor. Unlike a standard array, it is a view over a storage block engineered for deep learning. Today, a torch.Tensor wears three hats simultaneously:

1. The View (Metadata): Describes how to look at the data (shape, stride).
2. The Storage (Data): The actual raw memory pointer.
3. The Autograd Node (Graph): Pointers to the computational graph.

1.1) Memory Layout & Strides

A tensor is essentially a view over a storage block — a contiguous chunk of memory holding the raw numerical data. The stride attribute defines how to traverse this memory. Strides describe the number of steps needed in the flat storage array to move one index along a dimension.

import torch

# Create a 2D tensor (3 rows, 4 columns)
x = torch.arange(12).reshape(3, 4)
print(f"Shape: {x.shape}")   # torch.Size()
print(f"Stride: {x.stride()}") # (4, 1)
# Explanation: To move down one row, skip 4 elements. To move right one column, skip 1.

# Transpose the tensor
y = x.t()
print(f"Shape: {y.shape}")   # torch.Size([4, 3])
print(f"Stride: {y.stride()}") # (1, 4)
# Explanation: The shapes are swapped, but crucially, the strides are swapped too.
# PyTorch did NOT copy data. It just changed how it reads memory.

Memory Efficiency: Operations like transpose, slice, and view can create new tensor “views” without copying data. They just change the strides!

x = torch.randn(4, 4)
x_t = x.t()  # Almost instant, no copy!
x_reorder = x.view(2,8) # no copy!

1.2) The Mechanics of Broadcasting

Broadcasting is a mechanism that allows PyTorch to perform operations on tensors of different shapes without explicitly copying data. It automatically expands smaller tensors to match the shape of larger ones during arithmetic operations.

PyTorch checks dimensions from right to left, and each pair of dimensions must either:

• Be the same number, OR
• One of them is 1 (which can stretch)

# Example 1: Scalar with tensor
a = torch.tensor([1, 2, 3])
b = 5
result = a + b  # [6, 7, 8]
# 5 broadcasts to [5, 5, 5]

# Example 2: 1D with 2D
a = torch.tensor([[1, 2, 3],
                  [4, 5, 6]])  # shape: (2, 3)
b = torch.tensor([10, 20, 30])  # shape: (3,)
result = a + b
# b broadcasts to [[10, 20, 30],
#                  [10, 20, 30]]
# Result: [[11, 22, 33],
#          [14, 25, 36]]

# Example 3: Column vector with row vector
a = torch.tensor([[1],
                  [2],
                  [3]])  # shape: (3, 1)
b = torch.tensor([[10, 20, 30]])  # shape: (1, 3)
result = a + b  # shape: (3, 3)
# Result: [[11, 21, 31],
#          [12, 22, 32],
#          [13, 23, 33]]

Illusion of Copying Data: The interesting thing is that it never copies the data when it try to stretch the smaller tensor. It just reads the same data multiple times!!!

A = torch.randn(3, 4)
B = torch.tensor([[999]])

# Method 1: Direct broadcasting (no intermediate tensor) Broadcasts B (1,1) -> (3,4)
result1 = A + B

# Method 2: Manual expand (creates a view with zero strides)
# Pytorch knows that it needs to read the same data multiple times
result2 = A + B.expand(3, 4)

# Method 3: Full copy (wasteful) This doesn't happens!!
result3 = A + B.expand(3, 4).clone()

print(torch.equal(result1, result2))  # True
print(torch.equal(result1, result3))  # True

# But check what tensors exist in memory:
print(B.data_ptr())                    # Address X (1174703488)
print(B.expand(3, 4).data_ptr())       # Address X (1174703488)
print(B.expand(3, 4).clone().data_ptr()) # Address Y (1174437120)

2) Magic of AutoGrad

AutoGrad is PyTorch’s automatic differentiation engine. It’s the reason you don’t have to manually compute gradients for back-propagation.

Core Concept: As your forward pass executes, PyTorch builds a dynamic “tape” (computational graph) recording every operation. During backward(), it plays this tape in reverse to compute gradients.

The key insight is that PyTorch doesn’t compute derivatives symbolically at runtime. Instead, it uses a clever two-phase approach:

Phase 1: Forward Pass — Building the Computational Graph

During the forward pass, PyTorch does two things simultaneously:

1. Computes the result (the actual forward calculation)
2. Records what operation was used and stores a reference to the appropriate backward function

Here’s what happens under the hood when you write y = x ** 2:

# What you write:
y = x ** 2

# What PyTorch actually does:
result_data = x.data ** 2  # Compute the result
grad_fn = PowBackward(input=x, exponent=2)  # Store backward operation
y = Tensor(data=result_data, grad_fn=grad_fn)  # New tensor with link

The crucial part: PowBackward is not computing anything yet. It's just storing:

• A reference to the input tensor x
• The exponent value 2
• A pre-written backward() method that knows the power rule

Phase 2: Backward Pass — Applying Pre-Computed Formulas

When you call .backward(), PyTorch traverses the computational graph in reverse, and each operation applies its pre-computed gradient formula.

For the power operation, the gradient formula is hard-coded (Note: Actually it’s written in C++ and not in python — this is only for demonstration):

class PowBackward:
    """Backward pass for power: z = a ** n

    Gradient:
        dL/da = dL/dz * dz/da = dL/dz * n * a^(n-1)
    """

    def __init__(self, input_tensor: 'Tensor', exponent: float) -> None:
        self.input_tensor = input_tensor
        self.exponent = exponent
        self.name = "PowBackward"

    def backward(self, grad_output: NDArray) -> None:
        # Power rule: multiply by exponent and reduce power by 1
        if self.input_tensor.requires_grad:
            grad = grad_output * self.exponent * (self.input_tensor.data ** (self.exponent - 1))
            self.input_tensor.backward(grad)

The derivative formula n * x^(n-1) is already known and hard-coded into the PowBackward class. PyTorch isn't deriving this formula—it's applying it!

Let’s trace through a concrete example to see this in action.

x = Tensor(2.0, requires_grad=True)
w1 = Tensor(3.0, requires_grad=True)
w2 = Tensor(4.0, requires_grad=True)
b = Tensor(1.0, requires_grad=True)

# Forward pass: builds the graph
output = (x * w1 + x * w2 + b) ** 2

Computational Graph:

Each node stores its specific backward operation:

• [×] nodes → MulBackward
• [+] nodes → AddBackward
• [^2] node → PowBackward

Backward Pass:

output.backward()  # Start with grad = 1.0

# PowBackward at output:
#   ∂output/∂t3 = 2 * t3^1 = 2 * 15 = 30
#   Pass grad=30.0 to t3

# AddBackward at t3:
#   ∂t3/∂t2 = 1, ∂t3/∂b = 1
#   Pass grad=30.0 to both t2 and b
#   → b.grad = 30.0

# AddBackward at t2:
#   ∂t2/∂t1 = 1, ∂t2/∂t4 = 1
#   Pass grad=30.0 to both t1 and t4

# MulBackward at t1:
#   ∂t1/∂x = w1 = 3, ∂t1/∂w1 = x = 2
#   Pass grad=30*3=90 to x, grad=30*2=60 to w1
#   → x.grad += 90, w1.grad = 60

# MulBackward at t4:
#   ∂t4/∂x = w2 = 4, ∂t4/∂w2 = x = 2
#   Pass grad=30*4=120 to x, grad=30*2=60 to w2
#   → x.grad += 120, w2.grad = 60

# Final gradients:
# x.grad = 90 + 120 = 210.0  (gradient accumulated!)
# w1.grad = 60.0
# w2.grad = 60.0
# b.grad = 30.0

For demonstration purpose, I implemented key differentiation methods in python. Please refer [3] for more details.

3) Gradient Accumulation

One of the most important — and often misunderstood — aspects of PyTorch’s autograd is gradient accumulation. Understanding when we should accumulate gradients is crucial for correct training.

1. If your neural architecture has multiple losses, then we are actually doing gradient accumulation. They all backpropagate through the same neural network weights

model: Input → [Layer1: W1, b1] → [Layer2: W2, b2] → Output

# Multiple losses
loss1_classification = CrossEntropyLoss(output, target)
loss2_regularization = L2Penalty(W1, W2)  
loss3_auxiliary = AuxiliaryTask(hidden_layer)
# Combine into total loss
total_loss = loss1 + loss2 + loss3

# Single backward pass
total_loss.backward()  

Under the hood, PyTorch is doing this:

# Initialize gradients to zero
optimizer.zero_grad()

# Compute each loss
loss1 = classification_loss(output, target)
loss2 = regularization_loss(W1, W2)
loss3 = auxiliary_loss(hidden)

# Call backward on each loss separately
loss1.backward()          # w.grad = 5
loss2.backward()          # w.grad = 5 + 0.04 = 5.04
loss3.backward()          # w.grad = 5.04 + 2 = 7.04

# Update weights using accumulated gradient
optimizer.step()          # Update w using 7.04

2. If we need a larger batch size (32) but due to resource constraint we can only support smaller size (8). In this case, we will accumulate gradient for batches.

# Simulate batch size of 32 with only 8 at a time
accumulation_steps = 4
optimizer.zero_grad()

for i, (X_batch, y_batch) in enumerate(dataloader):
    # Forward pass
    output = model(X_batch)
    loss = criterion(output, y_batch)
    
    # Backward pass (accumulates gradients)
    loss.backward()
    
    # Update only every 4 batches
    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad() 

References

1. https://docs.pytorch.org/tutorials/beginner/basics/intro.html
2. https://sebastianraschka.com/teaching/pytorch-1h/
3. https://github.com/mandeep0405/pytorch-autograd-tutorial/blob/main/README.md