Hierarchical codebook classification is a small but useful model-building pattern: convert local image patches into discrete tokens, group neighboring tokens into coarser tokens, and classify by asking which label best explains the coarse pattern. It borrows the intuition of renormalization from physics: preserve spatial neighborhoods while repeatedly summarizing fine details into larger-scale structure.
The technique is not a replacement for convolutional networks or vision transformers. Its value is instructional and architectural. It shows how tokenization, vector quantization, spatial coarsening, Bayesian counting, and marginal inference can form a complete image classifier without gradient descent. That makes it a clean way to understand why hierarchy and locality matter in modern vision models.
The source model works on small digit images. A first-level codebook maps 4x4 pixel patches to token ids. A second-level codebook maps 2x2 neighborhoods of those tokens to supertokens. Training counts class priors, class-to-supertoken frequencies, and supertoken-to-token frequencies. Inference evaluates all possible supertokens for each block and marginalizes them with a log-sum-exp calculation.
The first layer divides an image into fixed-size patches and assigns each patch to the nearest prototype in a learned codebook. Each image becomes a grid of token ids, not a flat bag of pixels. The grid is critical because later coarsening depends on which tokens are neighbors.
A supertoken summarizes a small spatial block of lower-level tokens. In the toy model, a 2x2 block of patch tokens is represented by one id from a second codebook. The supertoken captures arrangements such as a curve in one corner and a vertical stroke in another, which are more informative than isolated patches.
Instead of backpropagation, the classifier estimates categorical distributions with smoothed counts. It learns how often each class appears, which supertokens are common for each class, and which lower-level tokens are likely under each supertoken. Dirichlet smoothing prevents zero probabilities during inference.
The classifier does not have to choose one best supertoken for a block. It can score every possible supertoken by combining its class likelihood and its ability to explain the observed lower-level tokens, then sum over those alternatives. This preserves uncertainty instead of committing too early.
Start by extracting patches, standardizing each patch, and assigning it to the nearest codebook centroid. With 16x16 input and 4x4 patches, the image becomes a 4x4 grid of token ids. Standardization keeps the codebook focused on local shape rather than absolute brightness.
def tokens_grid(img, codebook, patch=4):
rows = img.shape[0] // patch
cols = img.shape[1] // patch
ids = []
for i in range(rows):
for j in range(cols):
block = img[i * patch:(i + 1) * patch, j * patch:(j + 1) * patch]
vec = block.reshape(-1)
vec = (vec - vec.mean()) / (vec.std() + 1e-6)
distances = ((codebook - vec) ** 2).sum(axis=1)
ids.append(distances.argmin())
return np.array(ids).reshape(rows, cols)Training estimates three distributions: pi = p(class), theta = p(supertoken | class), and psi = p(token | supertoken). The first captures label frequency. The second says which coarse patterns belong to each class. The third says how a coarse pattern decomposes into fine tokens.
counts_c = np.zeros(num_classes)
counts_theta = np.zeros((num_classes, num_supertokens))
counts_psi = np.zeros((num_supertokens, num_tokens))
for image, label in zip(images, labels):
z = tokens_grid(image, token_codebook)
s = supertokens_for_grid(z, token_codebook, super_codebook)
counts_c[label] += 1
np.add.at(counts_theta[label], s, 1)
k = 0
for i in range(0, z.shape[0], 2):
for j in range(0, z.shape[1], 2):
block = z[i:i + 2, j:j + 2].ravel()
np.add.at(counts_psi[s[k]], block, 1)
k += 1
pi = normalize(alpha + counts_c)
theta = normalize_rows(beta + counts_theta)
psi = normalize_rows(gamma + counts_psi)During inference, each 2x2 token block contributes evidence for every class. For each class, the model sums over possible supertokens: how likely is the supertoken under the class, and how well does it explain the observed fine tokens? Computing this in log space avoids underflow.
This model is a compact demonstration of visual tokenization. Engineers can see how continuous pixels become discrete ids, how codebooks compress local detail, and how spatial hierarchy adds signal that a flat histogram would lose.
Counting models are useful baselines when debugging larger neural systems. If a small hierarchical model performs better than chance, the data likely contains local reusable structure. If it fails completely, the task may require position, scale, color, or long-range relationships that the hierarchy does not encode.
The reconstruction path provides a sanity check: if codebook tokens can roughly reconstruct the image, they preserve meaningful local information. That is the same reasoning behind modern discrete latent models, neural codecs, and image tokenizers, even though production systems use learned neural encoders instead of simple nearest-centroid assignments.
Keep token ids in a two-dimensional grid until the final classifier. Flattening too early destroys the neighborhood relationship that makes coarsening meaningful. A supertoken should summarize adjacent tokens, not arbitrary tokens with similar ids.
def block_features(token_grid, token_codebook):
features = []
for i in range(0, token_grid.shape[0], 2):
for j in range(0, token_grid.shape[1], 2):
ids = token_grid[i:i + 2, j:j + 2].ravel()
vec = token_codebook[ids].reshape(-1)
vec = (vec - vec.mean()) / (vec.std() + 1e-6)
features.append(vec)
return np.stack(features, axis=0)A class distribution over supertokens ignores where the supertoken appeared. That is sometimes too weak: two digits can share similar strokes but arrange them differently. A stronger variant learns p(supertoken | class, position), adding position as a conditioning variable while keeping the same counting machinery.
Hard nearest-centroid assignments are simple, but they discard uncertainty. A softer variant keeps a distribution over likely token ids for each patch and propagates those weights upward. This is closer to the spirit of hierarchical probabilistic modeling and can reduce errors around ambiguous strokes.