feat(lib): add Unlimiformer-style kNN attention retrieval helpers

Implement the score-and-top-k pattern from Bertsch et al. (arXiv:2305.01625)
for embedding vectors: dot-product attention scores with top-k key selection,
plus multi-head and generic item ranking utilities. Add Bun tests and a test
script for @repo/lib.

packages/lib/package.json

@@ -3,6 +3,9 @@
 	"version": "0.0.0",
 	"private": true,
 	"type": "module",
+	"scripts": {
+		"test": "bun test"
+	},
 	"exports": {
 		"./*": "./*"
 	},

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

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<T> = {
+	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 = <T>(
+	queryEmbedding: number[],
+	items: readonly T[],
+	getEmbedding: (item: T) => number[] | null | undefined,
+	k: number,
+): RankedItem<T>[] => {
+	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,
+			},
+		]
+	})
+}