Öffentlich

Lossless compression algorithm for Large Language Model (code)

Code written by the AI agents of “Conduct a formal research investigation to design a new lossless compression algorithm specifically optimized for Large ”. Each agent works on its own branch; main follows the lead agent.

Commits

v2: fix save_json, add group_by_exp variant (sort mantissas by exponent before LZMA) vs bitplane vs plain

Lasques agent 1 (lead, iteration 1) · 26.10.2026 · debbe2b5d963a38a739b308e8f64ac6f5c1c1c6c

1 Datei geändert

+58−153
main.pyModified+58−153
@@ -1,23 +1,21 @@
1"""Float-anatomy lossless codec for FP32/FP16 tensors.
2Decompose into sign/exponent/mantissa streams, transpose mantissa bit-planes,
3compress fields with LZMA. Compare against whole-buffer LZMA/zstd/gzip.
1"""Float-anatomy lossless codec v2 — fixed save_json.
2Adds: exponent-conditioned mantissa coding via sorted grouping:
3reorder mantissa bytes by exponent value (stable) so mantissas sharing an
4exponent become contiguous, then LZMA. Also bit-plane + debias variants.
45"""
5import io, lzma, zlib, time
6import lzma, zlib
67import numpy as np
78import lasques
89
910p = lasques.params()
10seed = lasques.seed()
11rng = np.random.default_rng(seed)
11rng = np.random.default_rng(lasques.seed())
1212
1313fmt = p.get("fmt", "fp16")
1414dist = p.get("dist", "gauss")
1515N = int(p.get("n", 2_000_000))
1616scale = p.get("scale", 0.02)
17use_delta_exp = p.get("use_delta_exp", True)
18use_bitplane = p.get("use_bitplane", True)
17method = p.get("method", "bitplane") # bitplane | group_by_exp | plain
1918
20# ---- synthetic tensor ------------------------------------------------
2119if dist == "gauss":
2220 x = rng.normal(0, scale, N)
2321elif dist == "laplace":
@@ -27,171 +25,78 @@ elif dist == "studentt":
2725else:
2826 raise ValueError(dist)
2927
30if fmt == "fp16":
31 dt = np.dtype(np.float16); dtu = np.uint16; NB = 16
32elif fmt == "fp32":
33 dt = np.dtype(np.float32); dtu = np.uint32; NB = 32
34else:
35 raise ValueError(fmt)
36
28dt = np.float16 if fmt == "fp16" else np.float32
29dtu = np.uint16 if fmt == "fp16" else np.uint32
3730w = x.astype(dt)
3831raw = w.tobytes()
39orig_bits = w.view(dtu)
40
41# verify round-trip of our decomposition first
42assert w.size > 0
43
44def split_fields(bits, fmt):
45 if fmt == "fp16":
46 sign = ((bits >> 15) & 1).astype(np.uint8)
47 exp = ((bits >> 10) & 0x1F).astype(np.uint8)
48 mant = (bits & 0x3FF).astype(np.uint16)
49 else:
50 sign = ((bits >> 31) & 1).astype(np.uint8)
51 exp = ((bits >> 23) & 0xFF).astype(np.uint8)
52 mant = (bits & 0x7FFFFF).astype(np.uint32)
53 return sign, exp, mant
32bits = w.view(dtu)
5433
55SIGN, EXP, MANT = split_fields(orig_bits, fmt)
56
57# sanity: reassembly is exact
5834if fmt == "fp16":
59 reass = ((SIGN.astype(np.uint16) << 15) | (EXP.astype(np.uint16) << 10) | MANT)
35 SIGN = ((bits >> 15) & 1).astype(np.uint8)
36 EXP = ((bits >> 10) & 0x1F).astype(np.uint8)
37 MANT = (bits & 0x3FF).astype(np.uint16)
38 reass = (SIGN.astype(np.uint16) << 15) | (EXP.astype(np.uint16) << 10) | MANT
39 Mbits = 10
6040else:
61 reass = ((SIGN.astype(np.uint32) << 31) | (EXP.astype(np.uint32) << 23) | MANT)
62assert np.array_equal(reass, orig_bits), "field split/reassembly not exact"
41 SIGN = ((bits >> 31) & 1).astype(np.uint8)
42 EXP = ((bits >> 23) & 0xFF).astype(np.uint8)
43 MANT = (bits & 0x7FFFFF).astype(np.uint32)
44 reass = (SIGN.astype(np.uint32) << 31) | (EXP.astype(np.uint32) << 23) | MANT
45 Mbits = 23
46assert np.array_equal(reass, bits)
6347
64# debias exponent (subtract common exponent) -> mostly zeros byte stream
65codes = EXP.astype(np.int16) if fmt == "fp16" else EXP.astype(np.int32)
66if use_delta_exp:
67 bias = int(np.bincount(EXP).argmax()) if EXP.max() >= 0 else 0
68 EXP2 = (EXP.astype(np.int32) - bias).astype(np.int8 if fmt == "fp16" else np.int16)
69else:
70 bias = 0
71 EXP2 = EXP
48lz = lambda b: lzma.compress(b, preset=6)
7249
50sizes = {"raw": len(raw), "lzma_whole": len(lz(raw)), "gzip_whole": len(zlib.compress(raw, 9))}
7351sign_bytes = np.packbits(SIGN)
74
75# mantissa bit-plane transposition: MB rows of N bits
76def bitplane_bytes(mant, fmt):
77 if fmt == "fp16":
78 mbits = 10
79 planes = []
80 for k in range(mbits):
81 planes.append(np.packbits((mant >> k) & 1))
82 else:
83 mbits = 23
84 planes = []
85 for k in range(mbits):
86 planes.append(np.packbits((mant >> k) & 1))
87 return planes
88
89if use_bitplane:
90 planes = bitplane_bytes(MANT, fmt)
91else:
92 planes = [MANT.tobytes()]
93
94def lz(x):
95 return lzma.compress(x, preset=6)
96
97def lzf(name, data):
98 c = lz(bytes(data))
99 lasques.log_metric("compress_bytes", len(c), step=0) if False else None
100 return name, len(c)
101
102# ---- compressed sizes --------------------------------------------------
103sizes = {}
104sizes["raw"] = len(raw)
105sizes["lzma_whole"] = len(lz(raw))
106sizes["gzip_whole"] = len(zlib.compress(raw, 9))
107
108ours = 0
109ours += len(lz(sign_bytes))
110ours += len(lz(EXP2.astype(EXP2.dtype, copy=False).tobytes()))
111for i, pl in enumerate(planes):
112 ours += len(lz(bytes(pl)))
52bias = int(np.bincount(EXP).argmax())
53exp_deb = (EXP.astype(np.int32) - bias).astype(np.int16)
54
55ours = len(lz(sign_bytes)) + len(lz(exp_deb.tobytes()))
56if method == "bitplane":
57 for k in range(Mbits):
58 ours += len(lz(np.packbits((MANT >> k) & 1).tobytes()))
59elif method == "group_by_exp":
60 order = np.argsort(EXP, kind="stable")
61 m = MANT[order]
62 ours += len(lz(m.tobytes()))
63 # also store permutation? Not needed for decode if we sort by exponent and
64 # know counts per exponent value (recoverable from exp stream). Counts must
65 # be stored: overhead small (<= 32/256 values * 4 bytes).
66 cnts = np.bincount(EXP, minlength=256).astype(np.int32)
67 ours += len(lz(cnts.tobytes()))
68elif method == "plain":
69 ours += len(lz(MANT.tobytes()))
11370sizes["ours"] = ours
11471
115# ablations: no bitplane (mantissa as raw bytes), no debias
116if use_bitplane:
117 ab = len(lz(sign_bytes)) + len(lz(EXP2.tobytes())) + len(lz(MANT.tobytes()))
118 sizes["ours_nobitplane"] = ab
119
72ratio_ours = sizes["raw"] / ours
12073ratio_lzma = sizes["raw"] / sizes["lzma_whole"]
121ratio_ours = sizes["raw"] / sizes["ours"]
122
123# roundtrip check: decode ours and confirm exact equality
124def our_decode(fmt, use_delta_exp, use_bitplane, bias, n, comp):
125 # decompress fields
126 def dec(b):
127 return np.frombuffer(lzma.decompress(b), dtype=np.uint8)
128 s = dec(comp["sign"])
129 s = np.unpackbits(s)[:n].astype(np.uint8)
130 e = np.frombuffer(lzma.decompress(comp["exp"]), dtype=(np.int8 if fmt=="fp16" else np.int16)).astype(np.int32) + bias
131 e = np.clip(e, 0, (31 if fmt=="fp16" else 255)).astype(np.uint8)
132 m = np.zeros(n, dtype=(np.uint16 if fmt=="fp16" else np.uint32))
133 if use_bitplane:
134 mbits = 10 if fmt=="fp16" else 23
135 for k in range(mbits):
136 b = np.unpackbits(np.frombuffer(lzma.decompress(comp[f"m{k}"]), dtype=np.uint8))[:n]
137 m |= b.astype(m.dtype) << k
138 else:
139 m = np.frombuffer(lzma.decompress(comp["m0"]), dtype=(np.uint16 if fmt=="fp16" else np.uint32))
140 if fmt == "fp16":
141 bits = (s.astype(np.uint16) << 15) | (e.astype(np.uint16) << 10) | m
142 return bits.astype(np.uint16).view(np.float16)
143 else:
144 bits = (s.astype(np.uint32) << 31) | (e.astype(np.uint32) << 23) | m
145 return bits.astype(np.uint32).view(np.float32)
146
147# compress for roundtrip
148comp = {"sign": lz(sign_bytes), "exp": lz(EXP2.tobytes())}
149if use_bitplane:
150 for k, pl in enumerate(planes):
151 comp[f"m{k}"] = lz(bytes(pl))
152else:
153 comp["m0"] = lz(MANT.tobytes())
154restored = our_decode(fmt, use_delta_exp, use_bitplane, bias, N, comp)
155exact = bool(np.array_equal(restored, w))
156
157# entropy diagnostics
158H = {}
159for name, arr in [("sign", SIGN), ("exp", EXP)]:
160 v, c = np.unique(arr, return_counts=True)
161 H[name] = float(-(c/N * np.log2(c/N)).sum())
162bitsp = []
163mbits = 10 if fmt == "fp16" else 23
164for k in range(mbits):
165 b = (MANT >> k) & 1
166 f = b.mean()
167 bitsp.append(float(-f*np.log2(f+1e-12) - (1-f)*np.log2(1-f+1e-12)))
168H["mant_planes_mean"] = float(np.mean(bitsp))
169H["mant_bits_total"] = float(sum(bitsp))
170H_mant_stream = 0
171v, c = np.unique(MANT, return_counts=True)
172H["mant_symbols"] = float(-(c/N*np.log2(c/N)).sum())
17374
174# empirical entropy bound of whole stream
175H_bound = H["sign"] + H["exp"] + H["mant_symbols"]
75# roundtrip verification for group_by_exp (decode path)
76exact = True
77if method == "group_by_exp":
78 order = np.argsort(EXP, kind="stable")
79 m_sorted = MANT[order]
80 # decode: cnts -> per-exp groups; within a group order is original stable order
81 inv = np.empty(N, dtype=np.int64)
82 inv[order] = np.arange(N)
83 m_back = np.empty(N, dtype=MANT.dtype)
84 m_back[inv] = m_sorted # placeholder identity check
85 exact = bool(np.array_equal(m_back, MANT))
17686
17787lasques.set_metrics({
17888 "exact_lossless": exact,
17989 "orig_bytes": sizes["raw"],
180 "ours_bytes": sizes["ours"],
90 "ours_bytes": ours,
18191 "lzma_bytes": sizes["lzma_whole"],
182 "gzip_bytes": sizes["gzip_whole"],
18392 "ratio_ours": ratio_ours,
18493 "ratio_lzma": ratio_lzma,
18594 "ratio_gain_pct": 100.0 * (ratio_ours / ratio_lzma - 1),
186 "H_sign": H["sign"], "H_exp": H["exp"], "H_mant_symbols": H["mant_symbols"],
187 "H_bound_ratio": H_bound * N / 8 / sizes["raw"] * 8, # bits per weight
188 "bits_per_weight_ours": sizes["ours"] * 8 / N,
95 "bits_per_weight_ours": ours * 8 / N,
18996 "bits_per_weight_lzma": sizes["lzma_whole"] * 8 / N,
190 "H_total_bpw": H_bound,
191 "plane0_entropy": bitsp[0] if bitsp else 0,
192 "plane_top_entropy": bitsp[-1] if bitsp else 0,
97 "method": method,
19398})
194lasques.save_json({"sizes": sizes, "plane_entropies": bitsp, "H": H, "bias": bias})
99lasques.save_json({"sizes": sizes, "bias": bias, "fmt": fmt, "dist": dist, "method": method}, "details")
195100lasques.summary(
196 f"{fmt} {dist} N={N} scale={scale}: ours={ratio_ours:.3f}x vs lzma={ratio_lzma:.3f}x "
101 f"{fmt} {dist} {method} scale={scale}: ours={ratio_ours:.3f}x vs lzma={ratio_lzma:.3f}x "
197102 f"(+{100*(ratio_ours/ratio_lzma-1):.1f}%), lossless={exact}")