amandeep singh

The Boon of Dimensionality

Tue, 03 Mar 2026 00:00:00 GMT

I recently watched Grant Sanderson's (3blue1brown) video about volume of high-dimensional spheres. He made a note that the high-dimensional space is also of peculiar interest to the ML field. I knew of the curse of dimensionality, but I wanted to dig deeper and here is what I found, the other side; the boon of dimensionality.

I would recommend you watch Grant's video for this. I will try my best here to explain the intuition behind the presented results.

Volume of a high-dimensional ball

The result that Grant's video centers on is that the volume of a ball with radius r in d dimensions is:

V_d = \cfrac{\pi^{d/2}}{\Gamma(d/2 + 1)}{r^d}

Where \Gamma(x) is the Gamma function. Plotting V_d for d from 1 to 50 we get something like this.

You can see the volume peaks at around d=5, roughly 5.26 and at d=50 becomes vanishingly small, since the Gamma function denominator grows much faster than the numerator. Even more strange is how this volume is distributed. Consider a ball of radius (1-\epsilon) where epsilon is an infinitesimal (read: vanishingly small value). The ratio of its volume to the full ball is

\cfrac{V(1-\epsilon)}{V(1)} = {(1-\epsilon)}^d \leq {e}^{-\epsilon d}

Since (1-x) \leq e^{-x} from Taylor expansion of e^{-x}. We can see that as d \rightarrow \infty, the ratio \rightarrow 0. This tells us that most of the already small volume is concentrated near its surface. This is a special case of a general principle called concentration of measure, the tendency for high-dimensional probability distributions to concentrate their mass in thin regions.

Equators

Yet another interesting thing is about the equators, but first what is the equator of a high-dimensional ball. we can pick a coordinate say x_1 and all the points with -1 \leq x_1 < 0 lie in one hemisphere and 0 < x_1 \leq 1 in another; the x_1 = 0 boundary will be the equator. The other such slices will be \sqrt{1 - {x_1}^2} (Pythagoras) and volume of such a slice will be \propto {(1-{x_1}^2)}^{(d-1)/2} , since we have fixed one dimension and are left with d-1 dimensions. This is again bounded by

{(1-{x_1}^2)}^{(d-1)/2} \leq e^{-(d-1){x_1}^{2}/2}

This shows that the slice volume decreases exponentially as we move away from 0, which means most of the volume is also concentrated near x_1 = 0. From the shell result and this we can see most of the points are concentrated on the equator.

Near Orthogonality

The value x_1 is just \langle x, e_1 \rangle, where e_1 is the unit vector across the chosen axis. If most points lie near the equator (small x_1), they are nearly orthogonal to e_1. Since balls are rotationally symmetric, this can be said about any direction. So pick one point, make it the north pole and we are very certain the next point we pick will lie on the equator.

The Johnson-Lindenstrauss Lemma

The near-orthogonality result has a famous companion that makes the capacity claim precise. In 1984, Johnson and Lindenstrauss proved the following ¹:

Take any n points in a high-dimensional space. You can project them into a space of just O(\epsilon^{-2} \log n) dimensions and preserve all pairwise distances up to a small factor (1 \pm \epsilon).

While this is originally stated as a compression result, we can read it backwards to see that a d dimensional space has the capacity to represent e^{\Omega(d)} points with geometry intact. Note that this is exponential in d.

This is not just a theoretical curiosity. Embedding models like word2vec ² or sentence transformers ³ pack millions of concepts into ~768 dimensions. The JL capacity result tells us why this works: even after accounting for the constant hidden in the \Omega, a 768-dimensional space has room for a huge number of near-orthogonal directions, far more than any vocabulary needs. Unrelated words get mapped to nearly orthogonal vectors and don't interfere with each other, which is why cosine similarity between embeddings tracks semantic similarity so well. The famous analogy arithmetic (king - man + woman ≈ queen) suggests that directions in this space can be meaningful, though in practice this only works for cherry-picked examples ⁴. What matters more is the geometric separation. Random projections exploit this directly too: you can project high-dimensional data down to O(\log n) dims via a random matrix and preserve distances, which is the basis of locality-sensitive hashing and streaming sketches.

Where next

The curse of dimensionality tells us that data gets sparse and distances lose meaning as dimensions grow. The boon is the flip side: exponential capacity, near-orthogonality for free, and distances that survive projection. Same geometry, two readings. These two operate at different scales though: the curse applies to the intrinsic dimensionality of the data (you still need samples proportional to the manifold dimension), while the boon applies to the ambient dimensionality (the room you have to embed things in). ML lives in the gap since real data is high-dimensional but structured (it sits on a manifold), and the boon wins over the curse.

There are two directions I want to explore from here. First is Cover's theorem (1965), which says data that is not linearly separable in low dimensions becomes separable when mapped to higher dimensions. This is why kernel methods and neural network hidden layers work, they buy room. Second is superposition in neural networks, the idea that a network with m neurons can represent far more than m features by packing them into near-orthogonal directions, which connects directly to the geometry above. But those are for another post.

Footnotes

W. B. Johnson and J. Lindenstrauss, "Extensions of Lipschitz mappings into a Hilbert space," Contemporary Mathematics, 26, 1984. ↩
T. Mikolov et al., "Efficient Estimation of Word Representations in Vector Space," arXiv:1301.3781, 2013. ↩
N. Reimers and I. Gurevych, "Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks," arXiv:1908.10084, 2019. ↩
O. Levy and Y. Goldberg, "Linguistic Regularities in Sparse and Explicit Word Representations," CoNLL 2014. See also Nissim et al., "Fair is Better than Sensational," arXiv:1905.09866, 2019. ↩

Twelve Attempts at an FP4 Kernel

Sat, 28 Feb 2026 00:00:00 GMT

I recently participated in GPU Mode's NVFP4 Kernel Hackathon. The goal was to write CUDA kernels for Blackwell's new 4-bit floating point format (NVFP4) and get as close as possible to the hardware's speed of light. The competition ran on B200 GPUs, and there were four tasks in total: matrix-vector multiplication (GEMV), matrix-matrix multiplication (GEMM), gated GEMM (with SiLU activation), and grouped GEMM. All tasks use block-scaled FP4 inputs where every 16 elements of e2m1 data share a single e4m3 FP8 scale factor, with fp16 output. I was able to complete the GEMV and GEMM tasks but ran out of time for the gated and grouped variants.

If you want background on CuTe's layout algebra (which underpins a lot of what follows), I wrote about it here.

The Task

For GEMV we were given a matrix A of shape M \times K \times L and a vector B of shape 1 \times K \times L (both in packed FP4 with FP8 scale factors), compute C = A \cdot B^T of shape M \times 1 \times L in FP16. The ranking metric was the geometric mean of benchmark times, measured against a speed-of-light analysis based on max(FFMA math throughput, DRAM memory throughput) of the B200 at 1.5GHz clock:

M	K	L	Speed of Light (µs)
7168	16384	1	8.6
4096	7168	8	17.3
7168	2048	4	4.3

Optimization Journey

This is where I spent most of my time. Twelve attempts, of which only one really worked well.

CuTe Python DSL (Attempts 1-4)

My first instinct was to use CUTLASS's Python DSL since I had been experimenting with CuTe's layout system already. The initial approach was straightforward. One thread per output row element, iterating over K, decoding FP4 to FP16, multiplying by scale factors, and accumulating.

@cute.kernel
def _kernel(self, a, b, sfa, sfb, c):
    bidx, bidy, bidz = cute.arch.block_idx()
    tidx, _, _ = cute.arch.thread_idx()

    # ... tensor setup, local_tile ...

    global_m = bidx * self.b_m + tidx
    if global_m < M:
        tCrC = cute.zeros_like(tCgC, cutlass.Float32)

        for k in range(k_tiles):
            # Load FP4 values, decode to FP16
            a_val = tAgA[tidx, None, k].load().to(cutlass.Float16)
            b_val = tBgB[0, None, k].load().to(cutlass.Float16)
            sfa_val = tAgSFA[tidx, None, k].load().to(cutlass.Float32)
            sfb_val = tBgSFB[0, None, k].load().to(cutlass.Float32)

            for i in cutlass.range_constexpr(self.b_k):
                tCrC += (a_val[i] * b_val[i]) * (sfa_val[i] * sfb_val[i])

        tCgC.store(tCrC.to(cute.Float16))

This produced correct results but the performance was not competitive. Each thread walks the entire K dimension sequentially with no parallelism along K. Over the next few attempts I experimented with K-dimension tiling and thread configuration but stayed within the same basic structure.

Split-K with Atomics in CuTe (Attempt 5)

To parallelize the K reduction I split the K dimension across threads and used atomicAdd to accumulate partial sums. This required defining custom @dsl_user_op functions for operations like atomic_add_fp32 that CuTe doesn't natively expose:

@dsl_user_op
def atomic_add_fp32(a, gmem_ptr, *, loc=None, ip=None):
    nvvm.atomicrmw(
        res=T.f32(), op=nvvm.AtomicOpKind.FADD,
        ptr=gmem_ptr.llvm_ptr, a=Float32(a).ir_value()
    )

The atomics were too expensive. The next attempt (attempt 6) replaced them with warp-shuffle reductions. 128 threads split into 4 warps, each warp handling one M row with 32-lane K tile splitting. This was better, but I was hitting a wall on the compute side.

The inner loop needs packed FP16 fused multiply-add (PTX: fma.rn.f16x2) to decode FP4 pairs, scale them, and accumulate efficiently. The CuTe Python DSL exposes fma_packed_f32x2 but has no equivalent fma_packed_f16x2 wrapper. Scalar FP16 arithmetic works through MLIR's arith dialect (arith.mulf, arith.addf), but whether the backend fuses a * b + c into an actual fma.rn.f16x2 is up to LLVM optimization passes, not something you can control or rely on. I tried writing llvm.inline_asm wrappers for the half2 operations but couldn't get them to work, and moved on to C++.

In hindsight this was a skill gap, not a DSL limitation. Another competitor placed in the top 10 (21.6µs) using a pure CuTe kernel that emits the entire decode-FMA-reduce pipeline as a single llvm.inline_asm block with cvt.rn.f16x2.e2m1x2 and fma.rn.f16x2, exactly what I was trying to do. The inlining is straightforward once you get the constraint strings right. I just didn't push through it at the time.

The Switch to Raw CUDA (Attempt 7)

I scrapped the CuTe approach and rewrote from scratch in raw CUDA C++, loaded via torch.utils.cpp_extension.load_inline.

The core idea is to assign 32 threads (one warp) per output row. Each thread handles a strided slice of the K dimension, accumulates locally, then the warp reduces via __shfl_down_sync. I used 128 threads per block with 4 rows per block.

For FP4 decoding I used Blackwell's new intrinsics. Each byte packs two FP4 values, and __nv_cvt_fp4x2_to_halfraw2 converts them to a half2 pair in a single instruction¹:

__device__ __forceinline__ __half2 decode_fp4x2(uint8_t byte) {
    __half2_raw raw = __nv_cvt_fp4x2_to_halfraw2(
        static_cast<__nv_fp4x2_storage_t>(byte), __NV_E2M1
    );
    return *reinterpret_cast<__half2*>(&raw);
}

The key function is the packed half2 dot product over a 4-byte (8 FP4 element) chunk. This is where the real speedup came from. The rest of the kernel structure (strided K loop, warp reduction) is fairly standard, but doing the decode-scale-multiply pipeline entirely in paired half2 operations avoids the scalar FP32 overhead of the CuTe attempts:

__device__ __forceinline__ __half2 dot_scaled_4bytes(
    uchar4 a4, uchar4 b4, __half2 scale_h2
) {
    __half2 acc0 = __hmul2(decode_fp4x2(a4.x),
                           __hmul2(decode_fp4x2(b4.x), scale_h2));
    __half2 acc1 = __hmul2(decode_fp4x2(a4.y),
                           __hmul2(decode_fp4x2(b4.y), scale_h2));
    acc0 = __hfma2(decode_fp4x2(a4.z),
                   __hmul2(decode_fp4x2(b4.z), scale_h2), acc0);
    acc1 = __hfma2(decode_fp4x2(a4.w),
                   __hmul2(decode_fp4x2(b4.w), scale_h2), acc1);
    return __hadd2(acc0, acc1);
}

Each scale factor group covers 16 FP4 elements = 8 bytes. The outer loop loads two uchar4 reads per group (4 bytes each, 8 FP4 values per load), calls dot_scaled_4bytes on each, and accumulates into FP32:

#pragma unroll 4
for (int sf = tid; sf < K_sf; sf += THREADS_PER_ROW) {
    float scale = decode_fp8(__ldg(&row_sfa[sf])) *
                  decode_fp8(__ldg(&batch_sfb[sf]));
    __half2 scale_h2 = __halves2half2(__float2half(scale),
                                       __float2half(scale));

    int byte_base = sf << 3;
    uchar4 a4_0 = *reinterpret_cast<const uchar4*>(&row_a[byte_base]);
    uchar4 b4_0 = *reinterpret_cast<const uchar4*>(&batch_b[byte_base]);
    uchar4 a4_1 = *reinterpret_cast<const uchar4*>(&row_a[byte_base + 4]);
    uchar4 b4_1 = *reinterpret_cast<const uchar4*>(&batch_b[byte_base + 4]);

    __half2 r0 = dot_scaled_4bytes(a4_0, b4_0, scale_h2);
    __half2 r1 = dot_scaled_4bytes(a4_1, b4_1, scale_h2);
    __half2 sum = __hadd2(r0, r1);
    float2 f = __half22float2(sum);
    acc += f.x + f.y;
}

After accumulation, the warp reduces with shuffles:

row_sum += __shfl_down_sync(0xffffffff, row_sum, 16);
row_sum += __shfl_down_sync(0xffffffff, row_sum, 8);
row_sum += __shfl_down_sync(0xffffffff, row_sum, 4);
row_sum += __shfl_down_sync(0xffffffff, row_sum, 2);
row_sum += __shfl_down_sync(0xffffffff, row_sum, 1);

if (lane == 0) {
    c[c_idx] = __float2half(row_sum);
}

Results:

M	K	L	My kernel (µs)	Speed of Light (µs)	Ratio
7168	16384	1	26.7	8.6	3.1x
4096	7168	8	45.1	17.3	2.6x
7168	2048	4	16.4	4.3	3.8x

Roughly 3x off speed of light. This became my baseline.

The Failed Experiments

Over the next five attempts, I tried every optimization I could think of. All of them made things worse or had zero effect.

In attempt 8, I went back to split-K with atomics, this time in C++. I tiled the K dimension across blocks and used atomicAdd to accumulate partial sums into an FP32 intermediate buffer with a final FP32-to-FP16 conversion pass. The atomic contention and extra memory traffic outweighed any parallelism benefit. The kernel is memory-bound, and adding more blocks just means more scheduling overhead and more atomic serialization at the same memory addresses.

Attempt 9 tried wider vectorized loads. Instead of two uchar4 loads (32 bits each), I used a single uint2 load (64 bits). This was 16-25% slower. Extracting individual bytes from a uint2 requires bitwise operations (& 0xFF, >> 8, etc.) and make_uchar4() calls. The compiler already optimizes two consecutive uchar4 loads into efficient memory transactions. I was adding instruction overhead to "save" an instruction the compiler was already handling.

Attempt 10 was the worst regression at +32-55%. I tried four independent accumulator chains to hide FMA latency through instruction-level parallelism. But the kernel is memory-bound, not compute-bound. FMA latency hiding is irrelevant when every cycle is waiting on memory. The four accumulator chains increased register pressure enough to cause spilling, and the strided K access pattern (THREADS_PER_ROW * 4 instead of THREADS_PER_ROW) reduced memory coalescing. I was optimizing for the wrong bottleneck.

In attempt 11 I tried tuning register count and block size. Reducing -maxrregcount from 80 to 64 had zero effect since the kernel naturally uses fewer than 64 registers. BLOCK_SIZE=256 with ROWS_PER_BLOCK=8 also changed nothing, because the 32-threads-per-row warp structure is what matters, not the block size. Bumping to #pragma unroll 8 instead of unroll 4 dropped performance by 5-87% due to register pressure and I-cache misses.

Finally, attempt 12 tried software pipelining with an explicit prologue to prefetch the next K tile into registers while computing the current one. On the B200's memory subsystem, the hardware prefetcher combined with __ldg cache hints is already doing this. Manual pipelining just doubled register pressure for data that was already on its way.

Why 3x Off?

The kernel is fundamentally memory-bound. FP4 data is tiny (4 bits per element), and even with block scaling overhead, the arithmetic intensity is low for GEMV. The speed-of-light analysis is based on DRAM bandwidth limits. Even the top solutions only reached ~2x off speed of light. The gap from 2x to 1x is genuinely hard for GEMV and may require approaches beyond what any competitor used.

The gap between my 3x and the winners' 2x came down to execution details I cover in the next section: cache policies, load widths, register pressure, and compile-time specialization. I didn't invest enough time into understanding the Nsight Compute profile for my kernel. Had I done that after attempt 7, it would have told me immediately that the kernel was memory-bound and that compute-side optimizations (ILP, wider instructions) were pointless. I instead learned this the hard way through attempts 8-12.

What the Top Solutions Did Differently

After the hackathon ended I studied the top 3 solutions (all clustered around 18.5µs geometric mean, roughly 2x speed of light vs my 3x). They shared several techniques I hadn't used.

All three wrote their load and decode paths in raw PTX rather than using C intrinsics. Where I used __nv_cvt_fp4x2_to_halfraw2, they wrote cvt.rn.f16x2.e2m1x2 directly. Where I used __ldg, they wrote ld.global with explicit qualifiers. This gives precise control over instruction selection and scheduling that the C intrinsics abstract away.

The biggest gap was cache policy control. All top solutions used different cache hints for A (the matrix, streamed once) vs B (the vector, reused across all rows). For A they used L1::no_allocate to avoid polluting L1 with data that won't be reused. For B they used L1::evict_last to keep it hot in cache since every row reads the same vector. My __ldg just requests a generic read-only cache path with no distinction between streaming and reusable data.

They also used much wider vectorized loads. I was loading with uchar4 (32-bit). The top solutions loaded with ld.global.v2.u64 (128-bit) and even ld.global.v4.u64 (256-bit), fetching 32 or 64 FP4 values in a single memory transaction. This is the kind of wider load that actually works. Unlike my attempt 9 with uint2, the top solutions decoded using PTX byte unpacking (mov.b32 {tmp0, tmp1, tmp2, tmp3}, %reg) which avoids the bitwise extraction overhead that killed my approach.

Rather than one generic kernel, the top solutions templated on the exact K dimension and dispatched at launch time. This lets the compiler fully unroll the K loop with a known trip count and select optimal register allocation per problem size. The rank 1 solution went further with per-K cache hint tuning, using different ld.global qualifiers for K=3584, K=8192, and K=1024.

Register budgets were also much tighter. The rank 1 solution used -maxrregcount=32, rank 3 used 45. I used 80. Lower register counts increase occupancy (more warps in flight), which is what a memory-bound kernel needs to hide memory latency. My attempt 11 tested this but only went down to 64, not aggressive enough.

The rank 2 solution also processed BLOCK_M rows per thread block where threads reading B data are shared across rows. This amortizes the B vector load cost. My kernel loaded B independently per warp, wasting bandwidth on redundant reads of the same vector.

One thing that surprised me; a pure PyTorch solution using torch._scaled_mm with multi-stream parallelism across the L dimension scored 22.4µs. No custom kernels at all, just calling into cuBLAS's FP4 path with the right scale factor layout. That's within 20% of the top 3 PTX solutions and faster than my hand-written C++ kernel.

Takeaways

The single most important thing I could have done after attempt 7 was run Nsight Compute and confirm the kernel was memory-bound. That would have saved me from attempts 8-12. Instead I optimized based on intuition, and intuition was wrong. Split-K doesn't help memory-bound kernels. Wider loads only help when data can be used directly without unpacking. ILP is irrelevant when the bottleneck is memory. Register tuning does nothing if you're already under the limit. Software pipelining is redundant when hardware prefetch is sufficient.

The full source code for all attempts is available on GitHub.

Footnotes

At the PTX level this maps to cvt.rn.f16x2.e2m1x2, which is a single instruction. However, the C intrinsic doesn't always compile down to it cleanly. The rank 2 hackathon winner noted that __nv_cvt_fp4x2_to_halfraw2() didn't produce the desired PTX, which is one reason all top solutions wrote inline PTX directly. ↩

Honey, I Tiled the Tensors

Thu, 26 Feb 2026 00:00:00 GMT

import Figure from "@/components/Figure.astro";

Layouts are a powerful abstraction introduced in NVIDIA's CuTe library for making operations on complicated Tensor configurations a little bit easier to understand. My goal here is to provide a good taste of how operations on these layouts work and an example matrix-matrix multiplication kernel to show the value of these abstractions and its drawbacks. This is not a full mathematical analysis of Layouts, for that a good reference is the official paper by CuTe's creator Cris Cecka, "CuTe Layout Representation and Algebra", or the earlier community note "A note on the algebra of CuTe Layouts"

GPUs don't know about tensors and their structures, operations are only performed on linear memory. Tensor structures need to be maintained by the kernels themselves. e.g. for a row-major matrix the matrix indices i, j will become index i*C + j where C is the number of columns. These kinds of mappings are pivotal in building GPU kernels since the shapes and structures of tensors can get very complicated.

A Layout we can easily define as nothing but a combination of the shape and stride (num of jumps to go from one element to next in a dimension) of a Tensor. It defines a mapping from the Tensor coordinate space to a flat array layout indices. Here stride and shape are tuples of matching dimensions. Calling the shapes S and the strides D we say

L = S : D

For example lets start with matrices, a row major n \times m matrix will be represented as (n, m) : (m, 1). Here we can see that increasing the first dimension (row) coordinate by 1 will increase the flat index by m and the second dimension (column) only increases it by 1. For getting the flat index/offset given i, j, we know we can get it by i \times m + j = i \times m + j \times 1. Or more generally we can see from how stride is defined that the flat index \text{offset} given indices i = \{ i_0, i_1, ..., i_d\} and stride s = \{s_0, s_1, ..., s_d\},

I = \sum i_k s_k = i \cdot s

Similarly the column-major matrix represented as (n,m):(1,m). Layouts can themselves be again made out of other Layouts. Lets analyze an example; ((2,2), (2,2)): ((1,4), (2,8)). Okay! Lets try some coordinates.

Coordinate `((i_0,i_1),(j_0,j_1))`	Calculation	Physical Offset
`((0,0), (0,0))`	`0(1) + 0(4) + 0(2) + 0(8)`	0
`((0,1), (0,0))`	`0(1) + 1(4) + 0(2) + 0(8)`	4
`((1,0), (0,0))`	`1(1) + 0(4) + 0(2) + 0(8)`	1
`((1,1), (0,0))`	`1(1) + 1(4) + 0(2) + 0(8)`	5
`((0,0), (1,0))`	`0(1) + 0(4) + 1(2) + 0(8)`	2
`((0,1), (1,0))`	`0(1) + 1(4) + 1(2) + 0(8)`	6
`((1,0), (1,0))`	`1(1) + 0(4) + 1(2) + 0(8)`	3
`((1,1), (1,0))`	`1(1) + 1(4) + 1(2) + 0(8)`	7

As we can see this is an interleaved layout (or swizzled). These kinds of layouts are generally used in GPU kernels to prevent memory bank conflicts.

This is okay, but it is inconvenient to do offset calculations of a particularly complex layout. Another way of visualizing this is to think about the groups individually first (2, 2) has a stride of (1, 4) creating a 0 \rightarrow 4 \rightarrow 1 \rightarrow 5 and this is repeated again in a (2,2) shape with each block in a 0 \rightarrow 2 \rightarrow 1 \rightarrow 3 layout. This can be seen by factorizing the expanded \text{offset}.

\text{offset} = i_0 \cdot 1 + i_1 \cdot 4 + j_0 \cdot 2 + j_1 \cdot 8

\text{offset} = \underbrace{(i_0 + 2 \cdot j_0)}_{r} \cdot 1 + \underbrace{(i_1 + 2 \cdot j_1)}_{c} \cdot 4

This reduces a 4D layout into a 2D layout (4, 4) : (1, 4) which is easy to visualize. But this is not possible for all layouts.

Layouts can also be operated on. This is the algebra part of the layout algebra. So we can have functions that can map from one layout to another layout. Such a function is equivalent to mapping of one flat index space to another. Hence

f : L_1 \rightarrow L_2 \implies f(i) : I(L_1) \rightarrow I(L_2)

Lets go through a few of these operations.

Coalescing

Coalescing is to simplify a layout. More formally, we can define Rank of a layout as the number of modes it has, coalesce operation reduces the number of modes i.e. to reduce its Rank. For example the layout (2,4):(1, 2) is the same as 8:1 for 1-D coordinates. So how do we get a coalesced layout? Say we have a layout L = (s_0, s_1, ...) : (d_0, d_1, ...) , (no nesting) how do we reduce a pair of layouts into one layout? We can check if the two modes are contiguous. When are they contiguous? When the next layout has a stride equaling the total flat indices of the former layout. This way they line up perfectly. so for s_0: d_0 and s_1:d_1 can reduce iff.

d_1 = s_0 \times d_0

and the two modes merge into s_0 \times s_1 : d_0. We can also see that 1 shape is an identity in reduction. i.e. 1: d' and s: d reduce to s: d. This is useful when we need to simplify some particularly gnarly tensor layouts. Coalescing can also be done by-mode i.e. we can only reduce some ranks e.g. from 4 \rightarrow 2.

Composition

Composition chains two layouts together. Given a layout A and a layout B, the composition B \circ A creates a new layout that first maps coordinates through A and then uses those resulting indices as coordinates into B. In other words, the output indices of A become the input coordinates of B.

(B \circ A)(i) = B(A(i))

Why is this useful? Composition is fundamental to the concept of tiling, where we partition a large data into smaller chunks that can be loaded into GPU's Shared Memory or further into Register Memory. We can do this with the composition operation as we'll see in later sections.

Lets work through an example. Take A = (3, 2) : (1, 3) and B = 6 : 2. Layout A maps 2-D coordinates to 1-D indices in a 3 \times 2 arrangement. Layout B maps a 1-D coordinate to offsets \{0, 2, 4, 6, 8, 10\}. The composition B \circ A should give us a layout that maps 2-D coordinates directly to the physical offsets.

Coordinate `(i, j)`	`A(i,j)`	`B(A(i,j))`
`(0, 0)`	`0`	`0`
`(1, 0)`	`1`	`2`
`(2, 0)`	`2`	`4`
`(0, 1)`	`3`	`6`
`(1, 1)`	`4`	`8`
`(2, 1)`	`5`	`10`

We can verify this is the layout (3, 2) : (2, 6). Now lets build up the general composition rules.

Single-mode `B`

In the simple case where B = s_B : d_B is a single mode, composition just scales the strides of A by d_B:

(s_B : d_B) \circ (S_A : D_A) = S_A : (d_B \cdot D_A)

This works because B is linear — B(k) = k \cdot d_B — so B(A(i)) = d_B \cdot A(i). This is the case we saw in the example above: (6:2) \circ ((3,2):(1,3)) = (3,2):(2,6).

Single-mode `A`

When B has multiple modes, the output of A needs to be decomposed into coordinates for B. Given B = (s_0, s_1, ...) : (d_0, d_1, ...) and A = s : d, the flat indices from A (which are \{0, d, 2d, ..., (s-1)d\}) are split across B's modes using its shapes. This can only be done if either d \mid s_0 or s_0 \mid d. Without this, A's indices don't cleanly align with B's mode boundaries and composition is undefined.

There are two cases:

Case 1: s_0 \mid d (stride skips past first mode). Since d is a multiple of s_0, every index k \cdot d has k \cdot d \mod s_0 = 0, so B's first mode coordinate is always zero. We skip it entirely and recurse with reduced stride:

(s_0, s_1, ...) : (d_0, d_1, ...) \circ (s : d) = (s_1, ...) : (d_1, ...) \circ (s : d / s_0)

Case 2: d \mid s_0 (stride fits within first mode). Let q = s_0 / d. The first q indices from A cycle through B's first mode before overflowing into the next.

If s \leq q: all indices fit in the first mode. Result: s : (d_0 \cdot d)
If s > q and q \mid s: the first mode fills completely, and the rest recurse: B \circ (s : d) = (q : d_0 \cdot d),\ ((s_1, ...) : (d_1, ...) \circ (s / q : 1)) ¹

Lets trace through a concrete example. Take B = (4, 3) : (1, 8) and A = 6 : 2.

We have d = 2 and s_0 = 4. Since d \mid s_0, we're in Case 2 with q = 4 / 2 = 2. Since s = 6 > q = 2 and q \mid s:

(4, 3) : (1, 8) \circ (6 : 2) = (2 : 1 \cdot 2),\ ((3,) : (8,) \circ (3 : 1))

For the recursive step, d = 1 and s_0 = 3, so q = 3 and s = 3 \leq q:

(3,) : (8,) \circ (3 : 1) = 3 : 8

Combining: B \circ A = (2, 3) : (2, 8).

General composition: left-associative reduction

For the fully general case where both A and B are multi-mode, we reduce it to the cases above by processing A's modes left to right. Given A = (s_0^A, s_1^A, ...) : (d_0^A, d_1^A, ...), we compose each mode of A one at a time (due to ¹):

B \circ A = (B \circ (s_0^A : d_0^A)),\ (B_\text{rest} \circ (s_1^A : d_1^A)),\ ...

Each single-mode composition consumes some of B's leading modes (via the cases above), and the unconsumed remainder B_\text{rest} carries forward for the next mode of A.

Composition of a Layout with a smaller Tile layout gives us a single first Tile from the original layout, we need some operation that can also do this for the rest of the layout.

Complement

Composition is always a subset of indices from the index space of A, what about the leftovers? This is where the complement operation comes into the picture. Complement is done with respect to a size M, i.e. given a layout A the complement A^*_M, is the layout that fills up the rest of M indices not covered by A. Complement only makes sense if the gaps left in M indices is shaped for filling up the space perfectly with a layout. For a layout A = (s_0, s_1, ..., s_n) : (d_0, d_1, ..., d_n), this means s_i \cdot d_i | d_{i+1} and s_n \cdot d_n | M. This means that the inner strides should fit into outer strides. We can calculate the complement of A as.

A^*_M = (d_0, \frac{d_1}{s_0 \cdot d_0}, \frac{d_2}{s_1 \cdot d_1}, ..., \frac{M}{s_n \cdot d_n}) : (1, s_0 \cdot d_0, s_1 \cdot d_1, ..., s_n \cdot d_n)

Here is an example for complement. Note we can reduce the s_0: d_0 = 1:1 from this example. The gray part is the original layout and colored indices

Division

Division operation splits a layout B into equal-sized tiles defined by layout A. As we have seen previously we can compose a layout and a tiler to get the first tile but to get all of them we need to also compose the complement with respect to a size M to get the rest of them.

B \oslash A := B \circ (A, A^*_M)

This splits each mode of B into two: the first mode is within tile from the composition and the second mode is indexing across tile from the complement. For a mode of B with shape s_B and stride d_B, and tiler shape t, we get

Within-tile: shape t, stride d_B
Across-tiles: shape s_B / t and stride t \cdot d_B.

The original mode becomes a pair: (t, s_B / t): (d_B, t \cdot d_B). This operation doubles the Rank of the layout. Lets work through a few examples, first a simple 1D case. B = 12:1 and A = 4:1, A should divide B into 3 pieces. Hence;

Complement: A^*_{12} = 3:4 (the 3 tile offsets: {0, 4, 8}).

Concatenating: (A, A^*) = (4, 3):(1, 4).

Composing: B \circ (4,3):(1,4) = (4, 3):(1, 4).

The first mode (size 4, stride 1) = within-tile indices {0,1,2,3}. Second mode (size 3, stride 4) = tile offsets {0,4,8}. Three tiles of four elements each, covering all 12.

Lets also work through a 2D example B = (4, 6):(1, 4) (4x6 column-major matrix), tiler A = (2, 3) (2x3 tiles). Mode-by-mode:

Mode 0: shape 4, tiler 2 → (2, 2):(1, 2) — 2 rows in-tile, 2 tile-rows
Mode 1: shape 6, tiler 3 → (3, 2):(4, 12) — 3 cols in-tile, 2 tile-cols Logical divide result: ((2, 2), (3, 2)) : ((1, 2), (4, 12))

This is rank-4. Each original mode became a nested pair of (\text{intra}, \text{inter}).

This layout is difficult to work with due to its nesting, we can reduce it to easier layouts

Zipped, Tiled and Flat Divides

These are just convenience representations of the logical divide operation. Lets say a layout L has shape (m, n, l, ....) and a tiler T with shape (t_m, t_n). The convenience forms are as below

Logical: ((t_m, r_m), (t_n, r_n), ...)
Zipped: ((t_m, t_n), (r_m, r_n, l, ...))
Tiled: ((t_m, t_n), r_m, r_n, l, ...)
Flat: (t_m, t_n, r_m, r_n, l, ...)

Tiled Divide is especially interesting, it regroups the modes by role: all intra-tile modes together, all inter-tile modes together. The result is a clean two-level hierarchy — first group is "what's inside a tile," second group is "which tile." This is the workhorse form that local_tile and partition_* use under the hood. It's what you want when different parts of the kernel need to reason about tiles independently — the CTA picks its tile via the inter-tile mode, then threads work within it via the intra-tile mode.

Product

Where division breaks a layout into tiles, product does the opposite — it replicates a layout to fill a larger space. Given a layout A (the atom to replicate) and a layout B (the replication pattern), the product is

A \otimes B := (A, A^* \circ B)

The result is a two-mode layout: mode 0 is the atom A itself and mode 1 is A^* \circ B which describes how the copies of A are arranged. The complement A^* finds the gaps between elements of A, and composing with B maps the replication pattern into those gaps.

Lets work through a 1D example. Take A = 4:1 (a contiguous atom of 4 elements) and B = 3:1 (replicate 3 times).

First we need the complement of A with respect to M = \text{size}(A) \times \text{size}(B) = 12. We get A^*_{12} = 3:4 — the three offsets \{0, 4, 8\} where copies begin.

A^* \circ B = (3:4) \circ (3:1) = 3:4

So the product is A \otimes B = (4, 3) : (1, 4). Notice anything? This is the same layout we got from dividing 12:1 by 4:1. That's the duality — division splits a layout into tiles, product builds one up from tiles. Two sides of the same coin.

Where product really shines is building thread-value layouts for distributing work across GPU threads. Say we have 4 threads each handling 2 values. The thread layout T = 4:1 and value layout V = 2:1, the product T \otimes V tells us which elements each thread owns.

T^*_8 = 2:4, so T^* \circ V = (2:4) \circ (2:1) = 2:4. The product is (4, 2) : (1, 4).

Thread 0 handles elements \{0, 4\}, thread 1 handles \{1, 5\} and so on — each thread's values are spaced apart by the number of threads. This cyclic distribution is exactly what CuTe's make_tiled_copy_tv builds internally from a thread layout and value layout.

Blocked and Raked Products

For multi-dimensional layouts, the logical product works mode-by-mode just like division. Given atom A with shape (a_m, a_n) and replication pattern B with shape (b_m, b_n), the result has structure ((a_m, r_m), (a_n, r_n)) where r denotes the replica modes. Blocked and raked products are two ways of reassociating these modes.

Blocked product groups like-modes with the atom inside: ((a_m, r_m), (a_n, r_n)). Each atom occupies a contiguous block and replicas tile these blocks across the space. Think of it as "each thread gets a contiguous rectangle of elements."

Raked product reverses the nesting: ((r_m, a_m), (r_n, a_n)). The replicas interleave within the atom dimensions, creating a cyclic distribution. Thread 0 gets element 0, thread 1 gets element 1, wrapping around. This is the classic GPU pattern for coalesced memory access — adjacent threads access adjacent memory locations.

Zipped and Tiled Products

Same convenience regrouping as with division. Given a layout L with shape (m, n, l, ...) and tiler T with shape (t_m, t_n):

Logical: ((m, t_m), (n, t_n), l, ...)
Zipped: ((m, n), (t_m, t_n, l, ...))
Tiled: ((m, n), t_m, t_n, l, ...)
Flat: (m, n, t_m, t_n, l, ...)

Putting it all together

Now lets see how all these layout operations come together in a real tiled GEMM kernel using CuTe's Python DSL. We compute C = A \times B where A is M \times K, B is K \times N (stored as N \times K, i.e. row-major transposed), and C is M \times N. The kernel tiles across all three dimensions with block tile sizes (b_M, b_N, b_K).

Setup and Tiled Copies

First we define shared memory layouts and tiled copy operations for moving data from global memory (GMEM) to shared memory (SMEM):

# Shared memory layouts - column-major for vectorized MMA access
sA_layout = cute.make_layout((b_m, b_k))  # (b_m, b_k) : (1, b_m)
sB_layout = cute.make_layout((b_n, b_k))  # (b_n, b_k) : (1, b_n)

# Thread layout for copy A: threads distributed across k-major order
tA = cute.make_layout(
    (num_threads // b_k, b_k), stride=(b_k, 1)
)
copy_atom_A = cute.make_copy_atom(
    cute.nvgpu.cpasync.CopyG2SOp(), dtype,
    num_bits_per_copy=dtype.width,
)
tiled_copy_A = cute.make_tiled_copy_tv(
    copy_atom_A, thr_layout=tA, val_layout=cute.make_layout((1, 1))
)

# Thread layout for copy B: with vectorized loads along n-major
num_vectorized = 4  # elements per vectorized load
copy_atom_B = cute.make_copy_atom(
    cute.nvgpu.cpasync.CopyG2SOp(), dtype,
    num_bits_per_copy=dtype.width * num_vectorized,
)
major_mode_size = b_n // num_vectorized
tB = cute.make_layout(
    (major_mode_size, num_threads // major_mode_size),
    stride=(1, major_mode_size),
)
tiled_copy_B = cute.make_tiled_copy_tv(
    copy_atom_B, thr_layout=tB,
    val_layout=cute.make_layout((num_vectorized, 1))
)

The make_tiled_copy_tv function takes a copy atom (the hardware copy instruction), a thread layout (how threads are mapped to tile elements), and a value layout (how many elements each thread copies per invocation). This is where layout composition shines — CuTe composes these layouts to determine exactly which elements each thread is responsible for copying.

MMA Setup

The tiled MMA is set up similarly. We distribute threads in a (T/16, 16, 1) layout across the M, N, K modes:

atoms_layout_mnk = cute.make_layout(
    (num_threads // 16, 16, 1), stride=(16, 1, 0)
)
tiled_mma = cute.make_tiled_mma(
    cute.nvgpu.MmaUniversalOp(abacc_dtype=acc_dtype),
    atom_layout_mnk=atoms_layout_mnk,
)

The Kernel

Inside the kernel, we first use local_tile to carve out each thread block's portion of the global tensors. The proj argument selects which modes the block tiler applies to:

bidx, bidy, _ = cute.arch.block_idx()
cta_coords = (bidx, bidy, None)

# gA: (b_m, b_k, k_tiles), gB: (b_n, b_k, k_tiles), gC: (b_m, b_n)
gA = cute.local_tile(mA, block_tiler, cta_coords, proj=(1, None, 1))
gB = cute.local_tile(mB, block_tiler, cta_coords, proj=(None, 1, 1))
gC = cute.local_tile(mC, block_tiler, cta_coords, proj=(1, 1, None))

Then we partition the tiles across threads for both copying and computation:

tidx, _, _ = cute.arch.thread_idx()
thr_mma = tiled_mma.get_slice(tidx)

# Partition copy source (GMEM) and destination (SMEM) for each thread
thr_copy_A = tiled_copy_A.get_slice(tidx)
tAgA = thr_copy_A.partition_S(gA)  # (cpy, cpy_m, cpy_k, k_tiles)
tAsA = thr_copy_A.partition_D(sA)  # (cpy, cpy_m, cpy_k)

thr_copy_B = tiled_copy_B.get_slice(tidx)
tBgB = thr_copy_B.partition_S(gB)  # (cpy, cpy_n, cpy_k, k_tiles)
tBsB = thr_copy_B.partition_D(sB)  # (cpy, cpy_n, cpy_k)

# Partition SMEM and output for MMA
tCsA = thr_mma.partition_A(sA)
tCsB = thr_mma.partition_B(sB)
tCgC = thr_mma.partition_C(gC)
tCrC = tiled_mma.make_fragment_C(tCgC)
tCrC.fill(0.0)

Each partition_* call internally uses layout composition and division to split the tile into per-thread pieces. The cpy mode in the copy partitions captures both the vectorization width and the number of copy operations each thread performs.

The Main Loop

The k-tile loop copies each (b_M, b_K) and (b_N, b_K) tile from GMEM to SMEM using async copies, then performs the MMA:

k_tiles = cute.size(tAgA, mode=[3])

for k in range(k_tiles):
    cute.copy(tiled_copy_A, tAgA[None, None, None, k], tAsA, pred=tApA)
    cute.copy(tiled_copy_B, tBgB[None, None, None, k], tBsB, pred=tBpB)

    cute.arch.cp_async_commit_group()
    cute.arch.cp_async_wait_group(0)
    cute.arch.sync_threads()

    cute.gemm(tiled_mma, tCrC, tCsA, tCsB, tCrC)
    cute.arch.sync_threads()

The pred arguments handle boundary conditions — when M, N, or K aren't exact multiples of the tile sizes, predicate tensors mask out-of-bounds accesses. These predicates are themselves built using layout operations on identity tensors.

Epilogue

Finally the accumulated results are written back to global memory, again with predication for bounds checking:

cC = cute.make_identity_tensor(gC.shape)
tCpC = thr_mma.partition_C(cC)
predC = cute.make_rmem_tensor(tCrC.layout, cutlass.Boolean)
residue_m = mC.shape[0] - b_m * bidx
residue_n = mC.shape[1] - b_n * bidy
for i in range(cute.size(tCrC.shape)):
    predC[i] = cute.elem_less(tCpC[i], (residue_m, residue_n))
atom = cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), mC.element_type)
cute.copy(atom, tCrC, tCgC, pred=predC)

The full kernel is available in the companion repo.

Notice how the layout algebra permeates the entire kernel — from defining how threads map to data (make_tiled_copy_tv), to carving tiles (local_tile), to splitting work across threads (partition_*), to bounds checking (make_identity_tensor + elem_less). Without these abstractions, we'd be manually computing thread-to-element mappings with error-prone index arithmetic. The layout algebra replaces all of that with composable, type-safe operations.

Closing thoughts

CuTe's layout algebra turns what would be pages of error-prone index arithmetic into a handful of composable operations — composition, complement, and division — that are easier to reason about. But I want to be honest about the tradeoff here.

GPU memory layouts are inherently complex. Swizzled shared memory, bank conflict avoidance, mixed-radix thread-to-data mappings — this complexity is intrinsic to the hardware. CuTe doesn't eliminate it; it repackages it into a different formalism. You're trading one kind of complexity (raw index math) for another (an algebra with its own rules, edge cases, and debugging challenges). As one Reddit commenter put it, CUTLASS is "a typical example of a library with so much abstractions that makes complicated things simple and simple things complicated."

The learning curve is steep. You need to internalize layouts, composition, complement, division, cosize, cotarget, tiled copies, MMA atoms — a whole vocabulary before you can write your first kernel. For someone who already understands raw CUDA well, the value proposition is debatable: you're investing significant effort to learn abstractions that, at the end of the day, generate the same PTX. The payoff comes when you need to support multiple GPU architectures, swap out MMA instructions, or restructure tiling strategies without rewriting everything — that's where the composability genuinely shines. But for a one-off kernel on a single architecture, you might be better off with raw CUDA and a good understanding of your hardware.

If you want to dig deeper, Cris Cecka's "CuTe Layout Representation and Algebra" is the definitive reference — it formalizes everything covered here and more, including inversion and cotarget. The CuTe documentation is thorough, and the CUTLASS Python DSL makes it easier to experiment with these ideas interactively — the Python interface significantly reduces the compile-time pain that plagues the C++ templates.

Footnotes

(s_0, s_1, ...): (d_0, d_1, ...) and listing them elementwise (s_0: d_0), (s_1: d_1), ... is the same. ↩ ↩²

Hierarchical Navigable Small Worlds

Sun, 08 Feb 2026 00:00:00 GMT

import Figure from "@/components/Figure.astro";

You might have heard of this algorithm during the RAG-mania of late. But not many sources on the web go into the inner workings and intuition building, which I would like to cover here.

Hierarchical Navigable Small Worlds (HNSW) is a top-performing and popular approximate nearest neighbor search algorithm. To find the exact nearest neighbor given a bunch of vectors, we need to find a vector that is the closest to a query vector. i.e. given a set of vectors X = \{x_i\}, and a query q, distance function \text{dist}(\cdot, \cdot) the goal is

x^* = \operatorname*{argmin}_{x_i \in X} \text{dist}(x_i, q)

The trivial but brute force way to do this is just iterate through all the vectors, find a minimum by distance. This is not scalable since the amount of vectors can be in billions. Another aspect of this is the number of vector dimensions which can be as large as 4096 in latest state of the art embedding models. We would also want to eliminate any extra distance calculations, since the time spent in these will not be small.

Delaunay Triangulation

A common approach when wanting to optimize for multiple queries, is to look for a data structure we can build on top of current data so we can spend some compute and memory to lower the per query cost exponentially. What if we could store a "region" around every vector, such that all points in that region are the closest to the vector. A query point q if it falls into p's region, p by definition is the nearest neighbor of q.

These regions are called Voronoi Cells and if we draw the region boundaries, this becomes a Voronoi Diagram. Here is a small visualization.

We can also represent each point as a node in a graph, and each adjacent cell as neighbors to the node. This is a mathematically equivalent representation of the Voronoi Diagram and is called Delaunay triangulation. This way for each point in the graph, finding the nearest neighbor is trivial, we just need to get the nearest out of the neighbors of the current node.

But our problem statement is for point q outside the given set of points, how do we use the Delaunay triangulation to answer this question? Imagine we are at a point p, such that the point q is not in the Voronoi cell of p. What is a good decision I can make given the local information I have about my neighbors? I can find a neighbor r such that dist(r, q) < dist(p, q). But will this work? Let's say we move to r and repeat this process. We can see that we strictly decrease the dist(\cdot, q) until we reach a point where we can find no such r' such that dist(r', q) < dist(p', q), which means p' is the nearest neighbor of q.

You might have noticed that we selected any closer neighbor and not the closest neighbor in the Delaunay Walk, since the process converges to produce the correct answer anyways, picking the closest adds extra computation for a reduction in number of steps the algorithm takes to converge.

The big problem with this is the curse of dimensionality, the amount of space needed to store the graph grows superpolynomially¹ and becomes infeasible to store beyond small dimensions. This can be easily skirted if we consider an approximate problem, we can cap the number of neighbors to M in which case, the memory requirements reduce to \mathcal{O}(nM). But we lose the monotonic routing behavior we discussed earlier.

Small Worlds and their Navigation

You might have heard about the six degrees of separation experiment in which social psychologist Stanley Milgram² asked people to send a letter to a target individual, but a person can only forward the letter to a single acquaintance that they know on a first name basis. About a third of the letters reached the target person with a median of 6 steps. These and similar experiments serve as the basic evidence of existence of short paths in the global friendship network, linking all (or almost all) of us together in society.

This kind of small world phenomena is also present in a lot of places outside society, e.g. the internet itself, power grids, the fully mapped biological neural network of the worm C. elegans and many more. All this happens because of "long-range links" between local clusters, e.g. most of your friends are local; the one overseas friend can shorten the total number of links needed to reach the destination.

The same concept can work if we add these long range links to the Delaunay graph. We can randomly add new links but adding too many of them, we lose local structure and the greedy walk will not easily converge since we can't determine if we are moving closer or not, we need to settle on some middle ground strategy. The optimal result comes from Jon Kleinberg³. Kleinberg analyzed a lattice in d dimensions, where long range links are added with the probability of link from u \rightarrow v, based on the distance d(u,v) between them.

P(u \rightarrow v) \propto \cfrac{1}{d(u,v)^{\alpha}}

If \alpha is low; the distance doesn't matter much. Long links are just as likely as short ones. Jumps can overshoot the target.
If \alpha is high; the penalty of distance is huge. We mostly just link to nearby links. Jumps can be too timid and might not help.

Kleinberg proved that the optimal result occurs at \alpha = d.

Intuition Behind `\alpha = d` and Runtime Cost

The number of points at distance r in a d-dimensional space grows \propto r^d. And the probability of linking to one such node is \propto r^{-\alpha}. This means the total probability of reaching distance r is \approx r^{(d-\alpha)}. Now at \alpha = d, this probability becomes constant, i.e. we can reach any distance with the same probability and creates a Scale Invariance in the graph.

Due to the Scale Invariance property of the long-range links, we can imagine a "ring" R_k between distance 2^{k-1} to 2^k. There would be \mathcal{O}(\log n) such rings. The probability of finding a link to a certain link in a ring is constant. Let's try to calculate this. We know P(u \rightarrow v) \propto d^{-\alpha}, but we need to normalize the probability. The normalization denominator would be.

Z = \displaystyle\sum_{u \ne v} \cfrac{1}{d(u, v)^{\alpha}}

substituting the ring distances we get

Z \approx \displaystyle\sum_{k=0}^{\mathcal{O}(\log n)}{2^{k(d-\alpha)}}

at d = \alpha, this becomes a sum of 1s.

Z \sim \mathcal{O}(\log n)

this implies the probability of reaching a ring is, ring's weight (1) divided by Z.

P(R_k) \approx \cfrac{1}{Z} \sim \cfrac{1}{\mathcal{O}(\log n)}

Our goal to move to the Ring R_k from a previous ring R_{k-1}, we have only two choices either we find a long-link and move or continue to explore in the current ring. The expected number of trials before we succeed is \mathbb{E}{[T]} = 1/p, ⁴ this implies expected number of steps to exit the ring is \mathcal{O}(\log n). Hence.

\text{Total Cost} = \text{\# of Rings} \times \text{\# of Steps per Ring} = \mathcal{O}(\log^2 n)

This is a good improvement, but the algorithm still is very sensitive to input data and has a tendency to get stuck in a local minimum when data is highly clustered. This means expensive restarts or adding more links, both of which increase runtime, we have to tradeoff latency to get a better recall. HNSW tries to fix robustness while providing an \mathcal{O}(\log n) query time.

Hierarchical Navigable Small Worlds⁵

The NSW greedy walk's run can be seen as a set of zoom in and zoom out phases. Zoom in phase goes to nodes that are highly clustered, on the other hand the zoom out phases travels long distances. Zoom in phases can get stuck in local minimas, forcing us to do costly restarts.

What if we separate out links of different length; start from the longest ones and progressively switch to shorter ones as we exhaust possibilities in the longer ones. This could add some robustness since we are not constantly zooming in and out. Another benefit of this approach is that the number of links will be a lot less in initial steps and we can explore a larger part of the graph, making our search cheaper.

This is the key idea behind the Hierarchy of HNSW, we create layers of graphs each with increasing coarseness, akin to having different zoom levels.

Layer 0: All the nodes, edges are short and local.
Layer 1: Fewer nodes, medium length edges.
...
Layer L: Top level; very few nodes, we can globally traverse the graph in handful of hops.

This is structurally very similar to Skip Lists.⁶ ⁷

Querying this structure is easy; we can start on the top layer; just do the same greedy walk on each layer and move to next layer when we exhaust our options.

Construction

Let's talk construction, so when adding a node to layer l, we need to (a) add it to all the layers below l down to 0. How do we find a layer for a new node? Also, what would be the best distribution of nodes across layers for fast queries. We can take a similar geometric distribution from previous section, since we will get exponentially decreasing nodes in each layer and about \mathcal{O}(\log n) layers. So P(L = l) \propto m^{-l}, ⁸ here m the tunable parameter. Note that this is constant per layer and not dependent on the number of nodes. This means

P(L = l) = C \cdot {m}^{-l}

To find C, we sum all the probabilities.

1 = \displaystyle\sum_{l=0}^{\infty} C{m}^{-l} = C \displaystyle\sum_{l=0}^{\infty} {m}^{-l}

Using sum of the geometric series, for m > 1.

1 = C \cdot \cfrac{1}{1 - \frac{1}{m}} \implies C = 1 - \cfrac{1}{m}

P(L = l) = \bigg(1 - \cfrac{1}{m}\bigg) \bigg(\cfrac{1}{{m}^l}\bigg)

We can also compute the tail probability P(L \geq l).

P(L \geq l) = \displaystyle\sum_{k=l}^{\infty} P(L=k) = \displaystyle\sum_{k=l}^{\infty}\bigg(1 - \cfrac{1}{m}\bigg){\bigg(\cfrac{1}{m}\bigg)}^k

Factoring out (1/{m})^l

P(L \geq l) = \bigg(\cfrac{1}{m}\bigg)^l\displaystyle\sum_{t=0}^{\infty}\bigg(1 - \cfrac{1}{m}\bigg){\bigg(\cfrac{1}{m}\bigg)}^t = {m}^{-l}

When sampling the level L, from uniform random distribution we see.

L = \bigg\lfloor \cfrac{-\ln U}{\ln m} \bigg\rfloor, U \sim \text{Uniform}(0,1)

It is also interesting to observe the the expected number of nodes in a layer will be

\mathbb{E}[N_l] = N \cdot P(L \geq l) = N m^{-l}

and the ratio of the layer sizes is

\cfrac{\mathbb{E}[N_l]}{\mathbb{E}[N_{l+1}]} = m

Therefore, m can also be thought of as a zoom factor.

Neighbor Selection Heuristic

While adding a node we also have to select neighbors to add, we can use the naive strategy of just picking the M nodes closest to us, but this strategy has the same issues with robustness as previously discussed. But we can easily add link diversity using a simple heuristic. We add a neighbor iff the neighbor is closer to us than any other neighbors, otherwise we skip. This strategy creates room in the M node cap for longer links to fill. To implement this we can Sort candidates by distance to the new node; keep a candidate if it is not too close to any already selected neighbor dist(candidate, selected) < dist(candidate, new).

Intuition Behind Runtime Cost

What is the runtime cost then? We know that we have \mathcal{O}(\log n) layers; we just need to find the amount of work needed per layer. Imagine we have just come down from layer l+1 to l. Search on this layer would always terminate before encountering a node that is also present in l+1, otherwise the new point in l+1 would have been the entrypoint to layer l. The probability of a point being present in both l+1 and l is

P(L \geq l+1 | L \geq l) = \cfrac{P(L \geq l+1)}{P(L \geq l)} = \cfrac{{m}^{-(l+1)}}{{m}^{-l}} = 1/{m}

Note that this is independent of layers and number of nodes. While on a walk of layer l we say p is probability of finding the point in l+1 (from above). The probability of continuous s successes is p(1-p)^s. The expected number of steps using similar logic as NSW follows as \mathbb{E}[s] = 1/p, since p = 1/m, we get \mathbb{E}[s] = m. This is independent of number of nodes hence the total search complexity is \sim \mathcal{O}(\log n).

Implementation

I am implementing this in Zig, all nodes are stored as indices of a global vector array. Here is the definition of a Node. You might have noticed that we have different sizes for neighbors in layer 0. In practice if we add more neighbors in layer 0 it can improve the recall. We generally set neighbor_size0 = 2*neighbor_size. Each node has neighbors in multiple layers.

pub const NodeIdx = u32;

const Node = struct {
    neighbors: []std.ArrayListUnmanaged(NodeIdx),

    pub fn initEmpty(
        alloc: std.mem.Allocator,
        layer: usize,
        neighbor_size: usize,
        neighbor_size0: usize,
    ) !Node {
        const neighbors = try alloc.alloc(std.ArrayListUnmanaged(NodeIdx), layer + 1);
        for (neighbors, 0..layer + 1) |*nbrs, l| {
            nbrs.* = .empty;
            const node_cap = if (l == 0) neighbor_size0 else neighbor_size;
            try nbrs.ensureTotalCapacityPrecise(alloc, node_cap);
        }

        return .{ .neighbors = neighbors };
    }
};

Now we define the main data structure for HNSW, here Params as tunable parameters for the algorithm.

max_neighbors_per_layer: (M) Max. number of neighbors of a node. Large values can increase recall but at the cost of runtime.
ef_construction: Construction entry factor, number of nearest candidates kept during construction.
ef_search: Search entry factor, number of nearest candidates kept during search.
layer_mult: (m) Layer Multiplier, usually set to approx. max_neighbors_per_layer.

pub const HnswIndex = struct {
    pub const Params = struct {
        max_neighbors_per_layer: usize,
        ef_construction: usize,
        ef_search: usize,
        num_words: usize,
        layer_mult: f64,
        seed: u64 = 2026,
    };

    // ...omitted boilerplate vars...

    entry_points: std.ArrayListUnmanaged(NodeIdx),
    layers: usize = 0,
    nodes: std.ArrayListUnmanaged(Node),
    visited: std.DynamicBitSetUnmanaged,
}

Let's add the function to search a single layer. We maintain two heaps: one to store the node candidates we want to consider, and farthest which keeps the current best ef results (its top is the worst among them). The rest is standard best first search.

fn searchLayer(
    self: *HnswIndex,
    layer: usize,
    idx: NodeIdx,
    entry_points: std.ArrayListUnmanaged(NodeIdx),
    entry_factor: usize,
) !std.ArrayListUnmanaged(NodeIdx) {
    self.visited.unsetAll();

    var candidates = std.PriorityQueue(SearchEntry, void,minCompareSearch).init(self.allocator,{});
    defer candidates.deinit();
    var farthest = std.PriorityQueue(SearchEntry, void, maxCompareSearch).init(self.allocator, {});
    defer farthest.deinit();

    for (entry_points.items) |entry_point| {
        self.visited.set(entry_point);
        try candidates.add(.{ .idx = entry_point, .dist = self.dist(idx, entry_point) });
        try farthest.add(.{ .idx = entry_point, .dist = self.dist(idx, entry_point) });
    }

    while (candidates.count() > 0) {
        const curr_candidate = candidates.remove();
        const curr_idx = curr_candidate.idx;
        const curr_distq = curr_candidate.dist;
        const curr_farthest = farthest.peek();

        if (curr_distq > curr_farthest.?.dist) {
            break;
        }

        for (self.nodes.items[curr_idx].neighbors[layer].items) |nbr| {
            if (!self.visited.isSet(nbr)) {
                self.visited.set(nbr);
                const nbr_dist = self.dist(nbr, idx);
                const curr_farthest_ = farthest.peek();
                if (farthest.count() < entry_factor or nbr_dist < curr_farthest_.?.dist) {
                    try candidates.add(.{ .idx = nbr, .dist = nbr_dist });
                    try farthest.add(.{ .idx = nbr, .dist = nbr_dist });
                    if (farthest.count() > entry_factor) {
                        _ = farthest.remove();
                    }
                }
            }
        }
    }

    const count = farthest.count();
    var results: std.ArrayListUnmanaged(NodeIdx) = .empty;
    try results.resize(self.allocator, count);
    var i: usize = count;
    while (farthest.count() > 0) {
        i -= 1;
        const entry = farthest.remove();
        results.items[i] = entry.idx;
    }

    return results;
}

For insertion, we assign a random layer using the sampling method discussed earlier. Then search for entry_points into the assigned_layer, with entry_factor = 1. Then add the new node from the assigned_layer down to 0.

pub fn insert(self: *HnswIndex, idx: NodeIdx) !void {
    const random = self.prng.random();

    const urandom = random.float(f64);
    const assigned_layer: usize = @intFromFloat(std.math.floor(
        -std.math.log(f64, std.math.e, urandom) / std.math.log(f64, std.math.e, self.layer_mult),
    ));

    const node = try Node.initEmpty(
        self.allocator,
        assigned_layer,
        self.params.max_neighbors_per_layer,
        self.max_neighbors_layer0,
    );

    try self.nodes.append(self.allocator, node);

    if (self.nodes.items.len == 1) {
        try self.entry_points.append(self.allocator, idx);
        self.layers = assigned_layer;
        return;
    }

    if (assigned_layer == self.layers) {
        try self.entry_points.append(self.allocator, idx);
    }

    var curr_entry_points: std.ArrayListUnmanaged(NodeIdx) = self.entry_points;
    var current_layer = self.layers;

    var current_nearest: std.ArrayListUnmanaged(NodeIdx) = undefined;

    while (current_layer > assigned_layer) {
        const new_entry_points = try self.searchLayer(current_layer, idx, curr_entry_points, 1);
        curr_entry_points = new_entry_points;
        current_layer -= 1;
    }

    current_layer = @min(assigned_layer, self.layers) + 1;
    while (current_layer > 0) {
        current_layer -= 1;

        current_nearest = try self.searchLayer(
            current_layer,
            idx,
            curr_entry_points,
            self.params.ef_construction,
        );

        try self.selectNeighbors(current_layer, idx, &current_nearest);

        for (current_nearest.items) |nbr_idx| {
            try self.nodes.items[idx].neighbors[current_layer].append(self.allocator, nbr_idx);

            const nbr_neighbors = &self.nodes.items[nbr_idx].neighbors[current_layer];
            try nbr_neighbors.append(self.allocator, idx);

            const max_neighbors = if (current_layer == 0) {
                self.max_neighbors_layer0
            } else {
                self.params.max_neighbors_layer0
            };
            if (nbr_neighbors.items.len > max_neighbors) {
                try self.selectNeighbors(current_layer, nbr_idx, nbr_neighbors);
            }
        }

        curr_entry_points = current_nearest;
    }

    if (assigned_layer > self.layers) {
        self.layers = assigned_layer;
        self.entry_points.clearRetainingCapacity();
        try self.entry_points.append(self.allocator, idx);
    }
}

Now we add the selectNeighbors function, this is the neighbor selection heuristic discussed earlier. We can sort the neighbors and skip nodes that are closer to another neighbor than to the query idx.

fn selectNeighbors(
    self: *HnswIndex,
    layer: usize,
    idx: NodeIdx,
    candidates: *std.ArrayListUnmanaged(NodeIdx),
) !void {
    const max_neighbors = if (layer == 0) {
        self.max_neighbors_layer0
    } else {
        self.params.max_neighbors_per_layer
    };

    var candidate_entries: std.ArrayListUnmanaged(SearchEntry) = .empty;
    defer candidate_entries.deinit(self.allocator);

    for (candidates.items) |c| {
        try candidate_entries.append(self.allocator, .{ .idx = c, .dist = self.dist(c, idx) });
    }

    candidates.clearRetainingCapacity();

    std.mem.sort(SearchEntry, candidate_entries.items, {}, minCompare);

    for (candidate_entries.items) |c| {
        var good = true;
        for (candidates.items) |s| {
            if (self.dist(c.idx, s) < c.dist) {
                good = false;
                break;
            }
        }
        if (good) {
            try candidates.append(self.allocator, c.idx);
        }
        if (candidates.items.len >= max_neighbors) {
            break;
        }
    }
}

Lastly we can add topK to find the K nearest neighbors. Here we just greedily search nodes in all layers except layer 0. On layer 0, we cast a larger net and filter out the top k nodes.

pub fn topK(self: *HnswIndex, idx: NodeIdx, k: usize) !std.ArrayListUnmanaged(NodeIdx) {
    var candidates: std.ArrayListUnmanaged(NodeIdx) = undefined;
    var curr_entry_points = self.entry_points;
    var curr_layer = self.layers;

    while (curr_layer > 0) {
        const new_entry_points = try self.searchLayer(curr_layer, idx, curr_entry_points, 1);
        curr_entry_points = new_entry_points;
        curr_layer -= 1;
    }

    candidates = try self.searchLayer(0, idx, curr_entry_points, self.params.ef_search);
    candidates.shrinkRetainingCapacity(k);
    return candidates;
}

This implementation is not going to be very performant, since I am not using a lot of optimal data structures and we have a lot of pointer chasing. On my AMD Ryzen 9 9950x3D (single core) using the GloVe Dataset's 50-dim sample (1.29M vectors), I am getting these results.

Index build time: 93.10s
Average recall: 75.00%
Average query time: 92.0µs (100 queries, k=10)

This seems good enough for a quick and dirty implementation. Code can be found on Github. Perhaps I will create another post about optimizing this in the future.

Footnotes

Why? See McMullen's Upper Bound Theorem ↩
Travers, Jeffrey, and Stanley Milgram. "An experimental study of the small world problem." Sociometry (1969): 425-443. ↩
Kleinberg, Jon. "The small-world phenomenon: An algorithmic perspective." Proceedings of the thirty-second annual ACM symposium on Theory of computing (2000). ↩
This follows from the geometric distribution. If each trial succeeds with probability p, then \mathbb{E}[T] = \sum_{k=1}^{\infty} k \cdot p(1-p)^{k-1} = \frac{1}{p}. Intuitively: with p = 0.5, you expect 2 coin flips to get heads; with p = 0.1, you expect 10 trials. ↩
Malkov, Yu A., and D. A. Yashunin. "Efficient and robust approximate nearest neighbor search using Hierarchical Navigable Small World graphs." IEEE transactions on pattern analysis and machine intelligence 42.4 (2018): 824-836. arXiv:1603.09320 ↩
Skip List: Wikipedia ↩
An Anlaysis of Skip Lists ↩
This is not the same in the original paper. Original paper uses exp(-l/{m_L}). But in my opinion m^{-l} framing leads to better intuition. ↩

The Rabbit Hole of Building a Filesystem Watcher

Tue, 23 Sep 2025 00:00:00 GMT

Some of the systems I work with are highly customized environments, and often need support engineers to maintain them. A lot of automation exists, but sometimes they need to manually go into a VM and change things. This is normal , but with these manual tasks, mistakes are inevitable. One such case is a service that would only work if all the files and directories it manages are owned by a special user. But sometimes people run commands in the service directories as root. This doesn't impact the service as it's running, but it won't restart. While the fix is simple, just chown -R the service directory. There are many easy ways to prevent this, e.g. setting file permissions, File ACLs. These are less strict as root user can override these. Setting SELinux policies would be a much stricter solution. These are very sensible solutions. But what is the fun it that? How about we build an entire filesystem event watcher ourselves?

Attempt 1 - `fanotify`

fanotify is a set of APIs in the Linux kernel by which we could get filesystem events sent to userspace. Let's dive in, according to man page, we first need to call fanotify_init with proper flags; This sets up a kernel-space notification group. We can set up the directories we need to watch via fanotify_mark. fanotify_init sets up a file descriptor for the event queue, which can be accessed by reading the file descriptor. This is a great built-in API, but we have a few issues.

We cannot monitor a directory recursively. This feature is only available for whole filesystem mounts.
Another limitation is that fanotify only gives us the PID of the process that triggered the event (metadata->pid), not the full credentials. If we want to know who (which UID/GID) actually performed the operation, we must do an extra lookup in /proc/<pid> (for example, reading /proc/<pid>/status) to fetch the task's credentials. That means for every single event, we would need to open and parse a /proc file, and then apply our filtering logic.

Attempt 2 - eBPF

This was when I put this idea on the back burner, but then I stumbled on eBPF while working on another project with Falco.

eBPF enables running programs in kernel space. Programs are first compiled into bytecode, then verified by an in-kernel static verifier, then run using JIT for native execution performance. To communicate with the user-space, we can instantiate various forms of data structures too. The official intro docs do a great job of explaining this, see What is eBPF?.

I have been fortunate enough to be writing this at a time when tooling around eBPF has evolved a lot. Earlier tools had to include kernel headers by either a) compiling the program with the exact kernel source present locally or b) compiling the program on the server where it will run. Thanks to improvements around adding lightweight type info BTF(BPF Type Format) , CO-RE (Compile Once - Run Everywhere) and libbpf loader. The user interface for writing eBPF programs is a bit easier.

Now the question comes, what do you hook into? We can directly hook into kernel VFS layer functions such as vfs_mkdir and vfs_create, which abstract out various filesystem implementations and expose a single filesystem interface to user-space. We could read the arguments and filter out the events shipped to userspace, saving on a lot of context switches.

This method again has its own slew of annoyances.

Using kprobes on functions like vfs_* does not guarantee a stable ABI, i.e the arguments can change anytime, or functions themselves can disappear across kernel releases. In my case, this is not a big deal since I would be running this in a standardized environment with consistent kernel versions. But this is a solvable problem, though requiring more engineering effort. See this section about handling kernel change in the BPF-CORE reference
We will have to write the path filtering logic in kernelspace using eBPF, since vfs_* probes will trigger for all events. We will have to walk the filesystem tree up and see if some dir matches our monitored dir. Aside from the complexity of writing this, each eBPF program is statically verified. It must not contain unbounded loops, and we have a limited stack size (typically 512 bytes).

Walking up the file tree in eBPF

With the generous help of Andrii Nakryiko's excellent BPF CO-RE reference guide, I was able to come up with a good enough solution. We can use the dentry struct to walk up the tree. But since we can't have unbounded loops in BPF, I had to truncate the walk at MAX_DEPTH, which is acceptable for my problem statement since the expected depth of the directory I want to monitor is known.

static bool is_monitored_dir(struct dentry *dentry, __u64 target_ino) {
  bpf_rcu_read_lock();
  struct dentry *curr_dentry = BPF_CORE_READ(dentry, d_parent);
  struct inode *curr_inode;
  __u64 curr_ino;
  bool result = false;

  #pragma unroll
  for(int i=0; i < MAX_DEPTH; i++) {
    if (!curr_dentry) {
      break;
    }

    curr_inode = BPF_CORE_READ(curr_dentry, d_inode);
    curr_ino = BPF_CORE_READ(curr_inode, i_ino);
    if (curr_ino == target_ino) {
      result = true;
      break;
    }

    struct dentry *parent_dentry = BPF_CORE_READ(curr_dentry, d_parent);
    if (curr_dentry == parent_dentry) {
      break; // curr_dentry is its own root, we have reached the top of
                    // the tree.
    }

    curr_dentry = parent_dentry;
  }

  bpf_rcu_read_unlock();
  return result;
}

Note the kernel RCU (Read, Copy, Update) locks are needed since the dentry tree can change while we are traversing it. The RCU mechanism lets the readers safely traverse without blocking the writers.

For a complete, working example of this approach, please refer to the fs-watcher GitHub repository . This repository contains the full source code.

Better Probes

LSM hooks provide a more stable and semantically meaningful API for monitoring filesystem events, since they are part of the kernel’s Linux Security Module framework. They can reduce the number of events you need to filter and eliminate some of the brittleness associated with probing low-level VFS functions. However, these hooks were not available in the kernel I was working with. With LSM hooks, we have access to the path struct with which we can resolve the name into a buffer using bpf_path_d_path. Then we can do a substring search to see if the path is monitored or not. I will be sure to try this out after our next infra update.

Wrapping Up

This little experiment turned out to be a great deep dive into Linux kernel internals, eBPF and various trade-offs of running kernel-space programs. eBPF is a very powerful tool, but also has very sharp edges if you are not careful. This has also been my most rigorous exercise in RTFM'ing. A lot of information about these tools exists, but it’s scattered across kernel docs, blog posts, and reference guides. Piecing it all together was a journey in itself.

Intuition behind Power of 2 Choices Load balancing

Sun, 17 Aug 2025 00:00:00 GMT

One of the hardest parts of balancing load across many targets is keeping an accurate view of their load. We have to check all the targets; this is expensive. Also, the process of checking is not instantaneous, which leads to stale data which in turn causes request herding. What if we assign a request to a random target? This works surprisingly well but causes hotspots, which is not ideal. A variant of this is where we select 2 targets at random and then assign the request to the one with less load, which works even better. This is called the Power of 2 Choices balancing.

This has been studied in Maths as the Balls & Bins problem. Where we place balls into bins, where a bin can house any number of balls. From here, we get that the typical max load on a target is exponentially better with 2 choices, i.e. \mathcal{O(\frac{\log n}{\log\log n})} for random and \mathcal{O(\frac{\log\log n}{\log d})} for d random choices, for 2 this will be \mathcal{O(\log \log n)}. From this result, we can also see that going from 1 to 2 improvement is far better than selections of 3 or 4, which are only marginally better. But it is not obvious intuitively, since we are still selecting the two targets randomly. This same exponential gap underlies why data structures like cuckoo hashing work so well: having two random choices dramatically reduces collisions and spreads keys more evenly.

Here is a visualisation of how the Power of 2 choices approach perform better, as you can see in the gif, the load appears more uniform as the number of requests increases.

If you want to go deeper into the practical side of these tradeoffs, Tyler McCullen’s talk "Load Balancing is Impossible" and Mark Booker's Blog Post are excellent resources.

I recently found a UWaterloo slide about this that had a very intuitive explanation of this. Say in our servers, N_k servers are already at max load out of a total n servers, say max load is k. The probability that the max load grows is just N_k/n since we can only choose one target. But when we get to choose two targets, what is the probability that the max load grows? What would it take to make k+1? We will have to choose two targets that are already at k, since if only one is at k, the other target will be selected, and max will not increase. Now, what is the probability that two targets are at max and we choose both of them is (N_k/n)^2.

Here, N_k would already be small. Furthermore to increase the max even more from k+1 to k+2, the number of servers with max load would have fallen even more, call it N_{k+1}, the probability will be (N_{k+1}/n)^2 which is event tinier than moving from k to k+1. This is why the tail of max load in case of 2 choices falls very rapidly, since with every iteration the probability of increasing the max falls faster than the single choice method.

Let's go through with an example

Say we have n/4 targets with a max of 4 requests each, the probability of selecting two of these is 1/16.
Now we should expect only n/16 targets to have the max 5 requests, and then only n/256 = n/(2^{2^3}) targets with max 6 requests
This amounts to \frac{n}{2^{2^{k-3}}} for a max of k requests.
To find the upper bound of k at a fixed n we can set the N_k to the minimum 1. This will give us k = \mathcal{O(\log \log n)}

Exploring Clojure for Advent of Code

Wed, 05 Jan 2022 00:00:00 GMT

Very early in my career I had the absolute delight of reading SICP ; this was my first introduction Lispy languages. I used Racket to go through the book at the time, and found it overall enjoyable to use. Although, I found the all the parentheses to be an acquired taste. Eventually could visually cut through the parentheses.

Clojure was one other language I wanted to explore in this genre. Aside form being Lispy, some of its selling points that intrigued me were that a) it is targeted to compiled to JVM bytecode b) had persistent data structures; that can change with mutating the original memory (by creating new structures efficiently).

I gave it a go for Advent of Code 2021. Here are my solutions, github/aoc2021. My overall experience was again also good, here are some of the language features I found very fun to use.

Threading Macros

These are a very convenient way to compose functions. I like the ergonomics of using them over the manual way of composing (c (b (a x))) which is difficult to understand when going through the program. There are two of them in Clojure; thread-first -> and thread-last ->> owing to the fact that the former threads into the first argument of each function and the latter the last argument.

Here is an example for file parsing. I also appreciate, when languages have simple IO APIs when sometimes 10-20% of AoC problem code is just input parsing.

(ns aoc.shared
  (:require [clojure.java.io :as io]
            [clojure.string :as str]))

(defn read-lines [file-resource]
  (->> (io/resource file-resource)
       (slurp)
       (str/split-lines)))

This was also my most overused feature when solving AoC since there is a huge correlation between composition and breaking a problem into smaller pieces, for example in Day 14 Problem, we apply a function 11 times and then calculate the difference of min frequency and max frequency of the characters of the resulting string. Complete Solution

;Part-1
(->> template
     (iterate apply-subs)
     (take 11)
     (last)
     (frequencies)
     (vals)
     (min-max-diff))

Collections

Clojure collections are persistent data structures, when mutating them we get a new structure, that may or may not share memory with the original. They provide \approx\mathcal{O}(log_{32}N) operations. They are based on Hash Array Mapped from Phil Bagwell's 2001 paper Ideal Hash Trees.¹. Here is an implementation Dijkstra's shortest path algorithm for Day 15. Note the assoc operator, this would be mutating key node in a map cost with curr-dist, but here you get the appearance of a returning a new map with the changed value. Complete Solution.

(ns aoc.day15
  (:require [clojure.data.priority-map :refer [priority-map]]))

(defn dijkstra [graph start-coords h w]
  (loop [q (priority-map start-coords 0)
         costs {}]
    (if (empty? q)
      costs
      (let [[node curr-dist] (peek q)
            dist (->> (valid-neighbors node h w)
                      (filter (complement costs))
                      (map #(vector % (+ (graph %) curr-dist)))
                      (into {}))]
        (recur
         (merge-with min (pop q) dist)
         (assoc costs node curr-dist))))))

Zippers

Input for Day 18 are nested vectors, which are tree like. This problem also wants us to edit the tree structure. Seems like a good place to use the Zipper APIs. A good deep dive for zippers is provided by this Ivan Grishaev post. But in short: a zipper is a data structure that represents not just a tree, but also a cursor into that tree; a focused position plus all the context needed to rebuild the whole structure after changes. In Clojure, (clojure.zip/vector-zip tree) turns a nested vector into such a navigable structure. Each zipper location contains/supports:

the current node (zip/node),
information about its siblings and parent path, and
functions to move (zip/down, zip/up, zip/left, zip/right, zip/next, zip/prev) or edit in place (zip/edit, zip/replace).

Because all of this is purely functional, edits return new zipper locations instead of mutating the original tree. You can freely descend, modify, and then climb back to the top with zip/root to recover the updated structure; exactly what we need for repeated tree rewrites in the problem.

Navigating the tree

Lets add a few helpers deal with moving through the nested vector and finding specific locations (like leaves or the root).

(ns aoc.day18
  (:require [clojure.zip :as z]
            [clojure.walk :as walk]))

(defn leaves-seq [loc step]
  (->> (iterate step loc)
       rest
       (take-while (complement z/end?))
       (remove z/branch?))) ;filter out branches

(defn next-leaves [number] (leaves-seq number z/next))
(defn prev-leaves [number] (leaves-seq number z/prev))

(defn root-loc [loc]
  (->> (iterate z/up loc)
       (take-while identity)
       last))

leaves-seq walks a zipper in depth-first order, collecting only the leaf nodes (the regular numbers, not pairs). next-leaves and prev-leaves specialize this for scanning forward or backward. Meanwhile, root-loc ensures we can always “rewind” to the top of the tree after edits since zipper edits keep you at the modified node.

Exploding deeply nested pairs

Exploding is the most complex operation: when a pair is nested inside four pairs, it “explodes” its left and right values are distributed to the nearest regular numbers on the left and right, and it’s replaced by 0.

(defn explode [number]
  (if-let [explode-loc
           (->> number
                next-leaves
                (filter #(> (count (z/path %)) 4))
                first
                z/up)]
    (let [[left-val right-val] (z/node explode-loc)
          explode-loc (z/replace explode-loc 0)
          explode-loc (if-let [left-loc (first (prev-leaves explode-loc))]
                        (-> left-loc (z/edit + left-val) next-leaves first)
                        explode-loc)
          explode-loc (if-let [right-loc (first (next-leaves explode-loc))]
                        (-> right-loc (z/edit + right-val) prev-leaves first)
                        explode-loc)]
      [explode-loc :continue])
    [number :done]))

Here we locate the leftmost pair nested deeper than 4 levels ((> (count (z/path %)) 4)). Once found, we destructure its values, replace the pair with 0, and then use prev-leaves and next-leaves to locate and increment the nearest neighboring numbers. Everything is done immutably, z/edit returns a new zipper each time.

To run explosions repeatedly until stable, we wrap it in full-explode:

(defn full-explode [number]
  (->> [number :continue]
       (iterate (fn [[number _]] (explode number)))
       (map-indexed vector)
       (filter (fn [[_ [_ state]]] (= state :done)))
       first
       ((fn [[i [number _]]]
          [number (if (> i 1) :changed :done)]))))

This repeatedly applies explode until it signals :done, then reports whether anything changed.

Reducing to a stable form

To fully normalize a snailfish number, we must apply explosions and splits repeatedly, explosions first, then splits, until neither changes the tree.

(defn normalize [number]
  (let [[exploded e-state] (full-explode number)
        [split-n s-state] (split (root-loc exploded))]
    (if (= :done e-state s-state)
      (z/root split-n)
      (recur (root-loc split-n)))))

This recursive loop ensures the number reaches its reduced form. We always start each pass from the root-loc, so the next traversal covers the entire tree correctly. z/root extracts the final value once normalization is complete.

Combining numbers and Magnitude

Adding two snailfish numbers simply wraps them in a new vector and normalizes:

(defn add [n1 n2]
  (normalize (z/vector-zip [n1 n2])))

Finally, we compute the magnitude; a recursive formula: magnitude([x, y]) = 3 * magnitude(x) + 2 * magnitude(y).

(defn magnitude [number]
  (walk/postwalk
   (fn [node]
     (if (number? node)
       node
       (+ (* 3 (first node)) (* 2 (second node)))))
   number))

clojure.walk/postwalk is perfect here, it processes the tree bottom-up, collapsing pairs into scalar magnitudes along the way.

In Part 1, we sum all snailfish numbers in sequence and compute the final magnitude:

(magnitude (reduce add input))

Each call to add wraps, normalizes, and returns the reduced result, just like the problem statement. The zipper operations handle all the nested mutations elegantly, keeping the code declarative and readable. Complete Solution

Chinks in the Armor?

So far I have been gushing praises for the language; but there are a lot of sharp edges. Errors have been a big pain for me and the Clojure community at large. Another one for me is the function discoverability, like how I am supposed to know remove exists? Clojure Docs help but this is not great; since all these are just clubbed in clojure.core.

Aside from these Clojure has been a joy to work with and I look forward to using it more in my projects. I have not yet explored a lot of tooling around dependency management, but I have seem people on the web complaining about it.

Footnotes

Some great resources to understand the inner workings Understanding Clojure's Persistent Vectors, pt. 1 and later posts from Jean Niklas ↩

Making Machine Learning models smaller

Tue, 10 Mar 2020 00:00:00 GMT

WARNING

Contents of this post were published in 2020, while some concepts are still relevant, most of provided code might not work due to framework changes.

Machine Learning models are getting bigger and expensive to compute. Embedded devices have restricted memory, computation power and battery. But we can optimize our model to run smoothly on these devices. By reducing the size of the model we decrease the number of operations that need to be done hence reducing the computation. Smaller models also trivially translate into less memory usage. Smaller models are also more power-efficient. One must think that a reduced number of computations is responsible for less power consumption, but on the contrary, the power draw from memory access is about 1000x more costly than addition or multiplication. Now since there are no free lunches i.e. everything comes at a cost, we lose some accuracy of our models here. Bear in mind these speedups are not for training but inference only.

Pruning

Pruning is removing excess network connections that do not hugely contribute to the output. Ideas of pruning networks are very old dating back to 1990s namely Optimal Brain Damage and Optimal Brain Surgeon. These methods use Hessians to determine the importance of connections, which also makes them impractical to use with deep networks. Pruning methods use an iterative training technique i.e. Train ⇒ Prune ⇒ Fine-tune. Fine-tuning after pruning restores the accuracy of the network lost after pruning. One method is to rank the weights in the network using the L1/L2 norm and remove the last x% of them. Other types of methods which also use ranking use the mean activation of neurons, the number of times a neuron's activation is zero on a validation set and many other creative methods. This approach is pioneered by Han et.al. in their 2015 paper.

Fig 1. Pruning in neural networks from Han et. al.

Even more recently in 2019, the Frankle et.al. paper titled The Lottery Ticket Hypothesis the authors found out that within every deep neural network there exists a subset of it which gives the same accuracy for an equal amount of training. These results hold for unstructured pruning which prunes the whole network which gives us a sparse network. Sparse networks are inefficient on GPUs since there is no structure to their computation. To remedy this, structured pruning is done, which prunes a part of the network e.g. a layer or a channel. The Lottery Ticket discussed earlier is found no to work here by Liu et.al. They instead discovered that it was better to retrain a network after pruning instead of fine-tuning. Aside from performance is there any other use of sparse networks? Yes, sparse networks are more robust to noise input as shown in a paper by Ahmed et.al. Pruning is supported in both TF (tensorflow_model_optimization package) and PyTorch (torch.nn.utils.prune).

To use pruning in PyTorch you can either select a technique class from torch.nn.utils.prune or implement your subclass of BasePruningMethod.

from torch.nn.utils import prune
tensor = torch.rand(2, 5)
pruner = prune.L1Unstructured(amount=0.7)
pruned_tensor = pruner.prune(tensor)

To prune a module we can use pruning methods (basically wrappers on the classes discussed above) given in torch.nn.utils.prune and specify which module you want to prune, or even which parameter within that module.

conv_1 = nn.Conv(3, 1, 2)
prune.ln_structured(module=conv_1, name='weight', amount=5, n=2, dim=1)

This replaces the parameter weight with the pruned result and adds a parameter weight_orig that stores the unpruned version of the input. The pruning mask is stored as weight_mask and saved as a module buffer. These can be checked by the module.named_parameters() and module.named_buffers(). To enable iterative pruning we can use just apply the pruning method for the next iteration and it just works, due to PruningContainer as it handles computation of final mask taking into account previous prunings using the compute_mask method.

Quantization

Quantization is to restrict the number of possible values a weight can take, this will reduce the memory a weight can reduce and in turn reduce the model size. One way of doing this is changing the bit-width of the floating-point number used for storing the weights. A number stored as a 32-bit floating-point or FP32 to an FP16 or an 8-bit fixed-point number and more increasingly an 8-bit integer. Bit width reductions have many advantages as below.

Moving from 32-bit to 8-bit gives us a 4x memory advantage straight away.
Lower bit width also means that we can squeeze me more numbers in registers/caches with leads to less RAM access and in-turn less time and power consumption.
Integer computation is always faster than floating-point ones.

This works because neural nets are pretty robust to small perturbations to their weights and we can easily round off them without having much effect on the accuracy of the network. Moreover, weights are not contained in very large ranges due to regularization techniques used in training, hence we do not have to use large ranges, say ~ −3.4×10^{38} to 3.4×10^{38} for a 32-bit floating -point. For example, in the image below the weight values in MobileNet are all very close to zero.

Fig 2. Weight distribution of 10 layers of MobileNetV1.

A Quantization scheme is how we transform our real weights to quantized one, a very rudimentary form of the scheme is linear scaling. Say we want to transform values in range [r_{min},r_{max}] to an integer range of [0, I_{max}], where I_{max} is 2^B -1 , 𝐵 being the bit-width of our integer representation. Hence,

r = \frac {r_{max} - r_{min}} {I_{max} - 0} {(q - z)} = s(q -z)

where r is the original value of the weight, 𝑠 is the scale, 𝑞q is the quantized value and 𝑧 is the value that maps to 0.0f. This is also known as an affine mapping. Since 𝑞 is integer results are rounded off. Now the problem arises how we choose r*{min} and r*{max}. A simple method to achieve this is generating distributions of weights and activations and then taking their KL divergences with quantized distributions and use the one with min divergence from the original. A more elegant way to do this is using Fake Quantization i.e. introduce quantization aware layers into the network during training. This idea is proposed by Jacob et. al..

Fig 3. (a) Normal conv layer, (b) Conv layer with fake quantization units added, (c) Comparison of quantized network's latency and accuracy. Image from Jacob et.al.

While training the Fake quantization node calculates the ranges for the weights and activations and store their moving average. After training, we quantize the network with this range to get better performance.

More drastic bit-width also explored in papers on XOR nets by Rastegari et.al, Ternary nets by Courbariaux et. al. and Binary nets by Zhu et. al. In PyTorch 1.3, quantization support was introduced. Three new data types are introduced for quantization operations torch.quint8, torch.qint8 and torch.qint32. It also offers various qunatization techniques included in torch.quantization package.

Post Training Dynamic quantization : Replaces float weights with dynamic quantized versions of them. Weight-only quantization by default is performed for layers with large weights size - i.e. Linear and RNN variants.
```
quantized_model = torch.quantization.quantize_dynamic(
    model, {nn.LSTM, nn.Linear}, dtype=torch.qint8
)
```

Post Training Static quantization : Static quantization not only converts float weights to int, but it also records the distribution of activations and they are used to determine the scale of quantization at inference time. To support this calibration type quantization we add QuantStub at the start of the model and DeQuantStub and the end of the model. It involves steps mentioned below.

myModel = load_model(saved_model_dir + float_model_file).to('cpu')
# Fuse Conv, bn and relu
myModel.fuse_model()

# Specify quantization configuration
# Start with simple min/max range estimation and per-tensor
# quantization of weights
myModel.qconfig = torch.quantization.default_qconfig

torch.quantization.prepare(myModel, inplace=True)

# Calibrate with the training set
evaluate(myModel, criterion, data_loader,
            neval_batches=num_calibration_batches)

# Convert to quantized model
torch.quantization.convert(myModel, inplace=True)

Quantization Aware Training : Uses fake quantization modules to store scales while training. For enabling QAT, we use the qconfig to be get_default_qat_qconfig('fbgemm') and instead of prepare use prepare_qat. After this, we can train or fine-tune our model and at the end of the training, get out the quantized model using torch.quantization.convert same as above.

Post-training quantization in PyTorch currently only support operations on CPU.

For detailed code examples visit the PyTorch documentation here. On Tensorflow side of things quantization can be done using TFLite's tf.lite.TFLiteConverter API by setting the optimizations parameter to tf.lite.Optimize.OPTIMIZE_FOR_SIZE. Fake quantization is enabled by tf.contrib.quantize package.

Low Rank Transforms

Low-rank transform means representing a matrix or tensor as a product of some lower rank components. These components often only approximate the original matrix but benefit hugely in space and computational efficiency. For example, fully connected layers can be compressed using plain old Truncated SVD as a fully connected layer can be represented as a matrix. In truncated SVD matrix M of size 𝑛×𝑚 is approximated by \tilde{M} = U\Sigma V^T, where U is n×t, Σ is a diagonal matrix of size 𝑡×𝑡 and V is t×m in size. A fully connected layer can be represented as Wx+b, where W is the weight matrix and b are the biases. We now represent the FC layer as

(U \Sigma V^T x) + b = U (\Sigma V^T x) + b

hence we can split our FC layer into two;

The first layer with shape n×t, having no biases and weights taken from \Sigma V^T.
Second layer with shape t×m, original biases and weights from U. This drops the number of weights from n×m to t(n+m). Time complexity is also reduced by the same factor.

This can be easily implemented as in PyTorch using the torch.svd method as depicted in the code snippet below. Here vgg16 is a pre-trained model picked from torchvision.models. I have applied SVD on Linear layers after training. It can also be applied before training, but that involves calculating the gradient of the SVD operation which is a hassle.

svd_classifier_layers = []
L = 50
for layer in vgg16.classifier:
  if isinstance(layer, nn.Linear):
    W = layer.weight.data
    U, S, V = torch.svd(W)
    W1 = U[:,:L]
    W2 = torch.diag(S[:L]) @ V[:,:L].t()
    layer_1 = nn.Linear(in_features=layer.in_features,
                        out_features=L, bias=False)
    layer_1.weight.data = W2
    svd_classifier_layers.append(layer_1)

    layer_2 = nn.Linear(in_features=L,
                        out_features=layer.out_features, bias=True)
    layer_2.weight.data = W1
    layer_2.bias.data = layer.bias.data
    svd_classifier_layers.append(layer_2)
  else:
    svd_classifier_layers.append(layer)

svd_vgg16.classifier = nn.Sequential(*svd_classifier_layers)

This results in size reduction from 528MB to 195M i.e. ~ 2.7x decrease.

This works so well works because the

majority of the weights in a VGG16 are in Fully Connected layers. For much newer network e.g. ResNets majority of the weights lie in the Conv layers, therefore it makes more sense to apply Low rank transforms to Conv layers. Since conv layers are 4D tensors i.e (batch, channels, width, height), SVD and its cousins will not work here. We need to apply specialized tensor decomposition techniques such as CP decomposition (Lebedev et.al. in 2015) and Tucker Decomposition (Kim et. al. in 2016). Not covering these papers in more detail because these techniques are now superseded by efficient architectures like SqueezeNet and MobileNet which are discussed in the next section.

Efficient network architectures

Rather than applying size reducing techniques to existing architectures, we try to create novel architectures that decrease the model size and try to preserve the accuracy of the network over the time there have been many such architectures, prominent of them being SqueezeNet, MobileNet V1 and MobileNet V2.

SqueezeNet

SqueezeNet by Iandola et.al. is presumably the first to explore a new architecture for smaller CNNs. At the core of SqueezeNet are Fire Modules. Fire modules use 1x1 filters rather than 3x3 filters as they have 9x lesser parameters and have a lesser number of channels than normal, which is called a squeeze layer. The lesser number of channels are recovered in the expand layer which consists of several zero-padded 1x1 filters and 3x3 filters. The number of filters in the squeeze layers and expand layers are hyper-parameters. If 𝑒*{3×3}+𝑒*{1×1} are the number of filters in expand layer and s*{1×1} is the number of filters in the squeeze layer. When using Fire module s*{1×1} < e*{3×3}+e*{1×1} works best.

Fig 4. Fire module with s_{1×1}=3, e_{1×1}=4 and e_{3×3}=4. (Source)

Code for the Fire Module adapted from torchvision.models. Here inchannels are the number of input channels, squeeze_planes are the number of output channels, expand1x1_planes and expand3x3_planes are the output channel number for the expand layer. They are generally same.

class Fire(nn.Module):

    def __init__(self, inplanes, squeeze_planes,
                 expand1x1_planes, expand3x3_planes):
        super(Fire, self).__init__()
        self.inplanes = inplanes
        self.squeeze = nn.Conv2d(inplanes, squeeze_planes, kernel_size=1)
        self.squeeze_activation = nn.ReLU(inplace=True)
        self.expand1x1 = nn.Conv2d(squeeze_planes, expand1x1_planes,
                                   kernel_size=1)
        self.expand1x1_activation = nn.ReLU(inplace=True)
        self.expand3x3 = nn.Conv2d(squeeze_planes, expand3x3_planes,
                                   kernel_size=3, padding=1)
        self.expand3x3_activation = nn.ReLU(inplace=True)

    def forward(self, x):
        x = self.squeeze_activation(self.squeeze(x))
        return torch.cat([
            self.expand1x1_activation(self.expand1x1(x)),
            self.expand3x3_activation(self.expand3x3(x))
        ], 1)

SqueezeNet also uses delayed-sampling to create larger activation maps towards the end layers, which in turn leads to greater accuracy. The full architecture can be visualized here.

MobileNets

MobileNets are specifically developed by Google to specifically run on mobile devices. MobileNets were first introduced in a paper by Howard et.al. in 2017, subsequently, in 2018 an improved version was introduced called MobileNet v2 in Sandler et. al.. The gist of optimization in MobileNet v1 lies in a special kind of convolution layer called Depthwise separable convolutions. For a simple convolution layer if 𝑘 is the dimension of the kernel, 𝑁𝑘 is the number of kernels, and the input is of size 𝑁_c×𝑊×𝐻, where 𝑁_𝑐 are the number of input channels. The total number of parameters and computations are k^2N_kN_cWH. MobileNet Convolutions work in two stages

Convolve a k×k for each channel of the input and stack N_c of them, creating an output tensor of size N_c×W×H. Total number of ops in this layer is k^2N_cWH
Convolve with a 1×1 filter with N_k channels to create the final output. Total number of computations in this stage is N_cN_kWH

Total computations in a MobileNet convolution are k^2N_cWH+N_cN_kWH. There total reduction in parameters in given by,

\frac {k^2 N_c W H + N_c N_k WH} {k^2 N_k N_c WH} = \frac {1} {N_k} + \frac {1} {k^2}

For k=3, 𝑁_k=16 we have a ~ 5.76x reduction in number of parameters for a layer.

Fig 5. Depthwise seperable convolution followed by pointwise convolutions (Source)

Implementing Depthwise conv. is quite simple. Checkout the code snippet below, inp donates the number of input channels and oup are the number of output channels.

def conv_dw(inp, oup, stride):
    return nn.Sequential(
        nn.Conv2d(inp, inp, 3, stride, 1, groups=inp, bias=False),
        nn.BatchNorm2d(inp),
        nn.ReLU(inplace=True),

        nn.Conv2d(inp, oup, 1, 1, 0, bias=False),
        nn.BatchNorm2d(oup),
        nn.ReLU(inplace=True),
    )

MobileNet v2 uses as an inverted residual block as its main convolutional layer. A Residual block taken from the ResNets includes bottleneck layers that decrease the number of channels followed by an expansion layer that restores the number of channels for the residual concat operation. The inverted block layer does the reverse of that it first expands the number of channels then reduce them. The last layer in the block is a bottleneck layer as it decreases the channels of the output. This layer has to non-linearity attached to it. This because authors found out that a linear bottleneck does not lose information when a feature-map is embedded into a lower dimension space i.e. reduced to a tensor with less number of channels. This is found to increase the accuracy of these networks. To calculate the number of parameters, presume N_{in} is the number of input channels, N_{out} the number of output channels and 𝑡 is the expansion ratio, the ratio between the size of the intermediate layer to the input layer. The number of computations and parameters are WHN_{in}t(N_{in}+k^2+N_{out}). But there is an extra 1x1 convolution component, still, we have a computational advantage because due to the nature of the block we can now decrease the input and output dimensions e.g. a layer with dimensions 112x112 can have only 16 channels and retaining the accuracy as compared to 64 for MobileNet v1.

Fig 6. MobileNet v2 primary convolution block.(Source)

The code for the InvertedResidual block is adapted from trochvision.models package.

class InvertedResidual(nn.Module):
    def __init__(self, inp, oup, stride, expand_ratio):
        super(InvertedResidual, self).__init__()
        self.stride = stride
        assert stride in [1, 2]

        hidden_dim = int(round(inp * expand_ratio))
        self.use_res_connect = self.stride == 1 and inp == oup

        layers = []
        if expand_ratio != 1:
            # pw
            layers.append(ConvBNReLU(inp, hidden_dim, kernel_size=1))
        layers.extend([
            # dw
            ConvBNReLU(hidden_dim, hidden_dim, stride=stride,
                       groups=hidden_dim),
            # pw-linear
            nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
            nn.BatchNorm2d(oup),
        ])
        self.conv = nn.Sequential(*layers)

    def forward(self, x):
        if self.use_res_connect:
            return x + self.conv(x)
        else:
            return self.conv(x)

Knowledge Distillation

Knowledge Distillation (KD) is a model compression technique by which the behaviour of a smaller (student) model is trained to replicate the behaviour of a larger (teacher) model. The term was first coined by none other than Geoffrey Hinton in his 2015 paper. KD involves training a smaller network on the weighted average of soft target output of the larger model and the ground truth. Soft target output can be obtained by calculating the softmax on the logits of the larger model, but this creates large divides between the probabilities of the correct label and the wrong label, thus not creating much information other than the ground truth. To remedy this problem Hinton introduces softmax with temperature given by

q_i = \frac {exp(\frac {z_i} T)} {\sum_j exp(\frac {z_j} T)}

where T the temperature parameter, T=1 gives the same result as a simple softmax. As Tgrows the probabilities grow softer, providing more information about the model. The overall loss function of the now student-teacher pair becomes

\mathcal{L} = \lambda \mathcal{L}_{gt} + (1 - \lambda) \mathcal{L}_{temp}

where \mathcal{L}_{gt} is the loss with ground truth outputs and \mathcal{L}_{temp} is the softmax temperature loss. Both \lambda and 𝑇are tunable hyperparameters. The loss configuration is as in the image below.

Fig 7. Knowledge distillation model configuration. (Source)

A major success story of KD is DistillBERT. Hugging Face managed to use KD to reduce the size of the BERT from 110M parameters to 66M parameters, while still retaining 97% of the performance of the original model. DistillBERT uses various additional tricks to achieve this such as using KD loss instead of standard cross-entropy to retain the probability distribution of the teacher model. The code to train a KD model will go like below. This code is adapted from DistilBERT training sequence itself.

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.optim import Optimizer

celoss = nn.CrossEntropyLoss
lambda_ = 0.5

def kd_step(teacher: nn.Module,
            student: nn.Module,
            temperature: float,
            inputs: torch.tensor,
            optimizer: Optimizer):
    teacher.eval()
    student.train()

    with torch.no_grad():
        logits_t = teacher(inputs=inputs)
    logits_s = student(inputs=inputs)

    loss_gt = celoss(input=F.log_softmax(logits_s/temperature, dim=-1),
                     target=labels)
    loss_temp = celoss(input=F.log_softmax(logits_s/temperature, dim=-1),
                       target=F.softmax(logits_t/temperature, dim=-1))
    loss = lambda_ * loss_gt + (1 - lambda_) * loss_temp

    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

There are many more methods to make ML models smaller which I cannot cover as the posts would become too long. More and more research is being done on this, to follow the research be sure to check to arixv-sanity. Will try to introduce a further reading section in future.

amandeep singh

The Boon of Dimensionality

Volume of a high-dimensional ball

Equators

Near Orthogonality

The Johnson-Lindenstrauss Lemma

Where next

Footnotes

Twelve Attempts at an FP4 Kernel

The Task

Optimization Journey

CuTe Python DSL (Attempts 1-4)

Split-K with Atomics in CuTe (Attempt 5)

The Switch to Raw CUDA (Attempt 7)

The Failed Experiments

Why 3x Off?

What the Top Solutions Did Differently

Takeaways

Footnotes

Honey, I Tiled the Tensors

Coalescing

Composition

Single-mode B

Single-mode A

General composition: left-associative reduction

Complement

Division

Zipped, Tiled and Flat Divides

Product

Blocked and Raked Products

Zipped and Tiled Products

Putting it all together

Setup and Tiled Copies

MMA Setup

The Kernel

The Main Loop

Epilogue

Closing thoughts

Footnotes

Hierarchical Navigable Small Worlds

Delaunay Triangulation

Small Worlds and their Navigation

Intuition Behind \alpha = d and Runtime Cost

Hierarchical Navigable Small Worlds5

Construction

Neighbor Selection Heuristic

Intuition Behind Runtime Cost

Implementation

Footnotes

The Rabbit Hole of Building a Filesystem Watcher

Attempt 1 - fanotify

Attempt 2 - eBPF

Walking up the file tree in eBPF

Better Probes

Wrapping Up

Intuition behind Power of 2 Choices Load balancing

Exploring Clojure for Advent of Code

Threading Macros

Collections

Zippers

Navigating the tree

Exploding deeply nested pairs

Reducing to a stable form

Combining numbers and Magnitude

Chinks in the Armor?

Footnotes

Making Machine Learning models smaller

Pruning

Quantization

Low Rank Transforms

Efficient network architectures

SqueezeNet

MobileNets

Knowledge Distillation

Single-mode `B`

Single-mode `A`

Intuition Behind `\alpha = d` and Runtime Cost

Hierarchical Navigable Small Worlds⁵

Attempt 1 - `fanotify`