Learned 3-Bit Codes on a Budget-Exact Allocation: GSQ and RCO on Apple Silicon

At 3 bits per weight, the gap between a quantized language model and its full-precision original comes from two decisions. The first is where the bits go: some linears tolerate 2 bits, others break below 5. The second is which code each weight gets: rounding to the nearest grid point, even with GPTQ's error feedback, is a greedy choice. Two recent papers attack each decision separately.

• RCO (Helcig and Alistarh, arXiv:2605.00649) chooses a bit-width for every linear under an exact average-bit budget, by optimizing allocation probabilities on the manifold where the expected cost equals the budget.
• GSQ (Dadgarnia et al., arXiv:2604.18556) learns which grid point each weight takes, relaxing the discrete choice with Gumbel-Softmax sampling and training it block by block against the full-precision outputs.

This post presents gsq-rco-mlx, our implementation of both papers in MLX for Apple Silicon, and the first run that chains them: RCO picks the allocation, GSQ trains the codes on it, and the result is exported as a bit-packed mixed-precision mlx-lm checkpoint. On Qwen2.5-0.5B the chain scores 16.67 WikiText-2 perplexity at 2.92 code bits, against 20.14 for uniform 3-bit GPTQ and 13.08 for the bf16 model.

The numbers are recorded single-seed runs on one small model, every one traceable to a committed run record in the repository. They are not general benchmarks, and the limits are spelled out below.

Mixed-precision allocation is a multiple-choice knapsack: every linear $i$ picks one option from a menu of bit-widths, and the weighted average must hit a budget $B$. The usual approach scores each option with a per-layer proxy (output MSE, a Hessian trace) and solves the knapsack once. The proxy ignores how errors in different layers interact.

RCO instead keeps a softmax distribution $p_i = \mathrm{softmax}(\alpha_i)$ over each linear's options and constrains the expected cost to equal the budget exactly:

where $w_i$ is the linear's weight count and $c$ the vector of option costs in bits. Every optimizer step stays on this level set: the gradient is projected onto the tangent space,

an Adam step is taken, and a bracketed binary search retracts the iterate back to $C(\alpha) = B$. The loss is evaluated on discrete allocations: a Gumbel-perturbed knapsack is solved exactly by dynamic programming, with a straight-through gradient back to $\alpha$. The objective is the full-model one from the paper's Appendix B, the mean token KL divergence from the bf16 model to the model assembled from a database of pre-quantized GPTQ weights. The allocation is therefore judged on the whole network, not layer by layer. After training, a final noise-free DP on $\log p$ decodes the allocation.

With bit-widths fixed, each weight is a group scale times a grid value, $w = s \cdot q$, with $q$ on a grid of $2^b$ levels. GSQ gives every weight a vector of logits $l$ over candidate grid points and samples a soft assignment:

At 2 bits the candidates are the full grid; at 3 bits and above only the five neighbours of the initial code are kept (the local-shift parametrization), which keeps memory linear in the weight count. The temperature $\tau$ anneals from 2 to 0.05 and the sharpness $\kappa$ from 100 to 500, so the soft weights harden into ordinary scalar codes by the end of training.

Training starts from a GPTQ solution and proceeds block by block in three stages, each with the paper's loss: qk (per-linear reconstruction), vo (attention output with q and k hardened) and mlp (block output with attention hardened). Logits and scales are trained with Lion at the paper's learning rates (1e-4 and 5e-5). The output is a standard group-wise quantized checkpoint: no codebooks, no new kernels.

The implementation is one Python package, gsq_rco, with a single gsq-rco command-line interface. Each stage is its own subcommand and the run command chains them with resumable, config-hashed steps:

build-db   GPTQ-quantize all 168 linears at 2/3/4/5/6/8 bits (group 64)
   ↓        → database of residuals + manifest
allocate   RCO: 100 Riemannian Adam steps, 8 Gumbel samples each,
   ↓        full-model KL on 16 × 512 calibration tokens
   ↓        → per-linear bit-widths, exact 3.0-bit expected cost
quantize   GSQ: GPTQ warm start, staged qk / vo / mlp training per block,
   ↓        Lion, 200 steps per stage, 128 × 512 calibration tokens
export     bit-packed mixed-precision mlx-lm checkpoint (per-linear bits)
   ↓
eval       WikiText-2 test perplexity, 146 windows × 2,048 tokens

Everything runs on the GPU of a 16 GiB M3 through MLX. GSQ streams calibration activations block by block and loads weights lazily, so the heaviest step peaks at 3.12 GiB. The exported checkpoint loads with a plain mlx_lm.load(), with each linear carrying its own bit-width.

All rows below use the same evaluation: the WikiText-2 raw test split under the GPTQ protocol, 146 non-overlapping windows of 2,048 tokens (298,862 scored tokens), Qwen2.5 tokenizer, seed 0. Bits per weight include one fp16 scale per group of 64 weights. GPTQ rows are scored directly from the quantized database; GSQ rows are scored from the exported mlx-lm checkpoint.

Because both methods were run alone and together, the gain can be split. Each row changes one factor and holds the other fixed.

The RCO allocation over the 168 linears (357.8M quantized weights) is far from uniform. Of the 168 linears, 35 get 2 bits, 77 get 3, 30 get 4, 11 get 5, 12 get 6 and 3 get 8. The expected cost stayed on the budget at every step (largest deviation 9.7e-9 bits), and the sampled KL objective fell from 1.40 to 0.41 over 100 steps.

This implementation started as a port that reported 1,418 perplexity. An audit against both papers found 23 defects, among them every quantized block applied twice, a per-layer proxy in place of the RCO objective, and a GSQ loop that could not change most codes. All were fixed before any number in this post was recorded; the repository lists each finding with its fix commit. Two more failures surfaced only when the methods were chained, and both are worth knowing about for anyone combining learned codes with mixed precision.

Lion's absolute step breaks high-bit scales

The first GSQ run on the RCO allocation scored 63.53, worse than plain GPTQ on the same allocation. Lion moves every parameter by roughly its learning rate per step, in absolute terms. A $b$-bit GPTQ scale is about $\max|w| / 2^{b-1}$, so scales halve with every extra bit. The paper's scale learning rate of 5e-5 is tuned for 2- and 3-bit grids; on 5-, 6- and 8-bit linears the same step is a large relative change. On layer 3, the mlp stage objective ended 3.97× above its GPTQ starting point.

The fix follows the GSQ paper's own Table 9, which sets the scale learning rate to zero for its GGUF K-quant formats: scales above 3 bits are frozen, logits still train. With the same allocation and configuration, the rerun scores 16.67, and all 72 stages end below their starting objective.

8-bit linears and the stage optimizer

RCO gave three linears 8 bits, an option uniform runs never exercise. The stage optimizer crashed on the first of them because it expected logits on every linear. At 8 bits GPTQ is already close to lossless, so 8-bit linears are now kept at their GPTQ codes and excluded from stage training; the optional end-to-end scale fine-tune can still adjust their scales.

Every step was run alone on an Apple M3 with 16 GiB of unified memory, MLX 0.32.4.

• One model, one dataset, one seed. Qwen2.5-0.5B on WikiText-2. No confidence intervals; run-to-run spread is not measured. GSQ on the MLX GPU is also not bit-reproducible, because the backward of mx.take is nondeterministic (the CPU path is deterministic).
• A much smaller training budget than the papers. GSQ calibrates on 128 sequences of 512 tokens with 200 steps per stage, group size 64 and micro-batch 2. The paper uses 4,096 sequences of 4,096 tokens, about 1,280 steps per block, group size 128 and batch 64. Our numbers are not comparable with the papers' tables.
• Deliberate deviations. Scales are frozen above 3 bits, 8-bit linears are GPTQ-only, and the paper's end-to-end scale fine-tune (Appendix H) is implemented but was not run.
• The RCO decode underspends. 2.92 of 3.00 bits. A budget-tight decode would make the allocation comparison iso-bit.

gsq-rco-mlx needs Apple Silicon, Python 3.12 and the bf16 Qwen/Qwen2.5-0.5B checkpoint in models/Qwen2.5-0.5B/:

uv venv --python 3.12 .venv && source .venv/bin/activate
uv pip install -e ".[dev]"

gsq-rco data                                                       # WikiText-2, downloaded once
gsq-rco run --run-dir results/runs/rco3 --avg-bits 3.0 --dry-run   # print the plan
gsq-rco run --run-dir results/runs/rco3 --avg-bits 3.0             # build-db → allocate → quantize → export → eval

gsq-rco run is resumable: finished steps whose config hash matches are skipped on a rerun. Each step is also its own subcommand, for example a uniform GSQ run without any allocation:

gsq-rco quantize --uniform-bits 3 --out results/runs/gsq_u3
gsq-rco eval --mlx-model results/runs/gsq_u3/mlx --out results/runs/eval_gsq_u3.json

The repository documents every flag and file format, maps each equation of both papers to the code, and commits the run records behind every number in this post.

Allocation and code training are complementary but not independent. On Qwen2.5-0.5B, RCO's budget-exact allocation and GSQ's learned codes each beat uniform 3-bit GPTQ on their own, and together they reach 16.67 perplexity at 2.92 bits, 17% below uniform 3-bit GPTQ, in an ordinary group-quantized checkpoint that loads in stock mlx-lm. The full-model objective is what makes the allocation work: the same knapsack solved on per-layer error is worse than not allocating at all.

The whole pipeline fits on a 16 GiB laptop in under three hours. The open questions are scale: larger models, the papers' full training budgets, a budget-tight decode, and several seeds.

gsq-rco-mlx is open-source at github.com/egesabanci/gsq-rco-mlx. It is an independent implementation; please cite the GSQ and RCO papers if you use it.