Lossless compression of an ML embedding cube

decorrelate_then_compress.py

"""Lossless compression of an ML embedding cube (H, W, C) by treating it as an
image with C channels -- and beating an image codec at it.

The trick is two stages, in this order:

  1. DECORRELATE THE CHANNELS. A photo codec models redundancy across *space*
     (neighbouring pixels) but not across *channels*. An embedding's C dimensions
     are heavily correlated (here 32 of 128 PCA directions hold 96% of the
     variance), so we first remove that with a reversible, PCA-like rotation:
     predict each channel from the already-seen channels via the Cholesky factor
     of the channel covariance. It is integer-reversible because the decoder
     rebuilds channels in order and replays the same rounded prediction -- no
     floating-point round-trip error, no lifting decomposition needed.

  2. COMPRESS SPATIALLY. Hand the decorrelated channels to JPEG XL (lossless),
     which is excellent at the spatial redundancy the first stage left behind.

On a 1135x733x128 int8 satellite-embedding cube this reaches ~1.66x lossless --
~20% smaller than the best GeoTIFF, and better than JPEG XL on its own (which,
built for 3-4 colour channels, treats dims 4+ as independent "extra channels"
and so never decorrelates across them).

    pip install numpy imagecodecs        # imagecodecs bundles libjxl
"""
import numpy as np
import imagecodecs


def _decorrelator(x):
    """Strictly-lower predictor weights from the channel covariance (LDL factor)."""
    c = x.shape[1]
    xc = x - x.mean(0)
    cov = (xc.T @ xc) / x.shape[0]
    cov += np.eye(c) * 1e-3 * np.trace(cov) / c            # ridge for conditioning
    lunit = (lambda L: L / np.diag(L))(np.linalg.cholesky(cov))   # unit lower-triangular
    return (np.linalg.inv(lunit) - np.eye(c)).astype(np.float32)  # strictly-lower part


def _predict(src, w, c):
    """Rounded prediction of channel c from channels 0..c-1 (bit-identical enc/dec)."""
    return np.rint(src[:, :c] @ w[c, :c]) if c else np.zeros(src.shape[0], np.float32)


def encode(cube):
    """(H, W, C) int8 -> (jxl_bytes, weights, offset, shape). Lossless."""
    h, w, c = cube.shape
    x = cube.reshape(-1, c).astype(np.float32)
    weights = _decorrelator(x)
    r = np.empty_like(x)
    for ch in range(c):                                    # stage 1: decorrelate channels
        r[:, ch] = x[:, ch] + _predict(x, weights, ch)
    offset = int(r.min())                                  # shift to non-negative ints
    res16 = np.ascontiguousarray((r - offset).astype(np.uint16).reshape(h, w, c))
    jxl = imagecodecs.jpegxl_encode(res16, lossless=True, effort=7)   # stage 2: spatial
    return jxl, weights, offset, (h, w, c)


def decode(jxl, weights, offset, shape):
    """Inverse of encode -> the original (H, W, C) int8 cube."""
    h, w, c = shape
    r = imagecodecs.jpegxl_decode(jxl).reshape(-1, c).astype(np.float32) + offset
    x = np.empty_like(r)
    for ch in range(c):                                    # undo stage 1, channel by channel
        x[:, ch] = r[:, ch] - _predict(x, weights, ch)
    return x.astype(np.int8).reshape(shape)


if __name__ == "__main__":
    import sys

    if len(sys.argv) > 1:                                  # real data: python … cube.npy
        cube = np.load(sys.argv[1]).astype(np.int8)
    else:                                                  # else a correlated synthetic cube
        rng = np.random.default_rng(0)
        latent = rng.standard_normal((256, 256, 16)).astype(np.float32)
        mix = rng.standard_normal((16, 128)).astype(np.float32)          # 128 dims from 16
        noise = 0.1 * rng.standard_normal((256, 256, 128)).astype(np.float32)
        cube = np.clip((latent.reshape(-1, 16) @ mix).reshape(256, 256, 128) + noise,
                       -127, 127).astype(np.int8)

    jxl, weights, offset, shape = encode(cube)
    back = decode(jxl, weights, offset, shape)
    raw = cube.size                                        # 1 byte/sample (int8)
    stored = len(jxl) + weights.nbytes                     # decorrelator travels with the file
    print(f"shape {shape}  lossless: {np.array_equal(back, cube)}")
    print(f"raw {raw:,} B  ->  {stored:,} B   ({raw / stored:.2f}x)")