# How Quantization Codebooks Are Trained in Cactus Quants: Lloyd-Max on Gaussian Data

> Discover how Cactus Quants trains quantization codebooks using the Lloyd-Max algorithm on Gaussian data. Learn about offline training and deterministic k-means clustering for optimal reconstruction levels.

- Repository: [Cactus Compute, Inc./needle](https://github.com/cactus-compute/needle)
- Tags: how-to-guide
- Published: 2026-09-09

---

**Cactus Quants trains its quantization codebooks offline using the Lloyd-Max algorithm on synthetic Gaussian data, implementing deterministic 1-D k-means clustering in [`needle/model/quantize.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/quantize.py) to generate optimal reconstruction levels without requiring real model weights.**

The [needle](https://github.com/cactus-compute/needle) repository by Cactus Compute provides an efficient quantization framework that relies on pre-computed codebooks rather than learning quantization levels from actual model parameters. This approach ensures deterministic, reusable quantization schemes that generalize across different neural network architectures.

## The Lloyd-Max Algorithm Implementation

Inside [`needle/model/quantize.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/quantize.py), the `_lloyd_max_gaussian` function implements a classic Lloyd-Max quantizer (equivalent to 1-D k-means) to derive optimal quantization centroids. This algorithm minimizes the mean squared error between original Gaussian samples and their quantized representations.

### Gaussian Sampling and Initialization

The training process draws **400,000 samples** from a standard normal distribution using `np.random.RandomState(seed).randn(samples)`[/needle/main/needle/model/quantize.py#L95-L100]. Rather than random initialization, the algorithm selects initial centroids by uniformly spacing them across the sorted sample array:

```python
c = x[((np.arange(levels) + 0.5) / levels * samples).astype(int)]

```

[/needle/main/needle/model/quantize.py#L98-L99]

This uniform initialization strategy ensures stable convergence properties regardless of the random seed used for sampling.

### Iterative Refinement of Quantization Levels

For a default of **200 iterations**, the algorithm performs three operations to refine the codebook:

1. **Boundary calculation**: Computes decision thresholds as midpoints between consecutive centroids using `bnd = (c[:-1] + c[1:]) / 2.0`[/needle/main/needle/model/quantize.py#L100-L101].

2. **Sample assignment**: Assigns each Gaussian sample to the nearest centroid via `idx = np.searchsorted(bnd, x)`[/needle/main/needle/model/quantize.py#L101-L102].

3. **Centroid update**: Updates each level to the mean of its assigned samples through `c[k] = x[m].mean()`[/needle/main/needle/model/quantize.py#L103-L106].

Upon completion, the final centroid array `c` is sorted and returned as the trained codebook[/needle/main/needle/model/quantize.py#L106-L107].

### Ternary Special Case Handling

When the requested bit-width equals the predefined `TERNARY_BITS` constant (approximately 1.58 bits), the Lloyd-Max training is bypassed entirely. The system returns a hardcoded ternary codebook:

```python
[-1.2240064, 0.0, 1.2240064]

```

[/needle/main/needle/model/quantize.py#L108-L110]

These values represent optimal reconstruction levels for ternary quantization of Gaussian-distributed data, eliminating the need for iterative computation in this specific configuration.

## Production Codebook Generation and Caching

The public API function `_cq_codebook_np(bits, group_size)` exposes trained codebooks to the rest of the library. This function utilizes `@lru_cache` to store computed codebooks, ensuring that identical configurations are retrieved instantly without recomputation.

Before returning results, the function scales the Gaussian-derived centroids by the square root of the group size (`/ sqrt(group_size)`) to accommodate grouped weight matrix quantization[/needle/main/needle/model/quantize.py#L112-L115]. This scaling maintains statistical consistency across different tensor shapes and group dimensions.

## Practical Usage Example

The following implementation demonstrates generating a Lloyd-Max trained codebook and applying it to weight matrices:

```python
import jax.numpy as jnp
from needle.model.quantize import _cq_codebook_np, cq_quantize

# Generate 4-bit codebook for group size 128 using Lloyd-Max training

codebook = _cq_codebook_np(bits=4, group_size=128)
print("Codebook shape:", codebook.shape)

# Quantize example weight matrix

weights = jnp.random.normal(size=(1024, 512))
q_weights = cq_quantize(
    weights, 
    bits=4, 
    group_size=128, 
    codebook=jnp.asarray(codebook)
)

# Verify quantization used the trained codebook

unique_values = jnp.sort(jnp.unique(q_weights))

```

The `cq_quantize` function applies Hadamard rotation to the weights, normalizes them against the codebook scale, and snaps values to the nearest centroid using `_cq_nearest`, then reverses the transformation to produce the final quantized output.

## Summary

- Cactus Quants trains **quantization codebooks** offline using the **Lloyd-Max algorithm** on 400,000 synthetic Gaussian samples rather than learning from real model weights.
- The implementation in [`needle/model/quantize.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/quantize.py) converges centroids through 200 iterations of boundary calculation, sample assignment, and mean updates according to classic 1-D k-means principles.
- **Ternary quantization** uses a hardcoded optimal codebook `[-1.2240064, 0.0, 1.2240064]` instead of iterative training when `TERNARY_BITS` is specified.
- The `_cq_codebook_np` function provides LRU caching and scales centroids by `sqrt(group_size)` to support efficient grouped matrix quantization across different model architectures.

## Frequently Asked Questions

### Why does Cactus Quants use synthetic Gaussian data instead of real model weights?

Cactus Quants uses synthetic Gaussian data because neural network weights typically approximate Gaussian distributions after transformations like Hadamard rotation. Training on synthetic data produces deterministic, reusable codebooks that generalize across different models without requiring access to proprietary training datasets or specific model checkpoints. This approach ensures consistent quantization behavior regardless of the target model architecture.

### How many iterations does the Lloyd-Max algorithm perform by default?

The `_lloyd_max_gaussian` function performs **200 iterations** by default when training quantization codebooks. This iteration count provides sufficient convergence for optimal Gaussian quantization levels while maintaining fast offline computation speed. The iteration count is configurable when calling the function directly, though the standard library API typically uses these defaults for consistency.

### What is the significance of the ternary codebook values `[-1.2240064, 0.0, 1.2240064]`?

These values represent the mathematically optimal reconstruction levels for ternary quantization (three discrete levels) of a standard Gaussian distribution. The constant `1.2240064` derives from statistical optimization to minimize mean squared error for Gaussian sources. By hardcoding these values in [`needle/model/quantize.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/quantize.py), Cactus Quants avoids expensive iterative computation for this common low-bit-width configuration while maintaining quantization accuracy.

### How does group size affect the generated quantization codebook?

The codebook values are scaled by `1/sqrt(group_size)` in the `_cq_codebook_np` function to account for variance reduction that occurs when grouping weights during quantization. This scaling ensures that the quantization levels remain statistically appropriate for the specific group dimensions used during model quantization, maintaining consistent signal-to-quantization-noise ratios across different tensor shapes and grouping configurations.