diff --git a/packages/lib/package.json b/packages/lib/package.json index 99b9d262..ee245167 100644 --- a/packages/lib/package.json +++ b/packages/lib/package.json @@ -3,6 +3,9 @@ "version": "0.0.0", "private": true, "type": "module", + "scripts": { + "test": "bun test" + }, "exports": { "./*": "./*" }, diff --git a/packages/lib/unlimiformer.test.ts b/packages/lib/unlimiformer.test.ts new file mode 100644 index 00000000..ae988a22 --- /dev/null +++ b/packages/lib/unlimiformer.test.ts @@ -0,0 +1,76 @@ +import { describe, expect, test } from "bun:test" +import { + attentionScores, + rankItemsByAttentionTopK, + topKAttentionKeys, + topKAttentionKeysMultiHead, +} from "./unlimiformer" + +const unit = (a: number, b: number, c: number) => { + const v = [a, b, c] + const norm = Math.hypot(a, b, c) + return v.map((x) => x / norm) +} + +describe("topKAttentionKeys", () => { + test("returns empty when k is zero or keys empty", () => { + const q = unit(1, 0, 0) + expect(topKAttentionKeys(q, [], 3)).toEqual([]) + expect(topKAttentionKeys(q, [unit(1, 0, 0)], 0)).toEqual([]) + }) + + test("orders by dot product / cosine on unit vectors", () => { + const q = unit(1, 0, 0) + const k0 = unit(1, 0, 0) + const k1 = unit(0, 1, 0) + const k2 = unit(-1, 0, 0) + const keys = [k1, k2, k0] + + const top = topKAttentionKeys(q, keys, 2) + expect(top.map((t) => t.index)).toEqual([2, 0]) + expect(top[0]?.score).toBeGreaterThan( + top[1]?.score ?? Number.NEGATIVE_INFINITY, + ) + }) + + test("caps k at number of keys", () => { + const q = unit(1, 0, 0) + const keys = [unit(1, 0, 0), unit(0, 1, 0)] + const top = topKAttentionKeys(q, keys, 10) + expect(top).toHaveLength(2) + }) +}) + +describe("attentionScores", () => { + test("matches per-key attentionScore", () => { + const q = unit(1, 1, 0) + const keys = [unit(1, 0, 0), unit(0, 1, 0)] + const scores = attentionScores(q, keys) + expect(scores).toHaveLength(2) + }) +}) + +describe("topKAttentionKeysMultiHead", () => { + test("runs independent top-k per query", () => { + const keys = [unit(1, 0, 0), unit(0, 1, 0), unit(0, 0, 1)] + const q0 = unit(1, 0, 0) + const q1 = unit(0, 1, 0) + const out = topKAttentionKeysMultiHead([q0, q1], keys, 1) + expect(out[0]?.[0]?.index).toBe(0) + expect(out[1]?.[0]?.index).toBe(1) + }) +}) + +describe("rankItemsByAttentionTopK", () => { + test("maps back to original indices and skips bad embeddings", () => { + const items = [ + { id: "a", e: unit(1, 0, 0) }, + { id: "b", e: null as number[] | null }, + { id: "c", e: unit(0, 1, 0) }, + ] + const q = unit(0, 1, 0) + const ranked = rankItemsByAttentionTopK(q, items, (x) => x.e, 2) + expect(ranked[0]?.item.id).toBe("c") + expect(ranked[0]?.originalIndex).toBe(2) + }) +}) diff --git a/packages/lib/unlimiformer.ts b/packages/lib/unlimiformer.ts new file mode 100644 index 00000000..a72a0c10 --- /dev/null +++ b/packages/lib/unlimiformer.ts @@ -0,0 +1,115 @@ +/** + * Unlimiformer-style kNN retrieval (Bertsch et al., 2023). + * + * @see https://arxiv.org/abs/2305.01625 — cross-attention is approximated by retrieving + * the top-k keys under dot-product scores. In this codebase, embeddings are treated as + * key/query vectors; for L2-normalized vectors, dot product equals cosine similarity, + * matching the ranking used elsewhere in `@repo/lib/similarity`. + */ + +import { cosineSimilarity } from "./similarity" + +export type AttentionTopK = { + index: number + score: number +} + +/** + * Dot-product attention score between one query vector and one key vector. + * For normalized embeddings this matches cosine similarity. + */ +export const attentionScore = (query: number[], key: number[]): number => + cosineSimilarity(query, key) + +/** + * Attention scores for `query` against every row in `keys` (same dimension as query). + */ +export const attentionScores = (query: number[], keys: number[][]): number[] => + keys.map((key) => attentionScore(query, key)) + +/** + * Retrieve the top-k keys by attention score (Unlimiformer's kNN over key index). + * Results are sorted by descending score. + */ +export const topKAttentionKeys = ( + query: number[], + keys: number[][], + k: number, +): AttentionTopK[] => { + if (k <= 0 || keys.length === 0) { + return [] + } + + const effectiveK = Math.min(k, keys.length) + const scored: AttentionTopK[] = keys.map((key, index) => ({ + index, + score: attentionScore(query, key), + })) + + scored.sort((a, b) => b.score - a.score) + return scored.slice(0, effectiveK) +} + +/** + * Per-head top-k retrieval: each query vector gets its own top-k over the same key set, + * analogous to multi-head cross-attention with separate query projections. + */ +export const topKAttentionKeysMultiHead = ( + queries: number[][], + keys: number[][], + k: number, +): AttentionTopK[][] => queries.map((q) => topKAttentionKeys(q, keys, k)) + +export type RankedItem = { + item: T + originalIndex: number + score: number +} + +/** + * Rank arbitrary items that carry embeddings, returning the top-k by attention score. + * Items with missing or empty embeddings are skipped. + */ +export const rankItemsByAttentionTopK = ( + queryEmbedding: number[], + items: readonly T[], + getEmbedding: (item: T) => number[] | null | undefined, + k: number, +): RankedItem[] => { + if (k <= 0 || items.length === 0) { + return [] + } + + const packed: Array<{ item: T; originalIndex: number; embedding: number[] }> = + [] + + for (let i = 0; i < items.length; i++) { + const item = items[i] + if (item === undefined) continue + const embedding = getEmbedding(item) + if (embedding && embedding.length > 0) { + packed.push({ item, originalIndex: i, embedding }) + } + } + + if (packed.length === 0) { + return [] + } + + const keys = packed.map((p) => p.embedding) + const top = topKAttentionKeys(queryEmbedding, keys, k) + + return top.flatMap(({ index: keyIndex, score }) => { + const row = packed[keyIndex] + if (!row) { + return [] + } + return [ + { + item: row.item, + originalIndex: row.originalIndex, + score, + }, + ] + }) +}