soulcraft/brainy
2
0
Fork 0
Code Issues 1 Pull requests Projects Releases 335 Packages Wiki Activity Actions

feat: add SQ8 vector quantization, lazy loading, and two-phase rerank to HNSW

Browse source 
- SQ8 scalar quantization (8-bit) for 4x vector storage reduction
- Lazy vector loading: evict float32 vectors after graph construction,
  load on-demand from storage via UnifiedCache
- Two-phase search: over-retrieve with SQ8 approximate distances,
  rerank top candidates with exact float32 distances
- Configuration surface: hnsw.quantization and hnsw.vectorStorage in BrainyConfig
- All features disabled by default (zero behavior change for existing users)
- 27 new tests covering quantization accuracy, lazy loading, reranking
- Remove GitHub Actions CI (build locally, cortex CI handles native builds)
This commit is contained in:
David Snelling 2026-01-31 12:41:53 -08:00
parent e384afcdac
commit 0f3a88429d
9 changed files with 1138 additions and 222 deletions
Download patch file Download diff file Expand all files Collapse all files

11
src/brainy.ts
View file

@ -6236,7 +6236,14 @@ export class Brainy<T = any> implements BrainyInterface<T> {
private setupIndex(): HNSWIndex | TypeAwareHNSWIndex {
const indexConfig = {
...this.config.index,
distanceFunction: this.distance
distanceFunction: this.distance,
// Wire HNSW optimization config (v7.11.0)
quantization: this.config.hnsw?.quantization ? {
enabled: this.config.hnsw.quantization.enabled ?? false,
bits: this.config.hnsw.quantization.bits ?? 8,
rerankMultiplier: this.config.hnsw.quantization.rerankMultiplier ?? 3
} : undefined,
vectorStorage: this.config.hnsw?.vectorStorage
}
const persistMode = this.resolveHNSWPersistMode()
@ -6368,6 +6375,8 @@ export class Brainy<T = any> implements BrainyInterface<T> {
reservedQueryMemory: config?.reservedQueryMemory ?? undefined as any,
// HNSW persistence mode - undefined = smart default in setupIndex
hnswPersistMode: config?.hnswPersistMode ?? undefined as any,
// HNSW optimization options (v7.11.0)
hnsw: config?.hnsw ?? undefined as any,
// Embedding initialization - false = lazy init on first embed()
eagerEmbeddings: config?.eagerEmbeddings ?? false,
// Integration Hub - undefined/false = disabled

12
src/coreTypes.ts
View file

@ -96,6 +96,10 @@ export interface HNSWNoun {
vector: Vector
connections: Map<number, Set<string>> // level -> set of connected noun ids
level: number // The highest layer this noun appears in
// SQ8 quantized vector for approximate distance computation (B1 optimization)
quantizedVector?: Uint8Array // 4x smaller than float32 vector
codebookMin?: number // quantization codebook minimum
codebookMax?: number // quantization codebook maximum
// ✅ NO metadata field - stored separately for optimization
}
@ -309,6 +313,14 @@ export interface HNSWConfig {
ml: number // Maximum level
useDiskBasedIndex?: boolean // Whether to use disk-based index
maxConcurrentNeighborWrites?: number // Maximum concurrent neighbor updates during insert. Default: unlimited (full concurrency)
// SQ8 vector quantization (4x memory reduction, ~0.4% accuracy loss)
quantization?: {
enabled: boolean // default: false — preserves current behavior exactly
bits?: 8 | 4 // default: 8 (SQ8). SQ4 requires brainy-cortex native.
rerankMultiplier?: number // default: 3 — over-retrieve 3x, rerank with float32
}
// Vector storage mode
vectorStorage?: 'memory' | 'lazy' // default: 'memory' — 'lazy' evicts vectors after insert
}
/**

137
src/hnsw/hnswIndex.ts
View file

@ -15,6 +15,8 @@ import { executeInThread } from '../utils/workerUtils.js'
import type { BaseStorage } from '../storage/baseStorage.js'
import { getGlobalCache, UnifiedCache } from '../utils/unifiedCache.js'
import { prodLog } from '../utils/logger.js'
import { quantizeSQ8, distanceSQ8 } from '../utils/vectorQuantization.js'
import type { SQ8QuantizedVector } from '../utils/vectorQuantization.js'
// Default HNSW parameters
const DEFAULT_CONFIG: HNSWConfig = {
@ -53,6 +55,12 @@ export class HNSWIndex {
private dirtyNodes: Set<string> = new Set() // Nodes with unpersisted HNSW data
private dirtySystem: boolean = false // Whether system data (entryPoint, maxLevel) needs persist
// SQ8 quantization support (B1 optimization)
private quantizationEnabled: boolean = false
private rerankMultiplier: number = 3
// Lazy vector storage (B2 optimization)
private vectorStorageMode: 'memory' | 'lazy' = 'memory'
constructor(
config: Partial<HNSWConfig> = {},
distanceFunction: DistanceFunction = euclideanDistance,
@ -67,6 +75,15 @@ export class HNSWIndex {
this.storage = options.storage || null
this.persistMode = options.persistMode || 'immediate'
// SQ8 quantization config (default: disabled, preserves current behavior)
if (config.quantization?.enabled) {
this.quantizationEnabled = true
this.rerankMultiplier = config.quantization.rerankMultiplier ?? 3
}
// Vector storage mode (default: 'memory', preserves current behavior)
this.vectorStorageMode = config.vectorStorage || 'memory'
// Use SAME UnifiedCache as Graph and Metadata for fair memory competition
this.unifiedCache = getGlobalCache()
}
@ -257,7 +274,11 @@ export class HNSWIndex {
id: original.id,
vector: [...original.vector], // Deep copy vector array
connections: connectionsCopy,
level: original.level
level: original.level,
// Copy SQ8 quantized data if present
quantizedVector: original.quantizedVector ? new Uint8Array(original.quantizedVector) : undefined,
codebookMin: original.codebookMin,
codebookMax: original.codebookMax
}
this.nouns.set(nodeId, nodeCopy)
@ -343,7 +364,7 @@ export class HNSWIndex {
// Generate random level for this noun
const nounLevel = this.getRandomLevel()
// Create new noun
// Create new noun with optional SQ8 quantization
const noun: HNSWNoun = {
id,
vector,
@ -351,6 +372,14 @@ export class HNSWIndex {
level: nounLevel
}
// Quantize vector if enabled (B1: 4x storage reduction)
if (this.quantizationEnabled) {
const sq8 = quantizeSQ8(vector)
noun.quantizedVector = sq8.quantized
noun.codebookMin = sq8.min
noun.codebookMax = sq8.max
}
// Initialize empty connection sets for each level
for (let level = 0; level <= nounLevel; level++) {
noun.connections.set(level, new Set<string>())
@ -574,6 +603,13 @@ export class HNSWIndex {
this.highLevelNodes.get(nounLevel)!.add(id)
}
// Lazy vector eviction (B2: graph-only memory after insert)
// After graph construction completes, evict the full vector from memory.
// Future searches will load vectors on-demand via getVectorSafe() + UnifiedCache.
if (this.vectorStorageMode === 'lazy' && this.storage) {
noun.vector = [] // Release float32 vector from memory
}
// Persist HNSW graph data to storage
// Respect persistMode setting
if (this.storage && this.persistMode === 'immediate') {
@ -631,11 +667,21 @@ export class HNSWIndex {
/**
* Search for nearest neighbors
*
* When SQ8 quantization is enabled and reranking is active:
* - Phase 1: Over-retrieve k*rerankMultiplier candidates using SQ8 approximate distances
* - Phase 2: Load full float32 vectors for candidates, compute exact distances, return top k
*
* @param queryVector Query vector
* @param k Number of results to return
* @param filter Optional filter function
* @param options Additional search options
*/
public async search(
queryVector: Vector,
k: number = 10,
filter?: (id: string) => Promise<boolean>
filter?: (id: string) => Promise<boolean>,
options?: { rerank?: { multiplier: number } }
): Promise<Array<[string, number]>> {
if (this.nouns.size === 0) {
return []
@ -687,6 +733,11 @@ export class HNSWIndex {
let currObj = entryPoint
// SQ8: Pre-quantize query vector for fast approximate distances during traversal
if (this.quantizationEnabled) {
this._querySQ8 = quantizeSQ8(queryVector)
}
// OPTIMIZATION: Preload entry point vector
await this.preloadVectors([entryPoint.id])
@ -754,9 +805,14 @@ export class HNSWIndex {
}
}
// Search at level 0 with ef = k
// Determine effective rerank multiplier
const rerankActive = this.quantizationEnabled && (options?.rerank || this.rerankMultiplier > 1)
const multiplier = options?.rerank?.multiplier ?? this.rerankMultiplier
const effectiveK = rerankActive ? k * multiplier : k
// Search at level 0 with ef = effectiveK
// If we have a filter, increase ef to compensate for filtered results
const ef = filter ? Math.max(this.config.efSearch * 3, k * 3) : Math.max(this.config.efSearch, k)
const ef = filter ? Math.max(this.config.efSearch * 3, effectiveK * 3) : Math.max(this.config.efSearch, effectiveK)
const nearestNouns = await this.searchLayer(
queryVector,
currObj,
@ -765,6 +821,35 @@ export class HNSWIndex {
filter
)
// Phase 2: Rerank with exact float32 distances if quantization is active (B3)
if (rerankActive && nearestNouns.size > 0) {
// Clear SQ8 cache before reranking (we need exact distances now)
this._querySQ8 = null
const candidates = [...nearestNouns].slice(0, effectiveK)
// Load full float32 vectors for the candidate set
const candidateIds = candidates.map(([id]) => id)
await this.preloadVectors(candidateIds)
// Recompute exact distances
const reranked: Array<[string, number]> = []
for (const [id] of candidates) {
const noun = this.nouns.get(id)
if (!noun) continue
const exactVector = await this.getVectorSafe(noun)
const exactDist = this.distanceFunction(queryVector, exactVector)
reranked.push([id, exactDist])
}
// Sort by exact distance and return top k
reranked.sort((a, b) => a[1] - b[1])
return reranked.slice(0, k)
}
// Clear SQ8 cache
this._querySQ8 = null
// Convert to array and sort by distance
return [...nearestNouns].slice(0, k)
}
@ -1074,6 +1159,19 @@ export class HNSWIndex {
* @returns number | Promise<number> - sync when cached, async when needs load
*/
private distanceSafe(queryVector: Vector, noun: HNSWNoun): number | Promise<number> {
// SQ8 fast path: use quantized distance when available (B1 optimization)
// This avoids loading full float32 vectors during graph traversal
if (this.quantizationEnabled &&
noun.quantizedVector &&
noun.codebookMin !== undefined &&
noun.codebookMax !== undefined &&
this._querySQ8) {
return distanceSQ8(
this._querySQ8.quantized, this._querySQ8.min, this._querySQ8.max,
noun.quantizedVector, noun.codebookMin, noun.codebookMax
)
}
// Try sync fast path
const nounVector = this.getVectorSync(noun)
@ -1088,6 +1186,9 @@ export class HNSWIndex {
)
}
// Cached SQ8 quantization of the current query vector for distanceSafe fast path
private _querySQ8: SQ8QuantizedVector | null = null
/**
* Get all nodes at a specific level for clustering
* This enables O(n) clustering using HNSW's natural hierarchy
@ -1207,14 +1308,25 @@ export class HNSWIndex {
continue
}
// Determine if vector should be kept in memory
const keepVector = shouldPreload && this.vectorStorageMode !== 'lazy'
// Create noun object with restored connections
const noun: HNSWNoun = {
id: nounData.id,
vector: shouldPreload ? nounData.vector : [], // Preload if dataset is small
vector: keepVector ? nounData.vector : [], // Preload if dataset is small and not lazy
connections: new Map(),
level: hnswData.level
}
// Restore SQ8 quantized data if quantization is enabled
if (this.quantizationEnabled && nounData.vector.length > 0) {
const sq8 = quantizeSQ8(nounData.vector)
noun.quantizedVector = sq8.quantized
noun.codebookMin = sq8.min
noun.codebookMax = sq8.max
}
// Restore connections from persisted data
for (const [levelStr, nounIds] of Object.entries(hnswData.connections)) {
const level = parseInt(levelStr, 10)
@ -1281,14 +1393,25 @@ export class HNSWIndex {
continue
}
// Determine if vector should be kept in memory
const keepVector = shouldPreload && this.vectorStorageMode !== 'lazy'
// Create noun object with restored connections
const noun: HNSWNoun = {
id: nounData.id,
vector: shouldPreload ? nounData.vector : [], // Preload if dataset is small
vector: keepVector ? nounData.vector : [], // Preload if dataset is small and not lazy
connections: new Map(),
level: hnswData.level
}
// Restore SQ8 quantized data if quantization is enabled
if (this.quantizationEnabled && nounData.vector.length > 0) {
const sq8 = quantizeSQ8(nounData.vector)
noun.quantizedVector = sq8.quantized
noun.codebookMin = sq8.min
noun.codebookMax = sq8.max
}
// Restore connections from persisted data
for (const [levelStr, nounIds] of Object.entries(hnswData.connections)) {
const level = parseInt(levelStr, 10)

19
src/hnsw/typeAwareHNSWIndex.ts
View file

@ -256,21 +256,22 @@ export class TypeAwareHNSWIndex {
queryVector: Vector,
k: number = 10,
type?: NounType | NounType[],
filter?: (id: string) => Promise<boolean>
filter?: (id: string) => Promise<boolean>,
options?: { rerank?: { multiplier: number } }
): Promise<Array<[string, number]>> {
// Single-type search (fast path)
if (type && typeof type === 'string') {
const index = this.getIndexForType(type)
return await index.search(queryVector, k, filter)
return await index.search(queryVector, k, filter, options)
}
// Multi-type search (handle empty array edge case)
if (type && Array.isArray(type) && type.length > 0) {
return await this.searchMultipleTypes(queryVector, k, type, filter)
return await this.searchMultipleTypes(queryVector, k, type, filter, options)
}
// All-types search (slowest path + empty array fallback)
return await this.searchAllTypes(queryVector, k, filter)
return await this.searchAllTypes(queryVector, k, filter, options)
}
/**
@ -286,7 +287,8 @@ export class TypeAwareHNSWIndex {
queryVector: Vector,
k: number,
types: NounType[],
filter?: (id: string) => Promise<boolean>
filter?: (id: string) => Promise<boolean>,
options?: { rerank?: { multiplier: number } }
): Promise<Array<[string, number]>> {
const allResults: Array<[string, number]> = []
@ -294,7 +296,7 @@ export class TypeAwareHNSWIndex {
for (const type of types) {
if (this.indexes.has(type)) {
const index = this.indexes.get(type)!
const results = await index.search(queryVector, k, filter)
const results = await index.search(queryVector, k, filter, options)
allResults.push(...results)
}
}
@ -320,13 +322,14 @@ export class TypeAwareHNSWIndex {
private async searchAllTypes(
queryVector: Vector,
k: number,
filter?: (id: string) => Promise<boolean>
filter?: (id: string) => Promise<boolean>,
options?: { rerank?: { multiplier: number } }
): Promise<Array<[string, number]>> {
const allResults: Array<[string, number]> = []
// Search each type's graph
for (const [type, index] of this.indexes.entries()) {
const results = await index.search(queryVector, k, filter)
const results = await index.search(queryVector, k, filter, options)
allResults.push(...results)
}

10
src/types/brainy.types.ts
View file

@ -728,6 +728,16 @@ export interface BrainyConfig {
// Cloud storage (GCS/S3/R2/Azure) should use 'deferred' for 30-50× faster adds
hnswPersistMode?: 'immediate' | 'deferred'
// HNSW optimization options (v7.11.0)
hnsw?: {
quantization?: {
enabled?: boolean // default: false — current behavior exactly
bits?: 8 | 4 // default: 8 (SQ8). SQ4 requires brainy-cortex native.
rerankMultiplier?: number // default: 3 — over-retrieve 3x, rerank with float32
}
vectorStorage?: 'memory' | 'lazy' // default: 'memory' — 'lazy' evicts vectors after insert
}
// Memory management options
maxQueryLimit?: number // Override auto-detected query result limit (max: 100000)
reservedQueryMemory?: number // Memory reserved for queries in bytes (e.g., 1073741824 = 1GB)

218
src/utils/vectorQuantization.ts Normal file
View file

@ -0,0 +1,218 @@
/**
* Vector Quantization Utilities
*
* Standalone SQ8 (Scalar Quantization, 8-bit) functions for HNSW vector compression.
* Each vector is independently quantized using its own min/max codebook.
*
* Storage format per vector:
* - quantized: Uint8Array (dimension bytes, 4x smaller than float32)
* - min: number (codebook minimum)
* - max: number (codebook maximum)
*
* Accuracy: SQ8 introduces ~0.4% error per dimension on normalized vectors.
* Combined with reranking (3x over-retrieval + float32 rerank), recall@100
* is expected to remain >99% for typical workloads.
*/
import type { Vector } from '../coreTypes.js'
/**
* Result of SQ8 quantization
*/
export interface SQ8QuantizedVector {
quantized: Uint8Array
min: number
max: number
}
/**
* Configuration for HNSW quantization
*/
export interface HNSWQuantizationConfig {
enabled: boolean
bits: 8 | 4
rerankMultiplier: number
}
/**
* Quantize a float32 vector to SQ8 (8-bit unsigned integers).
*
* Maps [min, max] of the vector to [0, 255].
* For zero-range vectors (all values equal), all quantized values are 128.
*
* @param vector - Input float32 vector
* @returns Quantized vector with codebook (min/max)
*/
export function quantizeSQ8(vector: Vector): SQ8QuantizedVector {
const len = vector.length
let min = vector[0]
let max = vector[0]
for (let i = 1; i < len; i++) {
const v = vector[i]
if (v < min) min = v
if (v > max) max = v
}
const range = max - min
const quantized = new Uint8Array(len)
if (range === 0) {
// All values are identical — map to midpoint
quantized.fill(128)
} else {
const scale = 255 / range
for (let i = 0; i < len; i++) {
quantized[i] = Math.round((vector[i] - min) * scale)
}
}
return { quantized, min, max }
}
/**
* Dequantize an SQ8 vector back to float32.
*
* @param quantized - Uint8Array of quantized values
* @param min - Codebook minimum
* @param max - Codebook maximum
* @returns Reconstructed float32 vector
*/
export function dequantizeSQ8(quantized: Uint8Array, min: number, max: number): Vector {
const len = quantized.length
const result = new Array<number>(len)
const range = max - min
if (range === 0) {
result.fill(min)
} else {
const scale = range / 255
for (let i = 0; i < len; i++) {
result[i] = min + quantized[i] * scale
}
}
return result
}
/**
* Compute approximate cosine distance between two SQ8-quantized vectors.
*
* Operates directly on uint8 arrays without full dequantization.
* Uses integer arithmetic where possible for speed.
*
* Cosine distance = 1 - (A.B / (|A| * |B|))
*
* For quantized values q_i with codebook (min, max):
* actual_i = min + q_i * (max - min) / 255
*
* The dot product and norms can be computed on quantized values
* and scaled by the codebook parameters.
*
* @param a - First quantized vector
* @param aMin - First vector codebook minimum
* @param aMax - First vector codebook maximum
* @param b - Second quantized vector
* @param bMin - Second vector codebook minimum
* @param bMax - Second vector codebook maximum
* @returns Approximate cosine distance [0, 2]
*/
export function distanceSQ8(
a: Uint8Array,
aMin: number,
aMax: number,
b: Uint8Array,
bMin: number,
bMax: number
): number {
const len = a.length
// Compute raw integer sums for efficiency
// actual_a[i] = aMin + a[i] * aScale, where aScale = (aMax - aMin) / 255
// actual_b[i] = bMin + b[i] * bScale, where bScale = (bMax - bMin) / 255
//
// dot = sum(actual_a[i] * actual_b[i])
// = sum((aMin + a[i]*aScale) * (bMin + b[i]*bScale))
// = n*aMin*bMin + aMin*bScale*sum(b[i]) + bMin*aScale*sum(a[i]) + aScale*bScale*sum(a[i]*b[i])
//
// normA^2 = sum(actual_a[i]^2)
// = n*aMin^2 + 2*aMin*aScale*sum(a[i]) + aScale^2*sum(a[i]^2)
const aScale = (aMax - aMin) / 255
const bScale = (bMax - bMin) / 255
let sumA = 0
let sumB = 0
let sumAB = 0
let sumAA = 0
let sumBB = 0
for (let i = 0; i < len; i++) {
const ai = a[i]
const bi = b[i]
sumA += ai
sumB += bi
sumAB += ai * bi
sumAA += ai * ai
sumBB += bi * bi
}
const dot =
len * aMin * bMin +
aMin * bScale * sumB +
bMin * aScale * sumA +
aScale * bScale * sumAB
const normASq =
len * aMin * aMin +
2 * aMin * aScale * sumA +
aScale * aScale * sumAA
const normBSq =
len * bMin * bMin +
2 * bMin * bScale * sumB +
bScale * bScale * sumBB
const normA = Math.sqrt(normASq)
const normB = Math.sqrt(normBSq)
if (normA === 0 || normB === 0) {
return 1.0 // Maximum distance for zero vectors
}
const cosine = dot / (normA * normB)
// Clamp to [-1, 1] to handle floating point imprecision
return 1 - Math.max(-1, Math.min(1, cosine))
}
/**
* Serialize an SQ8 quantized vector to a compact binary format.
*
* Format: [min:float32][max:float32][quantized:uint8[]]
* Total size: 8 + dimension bytes
*
* @param sq8 - Quantized vector
* @returns ArrayBuffer with serialized data
*/
export function serializeSQ8(sq8: SQ8QuantizedVector): ArrayBuffer {
const buffer = new ArrayBuffer(8 + sq8.quantized.length)
const view = new DataView(buffer)
view.setFloat32(0, sq8.min, true) // little-endian
view.setFloat32(4, sq8.max, true)
new Uint8Array(buffer, 8).set(sq8.quantized)
return buffer
}
/**
* Deserialize an SQ8 quantized vector from binary format.
*
* @param buffer - ArrayBuffer with serialized data
* @returns Deserialized SQ8 quantized vector
*/
export function deserializeSQ8(buffer: ArrayBuffer): SQ8QuantizedVector {
const view = new DataView(buffer)
const min = view.getFloat32(0, true)
const max = view.getFloat32(4, true)
const quantized = new Uint8Array(buffer, 8)
return { quantized, min, max }
}