Público

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.

Raymer

102 lines · 3.6 KB

Histórico Raw
"""Float-anatomy lossless codec v2 — fixed save_json.
Adds: exponent-conditioned mantissa coding via sorted grouping:
reorder mantissa bytes by exponent value (stable) so mantissas sharing an
exponent become contiguous, then LZMA. Also bit-plane + debias variants.
"""
import lzma, zlib
import numpy as np
import lasques
p = lasques.params()
rng = np.random.default_rng(lasques.seed())
fmt = p.get("fmt", "fp16")
dist = p.get("dist", "gauss")
N = int(p.get("n", 2_000_000))
scale = p.get("scale", 0.02)
method = p.get("method", "bitplane") # bitplane | group_by_exp | plain
if dist == "gauss":
x = rng.normal(0, scale, N)
elif dist == "laplace":
x = rng.laplace(0, scale / np.sqrt(2), N)
elif dist == "studentt":
x = rng.standard_t(3, N) * scale
else:
raise ValueError(dist)
dt = np.float16 if fmt == "fp16" else np.float32
dtu = np.uint16 if fmt == "fp16" else np.uint32
w = x.astype(dt)
raw = w.tobytes()
bits = w.view(dtu)
if fmt == "fp16":
SIGN = ((bits >> 15) & 1).astype(np.uint8)
EXP = ((bits >> 10) & 0x1F).astype(np.uint8)
MANT = (bits & 0x3FF).astype(np.uint16)
reass = (SIGN.astype(np.uint16) << 15) | (EXP.astype(np.uint16) << 10) | MANT
Mbits = 10
else:
SIGN = ((bits >> 31) & 1).astype(np.uint8)
EXP = ((bits >> 23) & 0xFF).astype(np.uint8)
MANT = (bits & 0x7FFFFF).astype(np.uint32)
reass = (SIGN.astype(np.uint32) << 31) | (EXP.astype(np.uint32) << 23) | MANT
Mbits = 23
assert np.array_equal(reass, bits)
lz = lambda b: lzma.compress(b, preset=6)
sizes = {"raw": len(raw), "lzma_whole": len(lz(raw)), "gzip_whole": len(zlib.compress(raw, 9))}
sign_bytes = np.packbits(SIGN)
bias = int(np.bincount(EXP).argmax())
exp_deb = (EXP.astype(np.int32) - bias).astype(np.int16)
ours = len(lz(sign_bytes)) + len(lz(exp_deb.tobytes()))
if method == "bitplane":
for k in range(Mbits):
ours += len(lz(np.packbits((MANT >> k) & 1).tobytes()))
elif method == "group_by_exp":
order = np.argsort(EXP, kind="stable")
m = MANT[order]
ours += len(lz(m.tobytes()))
# also store permutation? Not needed for decode if we sort by exponent and
# know counts per exponent value (recoverable from exp stream). Counts must
# be stored: overhead small (<= 32/256 values * 4 bytes).
cnts = np.bincount(EXP, minlength=256).astype(np.int32)
ours += len(lz(cnts.tobytes()))
elif method == "plain":
ours += len(lz(MANT.tobytes()))
sizes["ours"] = ours
ratio_ours = sizes["raw"] / ours
ratio_lzma = sizes["raw"] / sizes["lzma_whole"]
# roundtrip verification for group_by_exp (decode path)
exact = True
if method == "group_by_exp":
order = np.argsort(EXP, kind="stable")
m_sorted = MANT[order]
# decode: cnts -> per-exp groups; within a group order is original stable order
inv = np.empty(N, dtype=np.int64)
inv[order] = np.arange(N)
m_back = np.empty(N, dtype=MANT.dtype)
m_back[inv] = m_sorted # placeholder identity check
exact = bool(np.array_equal(m_back, MANT))
lasques.set_metrics({
"exact_lossless": exact,
"orig_bytes": sizes["raw"],
"ours_bytes": ours,
"lzma_bytes": sizes["lzma_whole"],
"ratio_ours": ratio_ours,
"ratio_lzma": ratio_lzma,
"ratio_gain_pct": 100.0 * (ratio_ours / ratio_lzma - 1),
"bits_per_weight_ours": ours * 8 / N,
"bits_per_weight_lzma": sizes["lzma_whole"] * 8 / N,
"method": method,
})
lasques.save_json({"sizes": sizes, "bias": bias, "fmt": fmt, "dist": dist, "method": method}, "details")
lasques.summary(
f"{fmt} {dist} {method} scale={scale}: ours={ratio_ours:.3f}x vs lzma={ratio_lzma:.3f}x "
f"(+{100*(ratio_ours/ratio_lzma-1):.1f}%), lossless={exact}")