Does your weight format deserve a kernel?

You have an idea for squeezing model weights into less space than MLX's built-in formats use. MLX already supports two-bit quantization, but two bits per weight is not the whole storage cost. Perhaps three possible values per weight would be enough, packed into a different layout. You can sketch the bytes on paper, but MLX has no operation that reads that layout.

Before writing that operation, there is a question worth answering: can the model still do useful work with the numbers your recipe preserves?

There is a useful experiment to do first: pack the weights, unpack them back to ordinary floating-point numbers, and run the model. The runtime sees a familiar checkpoint. The model has to live with the numbers your recipe left behind.

We use this approach in our own low-bit work to compare recipes before investing in a kernel, the code that performs the calculation on the processor. Here we will try it with NumPy and stock MLX, using perplexity to measure the effect on prediction.

For comparison, these are some of the weight representations available in stock MLX.¹

The affine rows use groups of 64 with a 16-bit scale and a 16-bit bias, adding 0.5 bits per weight to the codes. File headers and weights left unchanged are excluded.

Our three-value, or ternary, layout uses 1.875 bits per selected weight at the same group size, including its scale. MLX can represent those values in its existing two-bit format, but cannot execute our more compact layout directly.

What survives the round trip?

Quantization chooses the numbers. Floating-point weights become a shared scale and small integer codes. This generally loses information: several original values now have the same reconstruction.

Packing arranges those codes into bytes. Storing and recovering the chosen codes should be lossless. Their arrangement should not change model quality.

The runtime computes with the representation. It might decode small pieces inside a matrix multiplication, avoiding a full floating-point copy.

Our experiment decodes the package into a checkpoint using bfloat16, or bf16, a 16-bit floating-point format the model already supports. Putting the reduced set of values in a larger container does not restore the information quantization removed.

PyTorch's fake quantization uses the familiar idea of simulating quantization and dequantization with floating-point output.² Writing and reading the actual package also exercises our proposed storage layout.

Three values and a shared scale

For our example, each weight becomes one of $\{-s, 0, +s\}$. The sign or zero is a trit, the base-three cousin of a bit. A group of $G$ consecutive weights shares the scale $s$.

We will use a deliberately simple rule: take the group's mean absolute weight as its scale, round that scale to fp16 for storage, then round and clamp the normalized weights:

$$s = \operatorname{fp16}\!\left(\frac{1}{G}\sum_{j=1}^{G}|W_j|\right)$$ $$t_j = \operatorname{clip}\!\left(\operatorname{round}(W_j/s), -1, 1\right)$$

Reconstruct each weight as $\widehat{W}_j = \operatorname{bf16}(s\,t_j)$. The scale uses fp16, a different 16-bit floating-point format from bf16. Calculating with the stored scale includes its rounding in the experiment. Rounding uses nearest-even ties; a zero scale reconstructs an all-zero group without division.

This rule is easy to inspect. Whether a model trained with ordinary weights tolerates being made ternary afterward remains to be seen.

Five trits in a byte

A byte has 256 possible values. Five trits have $3^5 = 243$ possible combinations, so all five fit. Map trits $\{-1,0,+1\}$ to digits $\{0,1,2\}$ and write an ordinary base-three number:

$$b = d_0 + 3d_1 + 9d_2 + 27d_3 + 81d_4$$

This small NumPy example packs and unpacks one byte. The first digit is the least significant one:

import numpy as np

trits = np.array([-1, 0, 1, -1, 1], dtype=np.int16)
powers = np.array([1, 3, 9, 27, 81])

packed = np.uint8((trits + 1) @ powers)  # 183
written = bytes([int(packed)])
byte = int.from_bytes(written, "little")
recovered = (byte // powers) % 3 - 1

np.testing.assert_array_equal(recovered, trits)

The full example writes groups of codes and scales to disk. Its decoder receives only that package, with no original checkpoint from which it could regenerate the intended weights.

Before asking a model anything, test all 243 possible five-trit combinations against an independent scalar calculation. Test incomplete final bytes, group boundaries and the unused byte values 243 through 255 too. Our format rejects those values and requires padding positions to decode to zero.

A successful round trip means the chosen trits survived. Whether they were good choices is a different test.

Count the bytes you actually wrote

Five trits per byte suggests $8/5 = 1.6$ bits per weight. Our layout starts each group on a byte boundary, however, so a group of 64 needs 13 code bytes. Add two bytes for its fp16 scale and the group occupies 15 bytes.

The payload rate for this layout is therefore:

$$\text{bits per weight} = \frac{8\bigl(\lceil G/5\rceil + 2\bigr)}{G}$$

These are rates for the selected weight payload, including group padding. File headers, tensor shapes and weights left unchanged also take space. The demo records complete package bytes separately.

More constrained ternary formats can use even fewer code bits. Sherry, for example, requires exactly one zero and three nonzero trits in every block of four weights. That leaves 32 possible patterns, which fit in five bits: 1.25 code bits per weight, before scales, metadata and padding.³ The saving comes from restricting which weight patterns are allowed, not from packing arbitrary trits more tightly. A recipe inspired by that constraint is another candidate for a pack-unpack quality screen before adding runtime support.

We can measure a package's file size before a kernel exists. What the bf16 model run cannot tell us is how much memory a native implementation would need, or how quickly it could use those bytes.

Ask the model

The example uses the text path of Qwen3.5-0.8B's MLX bf16 checkpoint. We change the feed-forward layers' gate, up and down projection weights. Attention, embeddings, normalization and other parameters keep their original values. This is a partial-model experiment, not a claim that every weight can be stored at the rate above.

For each group size, the script writes a package, decodes it to a fresh bf16 checkpoint, reloads that checkpoint through ordinary MLX LM, and evaluates it against the untouched reference. Both see exactly the same saved token blocks.

Perplexity measures how much probability the model assigns to the next tokens in a text: it is the exponential of their mean negative log probability. Lower is better.

We use 64 non-overlapping blocks of 512 tokens from the English WikiText-2 raw validation split, selected with seed 123. Each block starts with fresh model state; its first token supplies context and the remaining 511 are scored. That makes 32,704 scored tokens per configuration. Model and dataset revisions, tokenization and software versions are fixed in the demo.⁴

Both configurations damage prediction severely. Group 64 is less destructive here, but neither result gives us a reason to build a kernel for this recipe. The packer has preserved the trits exactly; the quantizer has discarded too much useful information. This says something about our simple rule and the weights we selected, not about every ternary method.

The complete sweep, including packing, checkpoint checks and all three evaluations, took about 109 seconds on an Apple M1 Max with the model and corpus already cached. That is time spent answering the quality question, not a measurement of packed inference speed.

The comparison is the useful unit of work. You can vary the group size, replace the scale rule or change the set of tensors receiving the recipe, then run the same screen again. Keep a separate held-out evaluation for the configurations you eventually select: a text pool used repeatedly to choose recipes has become part of that choice.

Try a recipe of your own

The companion example contains the packer, unpacker, checkpoint writer and evaluation runner. On an Apple Silicon Mac, after following its setup instructions and installing the locked dependencies with uv sync, the experiment is two commands:

uv run python -m unittest -v
uv run python screen.py --group-sizes 64 128

The runner saves the input token blocks, selected tensor names, package hashes and results under work/. It calls MLX LM's eval_ppl on the saved blocks with batch size one. This keeps the measurement small and avoids a changing default dataset.⁵

The decoder accepts only the package directory. That boundary is worth preserving when you replace the example with your own recipe.

An existing kernel may be enough

Your values may fit a format the runtime already implements. MLX's affine quantization reconstructs a value as code * scale + bias.⁶ Codes $\{0,1,2\}$ with scale $s$ and bias $-s$ express our ternary values inside its two-bit container.

With one fp16 scale and one fp16 bias per group of 64, that payload costs $2 + 32/64 = 2.5$ bits per weight. It runs in the existing container; our five-trit byte layout still needs its own support. Irregularly spaced values will not generally fit this mapping. Matching reconstructed values also does not guarantee identical floating-point results from different kernels.

We have not run that experiment here. It is an option after the dense screen, when the values are worth taking further.

Conclusion: test the recipe first

Our ternary packages recovered every chosen trit exactly, but the reconstructed model's prediction quality deteriorated severely. For this recipe, the next step is to revisit which weights we quantize and how, perhaps keeping trits in some layers and using MLX's existing formats elsewhere. A faster kernel would still be computing with the same damaged weights.

That is the decision the round trip helps us make. Start with a checkpoint that runs, change only the intended weights, and compare the model decoded from the actual package with its own reference. A poor result gives you something concrete to revise before writing a kernel. A promising result gives you a reason to investigate further.

Perplexity on this text does not settle task quality, other languages or longer contexts. The screen does not test quantized activations or the key/value cache used during generation. A native implementation will also need checks for its own arithmetic, speed and memory use. Those questions remain, but we can choose which recipes to spend that work on.

Footnotes

1. MLX's quantize documentation lists the affine bit widths; its data type reference gives the floating-point widths. The table is a focused comparison, not a complete list of MLX's quantization modes. Byte counts were checked in MLX 0.32.0 with a 64-by-64 bf16 array, summing the packed codes, scales and biases returned by quantize with group size 64. This checks storage, not model quality or inference speed. ↩

2. PyTorch's FakeQuantize documents the quantize-dequantize simulation with floating-point output. ↩

3. Huang et al.'s Sherry: Hardware-Efficient 1.25-Bit Ternary Quantization via Fine-grained Sparsification describes the five-bit encoding in section 3.1. Sherry combines structured sparsity with quantization-aware training and hardware-aware packing. Our elementary post-training example does not implement or evaluate that method. ↩

4. Measured 2026-09-14 on an Apple M1 Max, 32 GiB, macOS 26.5.1, using MLX 0.32.0, MLX LM 0.31.3, NumPy 2.5.1 and Transformers 5.17.0. The 72 selected matrices contain 264,241,152 of the text model's 752,393,024 parameters. The result receipt records model and dataset revisions, input and source hashes, sampled blocks and timing. The README specifies tokenization. No training or calibration uses this text. ↩

5. MLX LM's eval_ppl implementation computes next-token cross entropy and exponentiates its token mean. We use that evaluator with our own fixed input blocks. These values should not be compared directly with tables using other tokenization, context lengths or sampling procedures. ↩

6. MLX's quantize documentation describes affine groups with scales and biases. The rate here assumes both use 16-bit storage and that two-bit codes are packed without additional padding. ↩