MaxSim

MaxSim is the late-interaction primitive used by systems such as ColBERT. Each query is a small matrix of token vectors. Each document is a small matrix of token vectors. The score between a query and a document is the sum of maximum cosine similarities between each query token and any document token. That is not a standard matrix multiply.

import NumKong

// 4 query tokens, each 16-dimensional
let queries = try Tensor<Float32>.full(rows: 4, columns: 16, value: 1.0)

// 8 document tokens, each 16-dimensional
let docs = try Tensor<Float32>.full(rows: 8, columns: 16, value: 1.0)

let queryPacked = try queries.maxSimPack()  // MaxSimPackedMatrix<Float32>
let docPacked   = try docs.maxSimPack()     // MaxSimPackedMatrix<Float32>

let score = try queryPacked.score(docPacked)  // Float64
assert(score.isFinite)

MaxSimPackedMatrix can also be constructed directly from a MatrixView:

let view = queries.view()
let packed = try MaxSimPackedMatrix<Float32>(packing: view)

Supported types and their output types:

Float16 support is unavailable on Intel macOS and Mac Catalyst, where Swift’s Float16 type does not exist; every other target, including Linux and Windows on x86-64, has it.