DKM: Differentiable
k-Means Clustering for Compression

Chapter 1: The Key Insight

Classic weight-clustering (DeepCompression) runs k-means once or periodically: it groups weights by nearest centroid, replaces each weight with its centroid, then fine-tunes. The assignment step — "weight w_i belongs to cluster c_j" — is a hard, discrete decision based purely on Euclidean distance. Discrete means non-differentiable. Non-differentiable means the network's task loss cannot reach back and say "actually, that boundary weight should have gone to the other cluster."

The one-line reframe

Treat each weight as a query, each centroid as a key and a value. The "attention score" is how close the weight is to the centroid. The clustered weight is the attention-weighted average of centroids — exactly the standard attention readout. That is the entire paper in one sentence.

Chapter 2: Background — k-means, Precisely

Before we soften it, let's be crisp about hard k-means, because DKM is k-means with exactly one operation swapped. We have N scalar weights W = \{w_1, ..., w_N\} and want k centroids C = \{c_1, ..., c_k\}. Lloyd's algorithm alternates two steps until convergence:

Every symbol, in plain words:

• w_i — one weight value (a single number for now; a sub-vector in Ch 7). Think: a marble on a number line.
• c_j — one centroid: a shared value many weights will round to. Think: a magnet on the line.
• \text{argmin}_j — pick the index of the closest magnet. This is the hard, discrete step.
• Update — slide each magnet to the average of the marbles stuck to it.

The data flow we're about to compress

Concretely, for one convolution layer we flatten the kernel tensor into a long vector of weights and feed it to a DKM "layer" that returns clustered weights of the same shape:

Marbles = weights, colored by their hard nearest centroid. Triangles = centroids. Step to assign-then-update. Notice marbles near a boundary flip color in a single discrete jump — there is no in-between. Hold that thought for Chapter 3.

Chapter 3: The Boundary Bug

Here is the failure that motivates the whole paper, drawn straight from Figure 1. Consider two boundary weights, call them i and j, sitting almost exactly between two centroids. Hard k-means assigns each to its marginally nearer magnet — i\to c_2, j\to c_1 — purely on distance.

But "nearest in Euclidean distance" is not the same as "best for the task." Maybe sending i\to c_0 would barely change the rounding error but would let the network keep a decision boundary it needs. Distance can't know that. Only the task loss knows that — and in hard k-means the loss has no way to express its preference, because the argmin wall blocks the gradient.

The second, sneakier bug: the gradient hack

Classic methods have a second problem. In the backward pass, the weights are tied to a centroid, so what gradient does a weight get? DeepCompression simply re-purposes the centroid's gradient as the weight's gradient — an ad-hoc substitution with no principled justification. It's a straight-through-estimator-style fudge. DKM has no such fudge: the gradient of a weight is exactly the product of attention values and centroid gradients, because the forward pass is genuinely differentiable.

Two centroids, one weight you drag between them. The bars show the hard one-hot assignment (snaps abruptly at the midpoint) vs the soft DKM attention (a smooth blend). Hard assignment is a step function — flat then a cliff. Soft is a ramp the loss can climb.

Chapter 4: k-means as Attention — The Math

Now we derive the soft assignment from scratch. We need a differentiable stand-in for "pick the nearest centroid." The argmin over distances has a famous smooth cousin: the softmax over negative distances. As distances shrink (closer), the negative distance grows (larger score), so softmax puts more weight there — and in the limit it becomes a one-hot at the nearest. This is the standard "soft argmax."

Step 1 — distance as an attention score

For every weight–centroid pair, define a score that is large when close:

The minus sign is the whole game: nearest centroid ⇒ smallest distance ⇒ largest (least-negative) score. f is any differentiable metric; the paper uses Euclidean.

Step 2 — softmax with temperature = soft assignment

Here a_{ij} is "how much weight i attends to centroid j," each row sums to 1, and \tau is the temperature — the softness knob we dissect in Chapter 6. Small \tau → spiky, near one-hot (almost hard k-means). Large \tau → flat, every centroid shares the weight.

Step 3 — the centroid update is also attention

Hard k-means updates a centroid to the mean of its members. The soft version is an attention-weighted mean, where weights that attend more pull harder:

The denominator \sum_i a_{ij} is the soft membership count of cluster j — the fractional number of weights leaning toward it. Compare to hard k-means where the count is an integer.

Worked numerical example (do it by hand)

One weight w=0.5, two centroids c_1=0.0, c_2=1.0, temperature \tau=0.5. Distances: d_1=-|0.5-0|=-0.5, d_2=-|0.5-1|=-0.5. Both equal ⇒ perfectly balanced:

Clustered weight: \tilde w = 0.5\cdot 0 + 0.5\cdot 1 = 0.5 — it reconstructs perfectly because it sits dead center. Now nudge w=0.6: d_1=-0.6, d_2=-0.4, scores e^{-1.2}=0.301, e^{-0.8}=0.449, so a_1=0.40, a_2=0.60 and \tilde w=0.60. The assignment continuously tilts toward c_2 — and that continuity is exactly what carries a gradient.

Chapter 5: The DKM Loop — Showcase

This is Figure 2 made interactive. A DKM layer doesn't just do one assignment — it runs the soft assign/update loop inside the forward pass until the centroids stop moving (|C-\tilde C|\le\epsilon) or it hits the iteration limit (the paper uses \epsilon=10^{-4}, limit 5). Then it produces the clustered weights \tilde W = A\cdot C for the rest of the network — and crucially, backprop runs through every iteration of this loop.

Play with the live DKM layer

Below is a real soft k-means loop on a layer of weights. Press Run loop to watch the attention reassign and the centroids glide to their attention-weighted means. Drop the temperature toward 0 and it behaves like hard k-means (centroids snap to cluster means); raise it and centroids collapse toward the global mean. Increase iterations to see convergence; change k to add codebook entries.

Each weight (dot, top row) draws faded lines to the centroids (triangles, bottom row); line opacity = attention a_ij. After each iteration the centroids slide to their attention-weighted mean. The status line reports total centroid movement ΔC — when it drops below ε, the loop has "converged" and W̃ = A·C is read out.

The whole loop in runnable PyTorch

This is the paper's core algorithm — a parameter-free, differentiable layer. Drop it on a flattened weight tensor; autograd handles the backward pass through the loop automatically.

import torch, torch.nn.functional as F

def dkm(W, C, tau=1.0, iters=5, eps=1e-4):
    # W: [N] weights to compress.  C: [k] current codebook (warm-started).
    W = W.unsqueeze(1)                       # [N, 1]
    for _ in range(iters):
        D = -(W - C.unsqueeze(0)).abs()   # [N, k]  d_ij = -|w_i - c_j|
        A = F.softmax(D / tau, dim=1)       # [N, k]  row-softmax = attention
        C_new = (A * W).sum(0) / A.sum(0)  # [k]   attention-weighted means
        if (C_new - C).abs().max() < eps:   # converged?
            C = C_new; break
        C = C_new
    W_tilde = (A * C.unsqueeze(0)).sum(1)  # [N]  W̃ = A · C, used in forward pass
    return W_tilde, C                       # gradients flow through the whole loop

Note what's not here: no requires_grad on new parameters, no custom backward, no regularizer added to the loss. It's just differentiable ops, so the existing training loop and its gradients do all the work.

Chapter 6: Temperature — The Hardness Knob

Soft assignment is wonderful for training, but at inference you must commit: each weight is stored as a single index into the codebook (that's where the 2-bit savings live). So the soft attention must harden into a one-hot. The temperature \tau is the dial between these two worlds.

The two limits, derived

Look again at a_{ij}=\text{softmax}(d_{ij}/\tau). The ratio between the top two scores is \exp((d_{\text{best}}-d_{\text{2nd}})/\tau):

• \tau \to 0^+: the exponent blows up, the nearest centroid dominates utterly ⇒ one-hot = hard k-means. Differentiable in name only (gradients vanish — that's why you don't train here).
• \tau \to \infty: d_{ij}/\tau\to 0, all scores equal ⇒ uniform attention, every weight becomes the global mean. Too soft to cluster anything.
• In between: a usable gradient and a meaningful clustering signal.

Worked example — the hardening gap

Boundary weight w=0.55, centroids 0.0 and 1.0. At \tau=0.3: scores e^{-0.55/0.3}=0.160 vs e^{-0.45/0.3}=0.223 → a=(0.42,0.58), soft \tilde w=0.58; hardened → 1.0, error 0.45. At \tau=0.05: a\approx(0.13,0.87), soft \tilde w=0.87; hardened → 1.0, error 0.13. Colder ⇒ the soft value already lives near the hard target ⇒ smaller surprise at inference time.

Three centroids and one fixed weight. Bars = attention a_ij. Slide τ down: the distribution sharpens toward a one-hot on the nearest centroid (hard k-means). Slide up: it flattens toward uniform (the global mean). The dashed line marks the hardened value the inference engine will store.

Chapter 7: Multi-Dimensional DKM

So far each weight was a scalar clustered on a line. But you can cluster groups of weights jointly. Split the N weights into N/d contiguous d-dimensional sub-vectors, and cluster those points in d-D space. Centroids become d-D vectors (C\in R^{k\times d}); the distance metric runs in R^d. This is product-quantization-style codebook compression, now made differentiable.

Why bother? The entropy argument (Figure 3)

The compression ratio for b-bit, d-dim clustering versus 32-bit floats is:

Read that: compression depends only on the ratio d/b, not b or d alone. So 4/4 and 8/8 and 2/2 all give the same compression (1 effective bit-per-weight). But they are not equally good for accuracy. Assuming an even split, the per-weight entropy is:

Entropy e_w=b increases with b at fixed d/b. Higher weight entropy ⇒ more learning capacity ⇒ (often) better accuracy. So at a target compression ratio, push b and d up together as far as memory allows.

Why this would WRECK ordinary k-means

In high d, a sub-vector can be near the boundary of many centroids at once (16 centroids at 4/4 vs only 2 at 1/1). The chance of a damaging hard misassignment grows combinatorially — prior methods often fail to converge at multi-dim configs. DKM survives because soft assignment + task-loss feedback can recover from an early bad choice over training. The 8/8 MobileNet-v1 result (64.3% at 1 effective bit-per-weight) is the proof.

Each point is a 2-D weight sub-vector (d=2). Centroids (large rings) are 2-D codebook vectors. Soft attention colors each point as a blend of centroid colors. Run the loop and watch centroids migrate to attention-weighted 2-D means.

Chapter 8: Experiments — Reading the Results

DKM is benchmarked on ImageNet1k (vision) and GLUE (NLP). The headline pattern: DKM is Pareto-superior — for any given model size, it gets higher accuracy than DeepCompression (DC), HAQ, and "And The Bit Goes Down" (ATB). The win is biggest exactly where prior methods struggle: very low bits and already-compact models.

The numbers that matter

The MobileNet-v1 row is the famous one: 6.8% higher top-1 and 33% smaller than the prior state of the art on a model everyone considered "already maximally efficient." And the 8/8 multi-dim config beats the 1/1 config (64.3% vs ~50%) at the same effective bit-budget — direct evidence for Chapter 7's entropy argument.

Honest limitations the paper admits

• Train-time memory. Keeping all A,D across iterations for backprop is O(r\,N\,2^b) — high-bit configs OOM. Inference is cheap; training is the expensive part.
• Activations untouched. The paper deliberately disables activation quantization to isolate weight compression. Full deployment still needs an activation story.
• Per-layer config search. Choosing b/d per layer (cv:4/4, fc:8/4 …) is a small design space the user tunes — not fully automatic.

Top-right is best (small model, high accuracy). DKM points (orange) sit above and to the right of prior methods (teal) — the definition of a Pareto-superior frontier. Toggle models to see the gap widen on the harder-to-compress MobileNet.

Chapter 9: Connections & Cheat Sheet

DKM lives at the intersection of three ideas you already know: k-means (the clustering it softens), attention (the mechanism it borrows), and quantization (the compression goal it serves). Its deepest lesson is a transferable pattern: any time a discrete argmax/argmin blocks a gradient, a temperatured softmax can turn it into a trainable soft choice — the same move behind Gumbel-softmax, soft attention, and differentiable rendering.

Cheat sheet — every symbol, one place

The five things to remember

Where to go next

• Hard k-means & Lloyd's algorithm — the foundation DKM softens.
• Attention / softmax — the exact mechanism DKM reuses; the query/key/value mapping is one-to-one.
• Quantization & bit-widths — the storage model (codebook + indices) that turns clusters into MB savings.
• Gumbel-softmax & VQ-VAE — sibling techniques for differentiating through discrete choices.