| """Geometry-preserving weight analogs for Q2_K bilinear factorization. |
| |
| At 2 bits per weight (4 levels ``q in {0, 1, 2, 3}`` and 4-bit sub-block |
| scales ``Ls, Lm in {0..15}``), raw weight blocks suffer from two structural |
| bottlenecks: |
| |
| 1. **Intra-sub-block outlier collapse**: a single heavy-tailed weight stretches |
| the sub-block span ``xmax - xmin``, forcing 13-15 of the 16 weights into |
| one or two quantization bins. |
| 2. **Inter-sub-block scale crushing**: a single high-variance sub-block inflates |
| the shared superblock scale ``d``, collapsing ``Ls_j`` of the remaining |
| 15 sub-blocks toward small integers with severe rounding loss. |
| |
| Rather than factoring raw weights alone, V2 constructs a bank of |
| **geometry-preserving weight analogs** ``x_tilde`` of the eligible weights: |
| |
| * :func:`subblock_isomoment_analog` maps each 16-weight sub-block through an |
| importance-weighted saturating transport blended with the optimal 4-level |
| quantile spacing, then orthogonally projects back onto the exact weighted |
| mean ``mu_j`` and weighted variance ``sigma_j^2`` of the original sub-block. |
| * :func:`vesica_gram_analog` operates in the D6 / Vesica fold basis |
| ``(x_p + x_{p+128}, x_p - x_{p+128})``, regularizing outliers in the whitened |
| eigen-basis and coloring back so the exact 2x2 weighted Gram matrix and means |
| of the two half-blocks are preserved. |
| * :func:`subscale_lattice_analog` tempers cross-sub-block dispersion toward the |
| 4-bit ``{0..15}`` scale lattice while preserving every sub-block's unit |
| direction, mean, and total block weighted energy. |
| |
| The discrete Q2_K lattice coordinates ``(Ls, Lm, q)`` discovered on each analog |
| are then projected back onto the **true weights** ``x`` via coupled bilinear |
| least-squares, guaranteeing that the final emitted block is strictly optimal |
| under the true objective. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import numpy as np |
|
|
| from .q2k import NS, QK_K |
|
|
| _Q2_IDEAL_QUANTILES = np.linspace(-1.35, 1.35, 16, dtype=np.float64) |
|
|
|
|
| def subblock_isomoment_analog( |
| x: np.ndarray, |
| w: np.ndarray | None = None, |
| tau: float = 1.65, |
| beta: float = 0.25, |
| quantile_blend: float = 0.30, |
| enforce_nonpositive_min: bool = False, |
| ) -> np.ndarray: |
| """Create a sub-block iso-moment analog preserving weighted mean & variance. |
| |
| Invariants per 16-weight sub-block ``j`` (when ``enforce_nonpositive_min`` |
| is ``False``): |
| * ``sum_k w_{j,k} x_tilde_{j,k} == sum_k w_{j,k} x_{j,k}`` |
| * ``sum_k w_{j,k} (x_tilde_{j,k} - mu_j)^2 == sum_k w_{j,k} (x_{j,k} - mu_j)^2`` |
| * strict rank-order preservation with respect to ``x_{j,k}`` (when ``beta == 0``) |
| """ |
| xb = np.asarray(x, dtype=np.float64).reshape(NS, 16) |
| wb = ( |
| np.ones((NS, 16), dtype=np.float64) |
| if w is None |
| else np.asarray(w, dtype=np.float64).reshape(NS, 16) |
| ) |
| wsum = np.maximum(np.sum(wb, axis=1, keepdims=True), 1.0e-30) |
| mu = np.sum(wb * xb, axis=1, keepdims=True) / wsum |
| centered = xb - mu |
| var = np.sum(wb * centered * centered, axis=1, keepdims=True) / wsum |
| sigma = np.sqrt(np.maximum(var, 1.0e-30)) |
| active = (sigma[:, 0] >= 1.0e-15) |
|
|
| z = centered / sigma |
| w_mean = wsum / 16.0 |
| w_rel = np.power(np.maximum(wb / w_mean, 1.0e-12), beta) if beta != 0.0 else 1.0 |
| z_comp = tau * np.tanh((z * w_rel) / max(tau, 1.0e-6)) |
|
|
| if quantile_blend > 0.0: |
| order = np.argsort(z, axis=1, kind="stable") |
| z_quant = np.empty_like(z) |
| rows = np.arange(NS)[:, None] |
| z_quant[rows, order] = _Q2_IDEAL_QUANTILES[None, :] |
| z_analog = (1.0 - quantile_blend) * z_comp + quantile_blend * z_quant |
| else: |
| z_analog = z_comp |
|
|
| mu_a = np.sum(wb * z_analog, axis=1, keepdims=True) / wsum |
| z_analog = z_analog - mu_a |
| var_a = np.sum(wb * z_analog * z_analog, axis=1, keepdims=True) / wsum |
| scale = np.where(var_a > 1.0e-30, np.sqrt(var / np.maximum(var_a, 1.0e-30)), 0.0) |
| out = np.where(active[:, None], mu + z_analog * scale, xb) |
|
|
| if enforce_nonpositive_min: |
| orig_min = np.min(xb, axis=1) |
| new_min = np.min(out, axis=1) |
| shift = np.where((orig_min <= 0.0) & (new_min > 0.0), new_min, 0.0) |
| out = out - shift[:, None] |
|
|
| return out.reshape(QK_K).astype(np.float32) |
|
|
|
|
| def vesica_gram_analog( |
| x: np.ndarray, |
| w: np.ndarray | None = None, |
| tau: float = 1.9, |
| ) -> np.ndarray: |
| """Create a D6 / Vesica fold analog preserving the 2x2 half-block Gram matrix. |
| |
| Decomposes ``x`` into symmetric ``v = 0.5*(x[:128] + x[128:])`` and |
| antisymmetric ``u = 0.5*(x[:128] - x[128:])`` modes under pair weights |
| ``wp = 0.5*(w[:128] + w[128:])``. Whitens ``(v, u)``, applies smooth |
| hyperbolic saturation in the principal axes, and colors back with the exact |
| 2x2 covariance square-root so ``Gram(v_tilde, u_tilde) == Gram(v, u)``. |
| """ |
| x_f64 = np.asarray(x, dtype=np.float64).reshape(QK_K) |
| w_f64 = ( |
| np.ones(QK_K, dtype=np.float64) |
| if w is None |
| else np.asarray(w, dtype=np.float64).reshape(QK_K) |
| ) |
| half = QK_K // 2 |
| x0, x1 = x_f64[:half], x_f64[half:] |
| wp = np.maximum(0.5 * (w_f64[:half] + w_f64[half:]), 1.0e-30) |
| wsum = float(np.sum(wp)) |
|
|
| v = 0.5 * (x0 + x1) |
| u = 0.5 * (x0 - x1) |
| mv = float(np.sum(wp * v) / wsum) |
| mu = float(np.sum(wp * u) / wsum) |
| Y = np.vstack([v - mv, u - mu]) |
|
|
| G = (Y * wp[None, :]) @ Y.T / wsum |
| evals, evecs = np.linalg.eigh(G) |
| if float(np.max(evals)) < 1.0e-28: |
| return x_f64.astype(np.float32) |
|
|
| pos = evals > 1.0e-28 |
| inv_sqrt = np.where(pos, 1.0 / np.sqrt(np.maximum(evals, 1.0e-30)), 0.0) |
| W_white = (evecs * inv_sqrt[None, :]) @ evecs.T |
| Z = W_white @ Y |
|
|
| Z_comp = tau * np.tanh(Z / max(tau, 1.0e-6)) |
| Z_mean = np.sum(Z_comp * wp[None, :], axis=1, keepdims=True) / wsum |
| Z_comp = Z_comp - Z_mean |
|
|
| G_z = (Z_comp * wp[None, :]) @ Z_comp.T / wsum |
| ez, Vz = np.linalg.eigh(G_z) |
| pos_z = ez > 1.0e-28 |
| if not np.all(pos_z): |
| return x_f64.astype(np.float32) |
| inv_sqrt_z = np.where(pos_z, 1.0 / np.sqrt(np.maximum(ez, 1.0e-30)), 0.0) |
| Z_orth = (Vz * inv_sqrt_z[None, :]) @ Vz.T @ Z_comp |
|
|
| sqrt_G = (evecs * np.sqrt(np.maximum(evals, 0.0))[None, :]) @ evecs.T |
| Y_tilde = sqrt_G @ Z_orth |
| v_new = mv + Y_tilde[0] |
| u_new = mu + Y_tilde[1] |
| return np.concatenate([v_new + u_new, v_new - u_new]).astype(np.float32) |
|
|
|
|
| def subscale_lattice_analog( |
| x: np.ndarray, |
| w: np.ndarray | None = None, |
| gamma: float = 0.65, |
| ) -> np.ndarray: |
| """Create a cross-sub-block analog whose sub-block dispersions fit 4-bit scales. |
| |
| Preserves each sub-block's mean ``mu_j`` and unit direction |
| ``(x_j - mu_j) / sigma_j``, while tempering extreme sub-block standard |
| deviations ``sigma_j`` toward the superblock root-mean-square dispersion |
| and rescaling to preserve total centered weighted energy. |
| """ |
| xb = np.asarray(x, dtype=np.float64).reshape(NS, 16) |
| wb = ( |
| np.ones((NS, 16), dtype=np.float64) |
| if w is None |
| else np.asarray(w, dtype=np.float64).reshape(NS, 16) |
| ) |
| wsum = np.maximum(np.sum(wb, axis=1, keepdims=True), 1.0e-30) |
| mu = np.sum(wb * xb, axis=1, keepdims=True) / wsum |
| centered = xb - mu |
| energy_j = np.sum(wb * centered * centered, axis=1, keepdims=True) |
| total_energy = float(np.sum(energy_j)) |
| if total_energy <= 1.0e-30: |
| return xb.reshape(QK_K).astype(np.float32) |
|
|
| sigma = np.sqrt(np.maximum(energy_j / wsum, 1.0e-30)) |
| sigma_ref = float(np.sqrt(total_energy / max(float(np.sum(wsum)), 1.0e-30))) |
| target_sigma = np.power(np.maximum(sigma, 1.0e-30), gamma) * ( |
| sigma_ref ** (1.0 - gamma) |
| ) |
| scaled = np.where(sigma > 1.0e-15, centered * (target_sigma / sigma), 0.0) |
| new_energy = float(np.sum(wb * scaled * scaled)) |
| if new_energy > 1.0e-30: |
| scaled = scaled * np.sqrt(total_energy / new_energy) |
| out = mu + scaled |
| return out.reshape(QK_K).astype(np.float32) |
|
|
|
|
| def build_block_analogs( |
| x: np.ndarray, |
| w: np.ndarray | None = None, |
| ) -> list[np.ndarray]: |
| """Return the portfolio of geometry-preserving weight analogs for a block.""" |
| a1 = subblock_isomoment_analog( |
| x, w, tau=1.60, beta=0.25, quantile_blend=0.30, enforce_nonpositive_min=True |
| ) |
| a2 = subblock_isomoment_analog( |
| x, w, tau=2.15, beta=0.15, quantile_blend=0.15, enforce_nonpositive_min=True |
| ) |
| a3 = vesica_gram_analog(x, w, tau=1.85) |
| a4 = subscale_lattice_analog(a1, w, gamma=0.70) |
| return [a1, a2, a3, a4] |
|
|
