Quality quantization at very low bpw (<= 2 bit): mixed allocation + custom packers

[magnuslundstedt]

I am fairly new to AI research but find it very interesting. My research goal is efficient local AI, optimizing the balance between quality and speed using MLX through aggressive, non-uniform bit allocation.

Recently in my experiments I found that compressing some weights in Qwen to 1 bit, ternary (1.58 bit, 1.6 to 1.7 bpw packed) and 2 bit, while spending the bit budget on the more crucial weights, can yield smaller models and lower RSS. But quality is a challenge. I've not managed to test end-2-end yet, rather custom quantizer and packer -> unpack to FP16 -> measure PPL against a regular flat allocation at matched size.

I take inspiration from GGUF and the llama.cpp project, for example K-quant hierarchical scaling, TQ1_0 packing 5 trits per byte for BitNet (3^5 = 243, base-243), the work by Unsloth, and in general inspecting how successful quantized models have allocated their bit budget through a custom tool I built as a side-quest for fun (tensorlens.dev). One inherent benefit in GGUF seems to be the K-quant hierarchical (superblock) scaling, which lets smaller quants still perform well. Besides the common challenge of bit allocation and rounding (prior knowledge, activation sensitivity, raw SSE, absmean/Lloyd-Max, etc), what about the capabilities of MLX itself?

That is the question I'm most curious about here: MLX quantize is affine group quantization with a single scale/bias per group and stops at 2-bit. Has anyone explored K-quant-style hierarchical scaling, or sub-2-bit / ternary packers, in MLX? Or more generally, what interest, research or experience is here on custom packers, quantizers and extending MLX in this direction?

Any interesting very-low-bpw MLX experiments, papers or discussions you can share or reference would be great.

[Diacod-I]

Hey @magnuslundstedt. I realise this is an old question, but here's my answer in case it's still useful. This is what I can see on current main:

Stock MLX's mx.quantize supports 2, 3, 4, 5, 6 and 8 bits (affine, one scale and bias per group) plus mxfp4, nvfp4 and mxfp8, but nothing below 2 bits and no superblock scales yet.

For ternary, mlx-lm already runs BitNet b1.58 models with a custom Metal kernel that computes directly on the packed weights (bitlinear_layers.py), which could be a good starting point for your own packer.

For K-quants, mlx-kquant adds Q2_K to Q6_K and IQ1/IQ2 types to MLX as native ops.

For non-uniform bits, mlx-lm's quant_predicate lets you set bits per layer, and DWQ and dynamic quant handle sensitivity-based allocation.

Also a tip: unpacking to FP16 only tells you about quality, so to see the RSS and speed gains, you'll need a kernel that reads your packed format directly.

[magnuslundstedt]

Thanks! Great suggestions. My original question was a while ago, but it's still highly relevant to me. Especially mlx-kquant is new to me, will check it out, thanks!

Since my original post I've evaluated a number of low bpw formats with and without custom kernels, and found them potentially very useful. I've tested both paths: unpacking to FP16 before loading weights, versus custom kernels that unpack in tile memory (with and without NAX). The good news is that it's possible to get good RSS and speed at <2 bpw during both prefill and decode with custom kernels. The sobering news is that it becomes exponentially more important to know exactly which weights need the bits and where it's safe to remove them at low bpw. 😅

So my current challenge in actually using these exotic formats has ended up being dynamic bit allocation (optimizing bit allocation and group sizes, scale/bias). Just mathematically minimizing the rounding error using various methods is not universally good enough.

So far I've tried families such as rounding-error minimization, activation-based (AWQ style), gradient/KL-guided, and copying bpw allocation maps from Unsloth releases; I haven't found the ideal method yet. But which weights get the bits is what moves PPL.

I now treat format and allocation as one problem: one global byte budget across tensors, where each tensor gets its own family (dtype) and group size, spanning 1.25-8.5 bpw. At runtime my kernels read all available dtypes directly, expanding only in tile memory. Tweaking DWQ is a piece I still have left.

I'm making good progress evaluating methods for allocation, but I needed more RAM to test the more advanced and promising ways of allocating bits, so right now my project is waiting on a new Mac Studio I ordered for this. 💸😂

Are there methods for bit allocation I'm missing, or papers worth reading?

[Diacod-I]

@magnuslundstedt That's awesome progress, and good luck with the new Mac Studio lol!

Since you're treating it as one global budget, you can maybe try something like this:

1. For each tensor, measure how much the end-to-end loss (e.g. KL on a calibration set) rises at each candidate format and group size.
2. Then solve a multiple-choice knapsack over those curves, which will essentially boil your problem down to "giving bits to whichever tensor gets the biggest loss drop per byte".
3. Layers aren't independent, so your errors can compound or cancel across them, hence do a final whole-model refinement pass after he knapsack.

I have listed some work in this direction you could check out if you are interested:

• HAWQ-V2 / HAWQ-V3: Hessian-trace sensitivity per layer, with bit-widths chosen by integer linear programming under a size budget.
• SqueezeLLM: Fisher-information weighted quantization, with a small set of outlier weights kept in full precision.
• SpQR / OWQ also keep the most sensitive weights or columns at higher precision. These might help at <2 bpw, where a few weights seem to dominate.
• SliM-LLM (ICML 2025): This does salience-driven mixed precision at the group level instead of per tensor.
• ExLlamaV2's EXL2 format: It measures each layer's error at several bitrates, then optimizes the mix to hit a target average bpw.
• CASA: This one is very recent and argues that collapsing each layer's sensitivity into a single number is unreliable. Instead, they use a Kronecker-factored Hessian metric, then refine the allocation with a cross-layer search on end-to-end loss specifically targeting below 3 bits.

Hope this answers your question! Would love to work with you on some of your experiments 😄.