A Discrete Weight Language for Large Language Models: Compressing Gemma 4 31B with a 5-Step Ternary Route
Arman Aubakirov
Independent research · April 2026
Abstract
We present Route B, a weight-compression scheme that treats every
weight tensor of a pretrained language model as a sequence of short
codes drawn from a shared, per-family code book — in effect, a
discrete language in which each scalar weight is spelled out by a
5-step ternary path through a learned ladder of amplitudes. Applied to
google/gemma-4-31B-it, the method replaces the 60 GB bf16 weights of
the 410 linear layers by 7.9 bits per weight (8-bit packed code + a
single fp32 scalar per layer, plus an optional 16-bit per-row scale),
reducing the working memory of inference to ~31 GB with no change to
the model's I/O interface. A fused Triton matmul that pulls the
243-entry layer code book into shared memory reaches 37 tokens / s
at batch size 4 on a single NVIDIA H200, without ever materialising
the bf16 weight matrix. On 32 k tokens of WikiText-2 at 2 k context the packed model is within
+0.08 nats of the dense bf16 baseline (loss 6.906 vs 6.829). Adding
a per-row scale reduces weight-reconstruction MSE by ~5× and matmul
output error by ~3×, but its end-to-end effect on Gemma 4 31B IT is
strictly harmful from 1 k context upwards: the gap vs dense widens
from +0.25 nats at 512 context to +0.64 nats at 8 k, even though
per-layer activation-weighted MSE is locally better with row scale.
A per-head prosody variant (one scalar ρ per attention head on
q/k/v_proj, no row scale elsewhere) eliminates the regression entirely:
it matches or beats the global-scale baseline at every context from
4 k to 16 k, and a brute-force end-to-end probe over the 3 locally
rejected k_proj layers shows that one of them should actually be
kept: force-keeping block 37 k_proj yields strict wins over the
global-scale baseline at all tested long contexts (−0.065 / −0.004 /
−0.008 nats at 4 k / 8 k / 16 k). This confirms that the damaging
component of naive per-row scale is within-head multiplicative noise
that compounds through softmax, while also showing that local
activation-MSE is still too conservative for the final keep/drop
decision. We analyse this regression in §6 and
argue that it reflects a fundamental limitation of weight-MSE — and
even of per-layer activation MSE — as PTQ objectives for
long-context transformers. The full pipeline is released as a 4-step
reproducible recipe on top of the public Gemma 4 31B IT weights.
We frame Route B as a first instance of a broader idea — representing
each layer as a sentence in a small alphabet whose letters are
multi-bit signed ladder steps — and discuss concrete paths towards
variable-depth, activation-aware and adaptive-alphabet variants that we
expect may further narrow the dense-vs-packed loss gap without leaving
the ~3-bit regime.
Production-grade LLMs at the 30 B – 70 B scale remain memory-bound on
mainstream GPUs. Even in bf16, Gemma 4 31B IT requires ~60 GB of
weights alone, leaving little room for KV cache, activations, or batch
parallelism on a single 80 GB card. Post-training quantization (PTQ) is
the pragmatic answer, and the literature spans from scalar INT8 all the
way down to ternary / binary weights (e.g. BitNet b1.58 [Ma et al.,
2024]), with a well-known trade-off: the lower the bit budget, the
more the reconstruction error — and, for sufficiently trained models,
the worse the downstream quality.
We revisit PTQ from a linguistic angle. Instead of asking
"into how few bits can I project a weight scalar?"
we ask
"what is the shortest sentence in a small alphabet that spells this
weight to within a chosen error?"
This is not a purely cosmetic rephrasing. It immediately suggests:
a small alphabet (a 5-letter ternary ladder) that is shared
across a whole family of layers and can therefore be stored at
essentially zero marginal cost per weight;
a word length (depth of the ternary path) that can vary per
layer, per row, or even per weight, trading bits only where they are
needed;
per-row prosody — a lightweight multiplier that lets the alphabet
span wider dynamic ranges without growing the alphabet itself;
a per-weight "spelling" that is discrete and so can be
interpreted, compared, searched, and (in future work) edited without
ever being re-floated.
Concretely, this paper presents the simplest instance of this
framework, which we call Route B: a fixed 5-step ternary word with
a shared per-family × depth-bucket ladder, packed into a single
uint8 per weight. We apply it to all 410 linear layers of Gemma 4
31B IT, ship a fused Triton kernel that decodes the LUT inside shared
memory, and release both the packed checkpoint and the full 4-step
reproducible recipe.
1.1 Contributions
A concrete, reproducible 3-bit weight-compression recipe for
Gemma 4 31B IT (410 linear layers, all-linear coverage including
the 10 shared-KV layers with v_proj = None), with code, packed
weights and a loading script.
A fused Triton matmul for 5-step ternary weights. On H200 with
bs=4 and 64 new tokens it reaches 37.2 tok/s at ~31 GB resident
in the default packed_triton_fused mode; packed_cached_decode
reaches ~42 tok/s only by re-materialising bf16 weights at ~60 GB.
A per-row scale extension that reduces weight-reconstruction
MSE by ~5× at a cost of one fp16 number per output row, with a
per-layer regression guard that only keeps the update where it
strictly helps.
An empirical diagnosis of where ternary codes fail and recover,
including the finding that a ~5× improvement in layer-level weight
MSE can still worsen end-to-end loss on Gemma 4 31B IT from 1 k
context onward, the per-head prosody correction that restores the
long-context regime, and a preliminary dense-vs-route downstream
comparison harness for early external validation.
1.2 Paper outline
§2 surveys related PTQ work and positions the discrete-language view.
§3 gives the mathematical formulation of Route B and the per-row
scale. §4 describes the runtime (Triton kernel, memory layout,
inference loop). §5 details the 4-step reproduction pipeline. §6
presents empirical results on Gemma 4 31B IT: weight MSE, PPL,
throughput / memory, and preliminary downstream comparisons. §7
discusses limitations and open problems. §8
outlines the broader "discrete weight language" programme and
concrete next steps.
2. Related work
Scalar PTQ. GPTQ [Frantar et al., 2022], AWQ [Lin et al., 2023],
SmoothQuant [Xiao et al., 2023] and INT4/INT8 group-wise schemes are
now the de-facto industry standard for serving large models. They
spend 4 – 8 bits per weight, usually with per-group scales of size
128.
Ternary and binary weights. BitNet b1.58 [Ma et al., 2024] and the
earlier ternary/binary literature show that models trained from
scratch with ternary or binary weights can match fp16 baselines. Post
hoc quantization of a pretrained model to ≤ 2 bits without retraining
remains challenging, with notable gaps on instruction-tuned models.
Look-up-table matmul. LUT-GEMM [Park et al., 2022] and SqueezeLLM
[Kim et al., 2023] materialise a small per-layer codebook and fuse the
dequantisation into the matmul; this is the direct ancestor of our
Triton kernel.
Product / vector quantization of weights. QuIP [Chee et al., 2023]
and QuIP# rotate weights into an incoherent basis before quantizing
with a structured lattice. Our work is compatible with such rotations:
the discrete "sentence" metaphor is independent of whether the
alphabet lives in the original weight basis or in a rotated one.
uses a single ladder per family × depth-bucket, not per layer or
per group,
emphasises a discrete, readable representation of the weights —
every scalar is spelled by an explicit word in {−1, 0, +1}⁵.
3. Route B — formal description
Let W ∈ ℝ^{O × I} be a weight matrix and let S = (s₁, s₂, s₃, s₄, s₅) be the ladder scales shared by the family of W.
3.1 Encoding
Given a real sample x = W[n, k], the 5-step greedy ternary
encoding produces a word c = (c₁, …, c₅) ∈ {−1, 0, +1}⁵ such that
c_i = argmin_{c ∈ {−1, 0, +1}} | r_{i−1} − c · s_i |
r_i = r_{i−1} − c_i · s_i (residual after step i)
r_0 = x / α (α is a per-matrix scalar)
Where α = max_{n,k} |W[n, k]| (the "global scale"). A small
zero-biasζ ∈ [0, 1] is used to prefer c_i = 0 when
|r_{i−1}| < ζ · s_i — this reduces the number of non-zero steps and
improves end-to-end quality.
3.2 Decoding and the layer code book
Decoding is by construction:
W_hat[n, k] = α · Σ_{i=1..5} c_i(n, k) · s_i
Because c ∈ {−1, 0, +1}⁵, there are only 3^5 = 243 possible words;
we enumerate them once per layer, pre-multiply by α, and store the
result as a 243 × 1 look-up table LUT ∈ ℝ^{243}. From that point on,
the weight matrix is fully described by the code matrix
codes[n, k] ∈ {0, …, 242} (base-3 encoding of c_1..c_5)
which we pack four or five per byte (see §4.1).
3.3 Per-row scale
A single scalar α has to cover the full dynamic range of W; rows
whose typical magnitude is smaller than α are under-resolved. We
introduce a per-row scale:
so the run-time cost is one fp16 multiplication per output scalar,
fused into the same Triton kernel.
3.4 Ladder fitting
For a group of layers (family × depth-bucket) we collect ~200 k
row-normalised samples and minimise, by coordinate descent on S:
L(S) = Σ_{x ∈ samples} (x − Σ_i c_i(x; S) · s_i)²
starting from the balanced-ternary initialisation
S₀ = (1, 1/3, 1/9, 1/27, 1/81). The resulting S is typically
close to but not equal to S₀; for Gemma 4 31B IT attention at the
early depth bucket we measure S ≈ (1.0, 0.33, 0.115, 0.037, 0.012).
4. Runtime
4.1 Packed storage
Base-3 packing fits five ternary digits in eight bits (243 ≤ 256), so
each weight costs
Per-layer overhead: one float32α (4 B), plus optionally a fp16
per-row scale vector ρ ∈ ℝ^O (2 × O B). For a Gemma 4 31B IT linear
with O = I = 5376, the per-row scale adds < 0.1 % to the storage.
4.2 Fused Triton matmul
The kernel loads the 243-entry per-layer LUT into shared memory once
per block, then runs a standard K-sliced GEMM where each load of the
packed code byte is immediately dereferenced through the LUT. Because
decoding is a single shared-memory lookup, no per-thread arithmetic
depends on the ternary structure — decoding is free.
The optional per-row scale ρ_n is applied to the output accumulator
before the bias add, so it does not increase memory pressure inside
the GEMM.
4.3 Inference modes
Four runtime modes are provided (throughput measured on H200, bf16,
bs=4, 64 new tokens, Gemma 4 31B IT):
mode
matmul
tok/s
mem (GB)
packed_decode_per_call
decode full W each call
1.5
31
packed_cached_decode
decode once, cache bf16
~42
60
packed_triton_fused(default)
LUT matmul
37.2
31
packed_triton_autotuned
LUT matmul, autotuned
26.9
31
packed_triton_fused is the default: it preserves the memory
advantage of the packed codes (~31 GB for the full model) while
reaching ~88 % of the throughput of a full bf16 decode.
4.4 Correctness
Output-level mean-relative error between the fused kernel and a
reference bf16 matmul is < 0.0 at fp32 precision; at bf16 we observe
~0 in the kernel itself, and all reconstruction error is attributable
to the ternary encoding step. Against a dense bf16 reference, a packed
Gemma 4 31B IT layer has mean-relative matmul error of 3 – 6 % at
layer 0 of the transformer stack (dropping to 1.5 % with per-row
scale).
5. Pipeline
We ship the end-to-end pipeline as four independent scripts:
01_fit_route_specs.sh — fit the 5-entry ladder for each
(family × depth-bucket) group from ~200 k samples per layer,
~2 min on one H200.
02_export_packed_checkpoint.sh — encode all 410 linear
layers and pack, ~40 min on one H200 (parallelisable across GPUs).
03_add_row_scale.sh — recompute per-row scale with a
regression guard, ~1 min.
04_validate_ppl.sh — measure WikiText-2 perplexity of the
packed checkpoint vs dense baseline, ~5 min.
The runtime folder (runtime/) contains a load_model.py that loads
the dense Gemma 4 31B IT and swaps every linear layer with its packed
counterpart; the resulting transformers model is a drop-in
replacement.
6. Experiments
6.1 Setup
Base model: google/gemma-4-31B-it, bf16.
Hardware: NVIDIA H200 (141 GB HBM), CUDA 12.4.
Packed coverage: all 410 linear layers inside the 60 transformer
blocks (including the 10 shared-KV blocks with v_proj = None,
whose k_proj doubles as v_proj).
Ladder fitting: family-depth, coordinate descent, 200 000 samples
per layer.
6.2 Weight reconstruction
Across the 410 layers, average per-layer weight-reconstruction MSE
(relative to Σ W²):
variant
mean relMSE
best layer
worst layer
Route B, global scale only
0.331 %
0.083 %
1.102 %
Route B + per-row scale (regression-guarded)
0.067 %
0.014 %
1.102 %
The per-row scale is kept in 381 / 410 layers. The 29 layers where it
regresses are concentrated in the very last block (layers.59) and
the output projection down_proj, where the per-row maximum is a
poor proxy for the typical magnitude of the matrix.
6.3 Output-level matmul error
On 8 test input batches, layer 0 of Gemma 4 31B IT:
projection
Route B only
+ per-row scale
reduction
k_proj
3.77 %
1.17 %
3.2 ×
q_proj
3.66 %
1.16 %
3.2 ×
v_proj
4.95 %
1.88 %
2.6 ×
o_proj
6.22 %
2.01 %
3.1 ×
gate_proj
4.70 %
1.52 %
3.1 ×
up_proj
3.59 %
1.43 %
2.5 ×
down_proj
5.90 %
2.16 %
2.7 ×
Mean output relative error drops from 4.54 % to 1.57 % (×2.89).
6.4 WikiText-2 perplexity across context lengths
We evaluate on the WikiText-2 test split, fixing total tokens at
~32 k and sweeping the context length from 512 up to 16 k. At ≥ 16 k
the 60 GB dense bf16 model no longer leaves room for the full
attention score tensor on a single H200 with math-SDPA; we therefore
report a packed-only number at 16 k. All runs use the Gemma 4 31B IT
tokenizer and do_sample=False.
setting (batches × ctx)
tokens
variant
loss
ppl
Δ vs dense
8 × 512
4 k
dense bf16
8.9196
7477.0
—
Route B (global scale only)
9.2035
9931.8
+0.2839
Route B + per-row scale
9.1666
9572.2
+0.2470
32 × 1024
32 k
dense bf16
7.8438
2549.8
—
Route B (global scale only)
8.0771
3219.7
+0.2333
Route B + per-row scale
8.1734
3545.3
+0.3296
16 × 2048
32 k
dense bf16
6.8292
924.4
—
Route B (global scale only)
6.9059
998.2
+0.0767
Route B + per-row scale
7.2812
1452.7
+0.4520
8 × 4096
32 k
dense bf16
6.8528
948.2
—
Route B (global scale only)
7.0365
1137.5
+0.1837
Route B + per-row scale
7.3864
1612.2
+0.5336
Route B + act-aware selection
7.3918
1620.9
+0.5390
4 × 8192
32 k
dense bf16
7.3337
1530.7
—
Route B (global scale only)
7.5483
1897.3
+0.2146
Route B + per-row scale
7.9786
2919.7
+0.6449
Route B + act-aware selection
7.9530
2846.9
+0.6193
2 × 16384
32 k
dense bf16
—
—
—
Route B (global scale only)
8.3255
4117.6
—
Route B + per-row scale
8.5604
5214.0
—
Route B + act-aware selection
8.5390
5107.7
—
Route B + per-head prosody
8.3461
4213.9
+0.02 vs baseline
Route B + per-head prosody + keep layer 37 k_proj
8.3172
4093.8
−0.008 vs baseline
For reference, the per-head prosody rows at shorter contexts:
ctx
Route B + per-head prosody
Δ vs dense
Δ vs baseline
4096
6.9735 (ppl 1067.97)
+0.1207
−0.0630
8192
7.5540 (ppl 1908.37)
+0.2203
+0.0057
16384
8.3461 (ppl 4213.85)
—
+0.0206
If we override the local selector and force-keep only
model.language_model.layers.37.self_attn.k_proj, the same variant
improves further:
ctx
Route B + per-head prosody + keep layers.37.k_proj
Δ vs dense
Δ vs baseline
4096
6.9718 (ppl 1066.10)
+0.1190
−0.0647
8192
7.5438 (ppl 1889.08)
+0.2101
−0.0045
16384
8.3172 (ppl 4093.75)
—
−0.0083
Four findings:
Route B without row-scale is the strong short-path default.
The gap to dense in nats is 0.28 → 0.23 → 0.08 → 0.18 → 0.21 at
512/1k/2k/4k/8k. The minimum gap (+0.08 nats) happens at 2 k —
longer contexts amplify compounded attention errors, but only
mildly.
Per-row scale is monotonically worse from 1 k context onwards.
The gap widens from +0.25 (512) → +0.33 (1 k) → +0.45 (2 k) →
+0.53 (4 k) → +0.64 (8 k). This happens despite per-row scale
reducing weight-space MSE by 5× and per-layer matmul error by
~3×. Weight-level MSE is therefore an unreliable proxy for
end-to-end loss at scale.
Naive activation-aware selection remains misaligned.
We implemented a per-layer activation-aware MSE criterion that
recomputes the keep/drop decision using a 4-batch calibration at
2 k context (experiments/select_row_scale_activation_aware.py).
It demotes only 9 / 381 layers (all k_proj in the middle third
of the stack) and recovers < 0.03 nats end-to-end. Conclusion:
local output error is not the right objective; the damage is
done by compounded noise through the softmax / residual stack,
and catching that requires either a hidden-state-trajectory
criterion or a short QAT phase on ρ (cf. §8.1, §8.4). The same
issue reappears in the successful per-head checkpoint: the local
selector drops 3 k_proj layers, but brute-force end-to-end probes
show that one of them (layers.37.k_proj) should actually be kept.
Per-head prosody eliminates the regression. A single ρ per
attention head (on q/k/v_proj only, 170 projections, head_dim = 256)
keeps between-head dynamic-range adaptation while erasing
within-head multiplicative noise. Activation-aware selection
then keeps 167 / 170 projections (3 dropped, all k_proj in
middle blocks 37, 40, 41). On WikiText-2 the base result is 6.9735
at 4 k (−0.063 nats vs baseline), 7.5540 at 8 k (+0.006),
and 8.3461 at 16 k (+0.021). If we force-keep only the locally
rejected layers.37.k_proj, this improves to 6.9718 / 7.5438 /
8.3172 — i.e. strictly better than the global-scale baseline at
all tested long contexts. This confirms the diagnosis: the
regression of per-row ρ is a softmax-compounding artefact of
within-head row-wise noise, not a property of multiplicative
prosody itself. See §8.2 for the broader prosody discussion and
experiments/add_per_head_prosody_to_checkpoint.py and
experiments/force_perhead_layers_checkpoint.py for the encoder
and the end-to-end override probe.
We also verified that the fp16 storage of ρ_n is not the
culprit: the row_scale values span only [0.05, 1.0] across all 381
kept layers (no subnormals, no saturation), so fp16 representation
is essentially exact at this dynamic range.
The current default checkpoint shipped in runtime/ is therefore
still the global-scale variant (checkpoint_baseline.pt); the
per-row variant, the activation-aware variant, the base per-head
variant, and the stronger long-context research checkpoint
checkpoint_perhead_force_k37.pt are included as research artefacts.
6.5 Memory and throughput
At inference, with packed_triton_fused:
quantity
dense bf16
packed
weight memory (410 linear layers)
60.1 GB
31.0 GB
peak GPU memory (bs=4, 64 new tok)
~60 GB
58.3 GB †
decode throughput (bs=4, 64 new tok)
n/a ‡
37.2 tok/s
† Peak includes KV cache and activations; resident weight memory after
warmup is 31 GB. ‡ A full dense bf16 decode of Gemma 4 31B IT does not
fit in 80 GB with meaningful batch, so this is not a realistic
baseline.
The single-batch latency penalty is ~1.5× (vs packed_cached_decode)
for a ~2× memory reduction.
6.6 Preliminary downstream comparisons
Beyond WikiText-2, we built a common OpenAI-compatible serving path so
that the dense baseline and the best current long-context Route
checkpoint (checkpoint_perhead_force_k37.pt) could be queried by the
same harness. We report these numbers as early dense-vs-route
comparisons, not as final benchmark claims: the completions-based
path currently underestimates absolute Gemma 4 31B IT quality relative
to the official dense product sheet, GPQA remains parser-limited, and
AIME remains prompt-limited in the current path.
benchmark
dense baseline
Route B
Route - dense
status
MMLU-Pro sample
0.5429
0.4500
-0.0929
usable relative regression
HumanEval
0.0854
0.0915
+0.0061
usable relative regression
GPQA Diamond sample
0.0625
0.0625
0.0000
parser-limited
AIME24 sample
0.0000
0.0000
0.0000
invalid under current prompt path
AIME25 sample
0.0000
0.0000
0.0000
invalid under current prompt path
Three observations matter here.
MMLU-Pro is the cleanest current external comparison. On the
capped sample harness, Route trails the dense model by 9.3 points.
HumanEval is roughly neutral at current scale. The same harness
gives a slight Route advantage (+0.0061 absolute), which is small
but at least shows that the packed model is not uniformly degraded.
The harness itself is still incomplete. A direct manual probe on
the first AIME24 sample returned the correct answer 33 for both
dense and Route once prompted with a stricter final-answer cue,
which means the current zeroes are a benchmark-path artefact, not a
clean model-quality measurement.
Separately, a full direct-HF MMLU-Pro evaluation for both dense and
Route was launched with max_length=8192 and max_gen_toks=64 while
this draft was being finalised. Those runs were still in progress at
the time of writing, so we do not cite them here as completed results.
7. Limitations
Post-training only. We do not fine-tune. Route B is evaluated
strictly as a PTQ method on top of the publicly released Gemma 4 31B
IT weights.
Weight-only. Activations and KV cache remain bf16. Reducing
these is orthogonal to Route B.
Evaluation coverage. The strongest settled quality table in this
paper is still the WikiText-2 sweep above. We now also have capped
dense-vs-route comparisons on MMLU-Pro sample and HumanEval (§6.6),
but these are still regression-oriented harnesses rather than final
leaderboard numbers. GPQA remains parser-limited, AIME remains
prompt-limited in the current path, and the full direct-HF MMLU-Pro
runs for both dense and Route were still in progress at the time of
writing. Extending this into a stable full benchmark suite remains a
necessary next step.
Per-row scale ambiguity. On Gemma 4 31B IT the per-row scale
improves weight MSE consistently, but its end-to-end effect is
mixed; the evidence now points away from naive per-layer
activation-aware selection and towards trajectory-level or end-to-end
selection criteria (§8.4).
Shared-KV layers. Ten transformer blocks in Gemma 4 31B IT have
v_proj = None (their k_proj doubles as v_proj). Our export
treats these correctly but does not jointly fit the K/V ladder;
we expect a small additional gain from joint fitting.
8. A discrete weight language — beyond Route B
Route B is the simplest instance of a broader programme. The
organising idea is that every weight scalar is a word in a small
alphabet, and the alphabet, the word length, and the prosody can all
be made adaptive.
8.1 Variable-depth words
Not every row of every layer needs five ternary steps. Looking at the
codebook histogram for Gemma 4 31B IT, the top 45 of the 243 codes
cover 99 % of all weights, and the top 18 cover 80 %. This directly
suggests:
Variable-depth encoding — stop the ternary path early for
weights well-approximated by the first 3 or 4 steps, storing the
stop-depth in a separate uint3 field. We estimate a further ~10
% compression with unchanged MSE.
7-bit truncated alphabet — if only 45 codes per layer are used,
we can store a code in 7 bits (plus a 243 → used-code table of ~8
KB per layer). This saves 3.66 GiB on Gemma 4 31B IT with no
quality change.
8.2 Per-row, per-column and per-head prosody
The per-row scale is a one-dimensional prosody on the output
direction. A symmetric per-column scale would do the same on the
input direction; together they lift the model to a rank-2
multiplicative envelope around the discrete codes.
For attention layers, a per-head prosody — one ρ per attention
head — turned out to be the right granularity in our experiments
(§6.4): it gives every head its own dynamic range but forces rows
within a head to share a scale, which eliminates the within-head
row-wise multiplicative noise that compounds through softmax. With
this single change (experiments/add_per_head_prosody_to_checkpoint.py,
head_dim = 256, activation-aware keep-or-drop at ctx = 2 k, 167/170
projections kept), we go from +0.53 nats at 4 k with per-row scale
to −0.06 nats at 4 k with per-head prosody (vs the global-scale
baseline). A one-layer end-to-end correction improves this further:
force-keeping only layers.37.k_proj yields 6.9718 / 7.5438 / 8.3172
at 4 k / 8 k / 16 k, strictly beating the global-scale baseline at all
three long-context points. We also tested the obvious follow-up of
stacking per-row scale back onto the MLP projections and found that it
still regresses at 4 k (7.0689 for gate+up, 7.0890 for down,
7.1284 for the full MLP stack). So, with current objectives, the
evidence says prosody should stay attention-local.
8.3 Adaptive alphabets
Instead of a single ladder S = (s₁, …, s₅) shared by a family × depth
bucket, one could fit per-row or per-group-of-32-rows ladders
at a cost of ~10 fp16 numbers per group. Preliminary analysis
(alternating least squares, fit_route_spec_alternating) suggests
that moving from one ladder per family to one ladder per 32-row block
reduces MSE by a further 1.5 – 2 ×, at a storage cost of < 0.5 %.
8.4 Activation-aware selection (and why the naive version fails)
The failure mode of §6.4 (better weight MSE, strictly worse end-to-end
loss at ≥ 1 k context) is a direct consequence of minimising the
wrong objective. We implemented the obvious fix — per-layer
activation-weighted MSE,
L_act(layer) = Σ_x || (W − W_hat) · x ||² / Σ_x || W · x ||²
computed on 4 calibration batches of 2 k context, and re-deciding
the keep/drop per layer. The selector
(experiments/select_row_scale_activation_aware.py) demotes only 9
of 381 layers (all k_proj in middle-block positions 20, 36–44) and
the resulting checkpoint recovers < 0.03 nats end-to-end at 4–16 k
context. Local output error is locally better with row_scale, as
the criterion correctly detects — but the damage at long context
accumulates through softmax and the residual stream, not through any
single layer's output. A correct criterion therefore needs to be
either trajectory-level (measure final hidden-state deviation
with the candidate layer toggled) or end-to-end (differentiable
surrogate objective on ρ). A small brute-force probe on the per-head
checkpoint makes the same point in miniature: among the 3 locally
dropped k_proj layers, only layers.37.k_proj should be restored;
the other two remain regressions. This is exactly the type of signal a
trajectory-level selector should capture automatically. §8.6 takes the
second route.
8.5 Joint attention fit
For the 10 shared-KV layers where v_proj = None and k_proj doubles
as v_proj, jointly fitting (q, k, o) with a single objective
(preserving softmax(qk^T) v · o) should close most of the residual
gap on those blocks. The repository already contains the primitives
for this (fit_attention_vo_joint).
8.6 Towards QAT-grade Route B
The above extensions are all post-hoc. A natural endpoint is a short
distillation phase on top of the packed model, with gradients flowing
only through the per-row prosody and the 45-odd used code-book entries
(a tiny fraction of the parameter count). This keeps the discrete
structure of the representation intact while eliminating the last of
the end-to-end loss gap.
8.7 What a "weight language" is actually good for
We end with a speculation. A ternary code is a finite discrete object.
Unlike a bf16 weight, a Route B code can be:
compared — two codes for the same position across two checkpoints
can be diffed;
searched — codes that realise a particular computational role
(e.g. copying a specific key vector) can be located and clustered;
edited — one can imagine targeted interventions at the level of
the "letters" of a weight, not the whole matrix.
We believe this angle — compression as a side-effect of obtaining a
manipulable discrete representation of the weights — is where the
long-term value of the programme lies.
9. Conclusion
We have presented Route B, a 3-bit weight-compression scheme with a
fused Triton matmul, applied end-to-end to Gemma 4 31B IT. The method
reduces weight memory from 60 GB to 32 GB and preserves 37 tok/s
throughput at batch 4 on a single H200. On WikiText-2 at 2 k context
the packed model is within +0.08 nats of the dense bf16 baseline
without any fine-tuning; across 1 k–16 k context it stays within
+0.08 to +0.21 nats. We add a per-row scale that drives weight
reconstruction MSE down 5× and matmul output error down 2.9× — and
empirically show that this weight-level improvement translates to a
strictly worse end-to-end model from 1 k context upward, even when
the keep/drop decision is made with a per-layer activation-weighted
MSE criterion. This is a direct refutation of weight-MSE as a
sufficient PTQ objective, and a warning that standard per-layer
activation-aware selection is insufficient on long-context models.
We then show that switching the attention prosody from per-row to
per-head (167 / 170 projections kept by activation-aware
selection) eliminates the regression, and that a one-layer end-to-end
override improves it further: force-keeping layers.37.k_proj yields
6.9718 / 7.5438 / 8.3172 at 4 k / 8 k / 16 k, beating the global-scale
baseline at all three long-context points. This is the first concrete
positive result from the discrete-language programme beyond the core
Route B encoder, and the new default direction for the 3-bit envelope.
We release the full 4-step pipeline, four variants of the packed
checkpoint (global-scale, per-row, per-head, per-head+k37), and a
ready-to-run runtime.
Outside WikiText-2, the current dense-vs-route harness already gives a
useful early signal: on a capped MMLU-Pro sample the best current Route
checkpoint trails the dense model by 9.3 points, while on HumanEval it
is roughly neutral/slightly positive. We do not yet treat GPQA or
AIME under the current completions path as final numbers, and the full
direct-HF MMLU-Pro evaluations for dense and Route were still running
when this version was prepared. The present paper should therefore be
read as a strong compression / long-context evaluation paper with
preliminary downstream coverage, not yet as a finished benchmark-paper
replacement for the official Gemma evaluation sheet.
Beyond the concrete result, we framed the method as a first instance
of a discrete weight language — a view in which every layer is a
sentence in a small alphabet, and the alphabet, word length and
prosody become legitimate design knobs. We see this framing, rather
than the specific ternary recipe, as the main contribution of this
work. The concrete improvement roadmap (variable-depth words,
per-head prosody, adaptive alphabets, trajectory-level
activation-aware selection, short QAT on ρ) is laid out in §8 and in
a separate release document, benchmarks/improvement_roadmap.md.
Appendix A. Reproducibility checklist
Base model: google/gemma-4-31B-it (bf16), downloaded from
HuggingFace under Gemma's public licence.
Hardware: NVIDIA H200 (141 GB HBM). Will run on any GPU with ≥ 40 GB.
Libraries (pinned): torch==2.5.1, triton==3.1.0,
transformers==5.5.0. Full list in runtime/requirements.txt.
Seeds: ladder-fitting sampling uses seed 0. Evaluation is
deterministic (do_sample=False).
Scripts: reproduce/01..04_*.sh, runnable individually or via
reproduce/run_all.sh. Total budget ≈ 50 min on one H200.
This repository contains a full all-linear Hugging Face export built on
top of google/gemma-4-31B-it from the Route B / discrete-weight-language
pipeline.
What it is
This project frames model compression as a move away from unconstrained
fp16 weights toward a small, structured, discrete weight language.
Each linear weight is represented through a compact ternary route and
can then be materialized back into a standard Hugging Face model.
The specific exported model in this release comes from the strongest
current long-context checkpoint, which uses per-head prosody in attention
and one end-to-end keep override for layers.37.k_proj. That source
checkpoint name should be treated as an implementation detail, not as the
public project title.
Contents
full all-linear materialized Hugging Face weights;
minimal runtime under runtime/;
minimal conversion/publishing entrypoints under kernel/;
short architecture note in kernel/ARCHITECTURE.md.
How to load
If this repo contains the packed .pt checkpoint, use the Route B
runtime in runtime/load_model.py.
If this repo contains a materialized Hugging Face export, load it as a
standard transformers model.
Important notes
This is a derived release based on google/gemma-4-31B-it.
Redistribution requires compliance with the Gemma license.
The local export metadata for this model records 410 materialized
weight modules across attention and MLP projections, i.e. the full
all-linear Route conversion rather than a one-layer toy example.
The strongest settled quality numbers are in the accompanying paper
and archiv.org/benchmarks/results.md.
Reference
Arman Aubakirov. A Discrete Weight Language for Large Language Models:
Compressing Gemma 4 31B with a 5-Step Ternary Route. Technical report,
April 2026.