| 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))); |
| } |
|
|
