File size: 3,345 Bytes
fc93158 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 | import { OMEGA_MEMORY_EMBEDDING_DIMENSIONS } from "./policy.js";
export type QuantizedEmbedding = {
scheme: "turboquant-lite-v1";
dimensions: number;
bits: number;
scale: number;
values: number[];
};
const DEFAULT_BITS = 6;
const MIN_BITS = 2;
const MAX_BITS = 8;
function isPowerOfTwo(value: number): boolean {
return value > 0 && (value & (value - 1)) === 0;
}
function normalize(vector: number[]): number[] {
const magnitude = Math.sqrt(vector.reduce((sum, value) => sum + value * value, 0));
if (magnitude === 0) {
return vector.map(() => 0);
}
return vector.map((value) => value / magnitude);
}
function resolveBitWidth(bits?: number): number {
const candidate =
typeof bits === "number" && Number.isFinite(bits) ? Math.round(bits) : DEFAULT_BITS;
return Math.max(MIN_BITS, Math.min(MAX_BITS, candidate));
}
function signMask(index: number): number {
let x = (index + 1) * 0x45d9f3b;
x ^= x >>> 16;
return (x & 1) === 0 ? 1 : -1;
}
function applySignScramble(vector: number[]): number[] {
return vector.map((value, index) => value * signMask(index));
}
function fastWalshHadamard(vector: number[]): number[] {
const out = [...vector];
for (let len = 1; len < out.length; len <<= 1) {
for (let start = 0; start < out.length; start += len << 1) {
for (let offset = 0; offset < len; offset += 1) {
const left = out[start + offset] ?? 0;
const right = out[start + len + offset] ?? 0;
out[start + offset] = left + right;
out[start + len + offset] = left - right;
}
}
}
const norm = Math.sqrt(out.length);
return out.map((value) => value / norm);
}
function clamp(value: number, min: number, max: number): number {
return Math.max(min, Math.min(max, value));
}
function quantizeScalar(value: number, bits: number, scale: number): number {
const levels = (1 << bits) - 1;
const normalized = clamp((value + scale) / (2 * scale), 0, 1);
return Math.round(normalized * levels);
}
function dequantizeScalar(value: number, bits: number, scale: number): number {
const levels = (1 << bits) - 1;
const normalized = clamp(value / levels, 0, 1);
return normalized * 2 * scale - scale;
}
export function quantizeNormalizedEmbedding(params: {
embedding: number[];
bits?: number;
}): QuantizedEmbedding | null {
const embedding = normalize(params.embedding);
if (
embedding.length === 0 ||
embedding.length !== OMEGA_MEMORY_EMBEDDING_DIMENSIONS ||
!isPowerOfTwo(embedding.length)
) {
return null;
}
const bits = resolveBitWidth(params.bits);
const rotated = fastWalshHadamard(applySignScramble(embedding));
const scale = Math.max(1e-6, ...rotated.map((value) => Math.abs(value)));
return {
scheme: "turboquant-lite-v1",
dimensions: embedding.length,
bits,
scale,
values: rotated.map((value) => quantizeScalar(value, bits, scale)),
};
}
export function dequantizeEmbedding(quantized: QuantizedEmbedding): number[] {
if (
quantized.scheme !== "turboquant-lite-v1" ||
quantized.dimensions <= 0 ||
!isPowerOfTwo(quantized.dimensions)
) {
return [];
}
const rotated = quantized.values.map((value) =>
dequantizeScalar(value, resolveBitWidth(quantized.bits), quantized.scale),
);
return normalize(applySignScramble(fastWalshHadamard(rotated)));
}
|