Shrinking a Model with a Look-Up Table
JLCM: Jointly Learnable Codebooks & Mappings

Chapter 1: The Key Insight

Plain integer quantization rounds every weight onto an evenly-spaced grid (e.g. 256 levels for 8-bit). But weights are not evenly distributed — they pile up near zero with rare large outliers. An even grid wastes most of its levels in empty regions.

The codebook idea fixes this. Instead of a fixed grid, you keep a small table of learned representative values — the codebook C = (c₀, c₁, …, c_K−1) — and replace every weight with the index of its nearest table entry. The table can put its values exactly where the weights actually cluster. This is just k-means on the weights, also called vector quantization or weight sharing.

JLCM takes this old idea and removes its two blockers, which become the spine of this lesson:

Get those two right and you can push transformers to nearly ternary (3-value) weights with under 1% accuracy loss. That is the whole paper.

Chapter 2: LUT Basics — the memory math

Let us be exact about why a look-up table saves memory, with a worked example. Take one weight matrix W ∈ ℝ^n_o×n_i with Ω(W) total scalars. We notate the codebook size Ω(C)=K.

Baseline cost (fp16)

Every scalar stored at 16 bits: 16 × Ω(W) bits. For a 4096×4096 layer that is 16 × 16.78M ≈ 268 Mbit ≈ 33.5 MB.

Codebook cost

Two parts. (1) The table itself: K values at 16 bits = 16K bits — paid once. (2) The indices: each of the Ω(W) weights needs ⌈log₂K⌉ bits to name its table slot.

The teal term is the codebook (small, fixed). The warm term is the mapping (the indices — billions of them). For any real layer the warm term dominates utterly, because Ω(W) is in the millions and K is tiny. So the lever that controls memory is log₂K: the bits per index.

The compression goal α

The paper defines one knob: α = M_ref / M, the target compression ratio (reference fp16 footprint ÷ final footprint). Everything else — codebook size, number of codebooks — is derived from α. The recommended table size is

Read it: a bigger α (more compression) means fewer bits per index, hence a smaller table. At α=5.33 you get 2^⌊3⌋=8 entries (3-bit). At α=7.5 you get 2^⌊2.13⌋=4 entries (2-bit, nearly ternary). One hyperparameter sets the whole regime.

Left: the stored stream of indices (each a few bits). Middle: the shared codebook. Right: the reconstructed fp16 weight matrix. Click to decode one index at a time — every weight is just table[index]. Slide the bit-width: fewer bits = smaller table, coarser values.

Chapter 3: The Granularity Problem

One codebook for an entire weight tensor assumes all the weights share one distribution. They don't. Different output neurons (rows of W) often have wildly different scales — one row might live in [−0.01, 0.01], another in [−2, 2]. Force them through the same 8-entry table and the small-scale row gets crushed to near-zero garbage.

So you'd like finer granularity. Two natural fixes — and the paper studies both side by side.

The notation, to anchor it. For a single codebook the reconstruction is W_i,j = ⟨C; I_i,j⟩ — read ⟨C; I⟩ as "the entry of table C selected by one-hot index I." The cost was 16K + log₂(K)·Ω(W). Adding a separate codebook-id index pushes that to log₂(K·k_books)·Ω(W) — strictly worse. We need Option B's expressiveness at Option A's price. Chapter 4 is the trick that gets it.

Two rows of weights with different scales (blue = big-scale row, green = small-scale row), and the codebook levels (orange ticks). Widen the gap: with one shared table, the small-scale row's weights all snap to the same level — total information loss. Toggle to per-row codebooks and watch each row get levels that fit it.

Chapter 4: The Reordering Trick — free multi-codebooks

Here is the move that makes multiple codebooks free. If a weight's row index already tells you which codebook it uses, you don't need to store a codebook-id at all. The codebook index is implicit in position.

Formally, split the n_o rows into k contiguous blocks and assign codebook C_{⌊i·k/no⌋} to row i:

Read it: rows 0…(block boundary) use table 0, the next block uses table 1, and so on. The decoder computes the table id from the row number with a floor-divide — zero extra bits stored. You get Option B's expressiveness at Option A's price. Granularity problem solved.

How do you reorder without breaking the network?

This is the data-flow subtlety. You can't just shuffle one layer's output neurons — the next layer expects them in the original order. The fix: permute the output rows of layer l by a permutation σ, and apply the same permutation to the input columns of layer l+1. The two permutations cancel; the network computes exactly the same function. (This works cleanly when there's no skip connection in between — a documented assumption.)

The permutation σ itself comes from clustering the rows: treat each row (a length-n_i vector) as a sample and run a clustering method C on the n_o rows, then order rows so cluster-mates are contiguous. The final reconstruction folds the permutation in: W_i = ⟨ C_{⌊σi·k/no⌋} ; I_i ⟩.

Each row of the weight matrix is colored by its true distribution-type. Before reordering, types are scrambled — any contiguous block of rows is a mess, so a per-block codebook can't fit it. Click to cluster & reorder: now each block is one color, and one small codebook nails it. The dashed lines show the block→codebook boundaries.

Worked permutation cancellation

# Reorder output neurons of layer l, undo it on layer l+1's inputs.
# sigma: a permutation of the n_o output neurons (from clustering).
W_l_new  = W_l[sigma, :]        # permute ROWS  (output neurons)
b_l_new  = b_l[sigma]           # permute bias too
W_l1_new = W_l1[:, sigma]      # permute next layer's COLUMNS to match
# For input x:  W_l1_new @ act(W_l_new @ x + b_l_new)  ==  original output.
# The two permutations cancel: zero accuracy change, but now contiguous
# row-blocks share a distribution, so per-block codebooks are nearly free.

Chapter 5: Jointly Learning table + pointers

Initialization by clustering is a good start, but it minimizes the wrong thing — it minimizes weight reconstruction error, not output error. What you actually care about is that the compressed layer produces the same activations as the original on real inputs. So JLCM optimizes the codebook C and the mapping I jointly, layer by layer, on a small calibration set — the same self-distillation recipe as AdaRound/GPTQ.

The obstacle: an index is discrete

"Pick table slot 2" is not differentiable — you can't take a gradient through an argmax. The standard relaxation: replace the hard one-hot pick with a soft assignment, softmax(I), a probability vector over the K table slots. The reconstructed weight is then a soft blend C·softmax(I), which is differentiable. We push the softmax toward a hard one-hot over training so the final, stored mapping is a clean integer index.

The objective minimizes the gap between the compressed layer's output ƒ̃(X̃) and the original Y, plus two regularizers:

• Output term ‖ƒ̃(X̃)−Y‖²: the compressed layer should reproduce the real output on calibration data. ƒ̃ uses C·softmax(I) as its weights.
• L₁ = ‖C·softmax(I)−W‖²: keep the reconstruction near the original weights — an anchor that fights overfitting to the tiny calibration set.
• L₂ = 1−|2·softmax(I)−1|^β: a "be decisive" penalty. It is minimized when softmax outputs are near 0 or 1 — i.e. it drives the soft assignment toward a hard, single-slot pick, weighted by the single knob λ that is annealed up over training.

Data flow, with shapes

# One layer, calibration batch. Shapes for a 4096->4096 linear.
W      # (4096, 4096) fp16, frozen original
C      # (K,) learnable table, e.g. K=8
I      # (4096, 4096, K) learnable logits  ->  softmax over last axis
soft   = softmax(I, axis=-1)        # (4096,4096,K) soft assignment
W_hat  = (soft * C).sum(-1)        # (4096,4096) reconstructed weights
Y_hat  = act(X @ W_hat.T)          # X:(B,4096) -> (B,4096)
loss   = mse(Y_hat, Y) + mse(W_hat, W) + lam * decisiveness(soft)
# At export: index = argmax(I, -1)  ->  store as ceil(log2 K)-bit ints.

Note the cost: I is K× bigger than W during optimization (you keep logits for every slot of every weight). That is fine — it is transient, per-layer, on a calibration set; only the final argmax indices ship.

Chapter 6: The Proximal Gradient Fix

Why does naive joint learning fail? Walk through the gradient on the mapping I. By the chain rule, the gradient of the loss w.r.t. I contains the factor C · ∂softmax(I)/∂I. That C factor is the problem: the update to "move to slot j" is scaled by the value c_j. Large-magnitude table entries therefore attract the largest gradients — regardless of whether they are close to the current weight.

The fix: borrow the STE idea

Quantization's classic trick is the straight-through estimator (STE): the forward pass uses a non-differentiable op (rounding), and the backward pass swaps in a hand-designed gradient that behaves how we want. JLCM does the same for the index. It replaces the troublesome C·∂softmax/∂I term with a custom matrix D that is a decreasing function of distance between codewords:

where a = argmax(I_i) is the weight's current slot. Decode it slowly:

• The denominator 1 + |c_a − c_j|: as a candidate slot j gets farther in value from the current slot a, this grows, so the term shrinks. Near slots get strong pull; far slots get weak pull. This is the whole point — proximal search.
• The sign sign(c_a − c_j): points the update toward the neighbor on the correct side, so a positive gradient moves to the next-lower-value slot and vice versa.
• The identity subtraction: zeroes out the current slot's own entry so the weight is pushed to move, not to stay.

The codebook gradient (first row of the update) is left untouched — it was already well-behaved. Only the index gradient is surgically replaced. With this single change, joint learning finally works: Table 4 shows it lifting ResNet18 from 20.3% to 41.8% and ViT from 57.8% to 70.0% at the same compression.

The ternary codebook (−2, 0, 0.15, 1.3) on a number line; the white marker is the current weight. Bars show the gradient pull toward each slot. Naive: pull scales with |value|, so −2 dominates — a wild jump. Toggle to Proximal: pull scales with 1/(1+distance), so the nearest slot wins. Drag the weight and watch which slot the optimizer would jump to.

Chapter 7: Showcase — Compress a Layer End-to-End

Now put every piece together on one weight matrix. This is the JLCM pipeline you can drive: pick a compression goal α (which sets the table size K = 2^⌊16/α⌋), choose how many codebooks to split the rows into, optionally reorder, and watch the live memory budget, the reconstructed matrix, and the resulting error.

Left: original 16×16 weight matrix (color = value). Right: the reconstruction after compression. The readout shows live K, bits/index, the memory split (table vs indices), the achieved α, and reconstruction error. Toggle Reorder to group similar rows before splitting into per-block codebooks; toggle Proximal learn to run a few proximal-gradient refinement steps and watch the error drop. Push α high and the matrix gets blocky — that's the accuracy/size tradeoff made visible.

Chapter 8: Experiments & what they prove

The ablation (Table 4, α=7.5) is the cleanest evidence, because it isolates each contribution. Read it top to bottom — each row adds one component:

The story is dramatic. Learning only the codebook helps a little (+1.5 / +3.2). Naively adding the mapping hurts (the wild-gradient failure from Ch 6). The proximal gradient then nearly doubles ResNet18 accuracy. The headline number isn't joint learning — it's fixing the gradient so joint learning is possible.

Versus the field, and the payoff

On ImageNet at α≈5.33 (fp16→int3), JLCM beats data-free methods (DFQ, SQuant, PowerQuant) by large margins and is competitive with data-driven AdaRound/BrecQ/NUPES — while targeting memory, not latency. The flagship result: Llama 7B compressed ~7.2× fits in 2 GB, retaining ~95% of its common-sense-reasoning performance — enough to load on an iPhone 8 or Galaxy A10. A Stable Diffusion model holds its CLIP score under 4× compression too.

Step through the ablation. Watch the ViT and ResNet bars: codebook-only helps, naive joint learning dips, and proximal search leaps. The dip is the whole reason the proximal gradient exists.

Chapter 9: Connections & Cheat Sheet

JLCM sits at the intersection of three lineages. Vector quantization / weight sharing (Deep Compression, Han 2015) gave the codebook idea. Gradient-based PTQ (AdaRound, BrecQ, GPTQ) gave the per-layer self-distillation recipe. Straight-through estimation (Bengio 2013) gave the trick of hand-designing a backward pass. JLCM's novelty is the implicit-codebook-via-reordering and the proximal index gradient — two small ideas that unlock the combination.

Where it fits in the compression map

Cheat Sheet

Reference implementation of the core insight

import numpy as np

def proximal_index_grad(C, cur_slot, upstream):
    """The paper's fix: gradient on the index that prefers NEAR codewords.
    C: (K,) codebook values.  cur_slot: int, the weight's current slot.
    upstream: scalar dL/d(reconstructed weight). Returns pull per slot (K,)."""
    ca = C[cur_slot]
    dist = np.abs(ca - C)                 # |c_a - c_j| for every slot
    D = np.sign(ca - C) / (1.0 + dist)  # DECREASING in distance -> proximal
    D[cur_slot] = 0.0                    # identity subtraction: don't reward staying
    return upstream * D                    # near slots get the strongest pull

# Compare to the NAIVE gradient, which the paper shows diverges:
def naive_index_grad(C, cur_slot, upstream):
    return upstream * C   # scaled by codeword VALUE -> favors extreme |c_j|, not nearby

C = np.array([-2.0, 0.0, 0.15, 1.3])     # the paper's ternary example
# weight=0.12 sits in slot 2 (value 0.15); a +grad SHOULD nudge it to slot 1 (value 0).
print("proximal", np.round(proximal_index_grad(C, 2, 1.0), 2))
# -> [ 0.31  0.87  0.    0.47]  : slot 1 (the NEIGHBOR) gets the biggest pull. Correct.
print("naive   ", np.round(naive_index_grad(C, 2, 1.0), 2))
# -> [-2.    0.    0.15  1.3 ] : slot 0 (value -2) dominates by magnitude. Wild jump.