Feature: Add high level OP for MaxSim calculation #272
Description
Describe what you are looking for
MaxSim was introduced in the Colbert paper and is used by almost all Late Interaction models like Colbert, Colpali etc.
Unfortunately there is not native or fast implementation (at least I am not aware of one) to compute MaxSim fast over batch of documents.
I already tried to optimize it a bit. Here my approach with using Numpy:
def max_sim(q,d):
# flatten then compute; may advantage due to better memory continuous
K,M,D = d.shape
Q = q.shape[0]
scores = d.reshape(-1,D) @ q.T # (K*M, Q)
max_scores = np.max(scores.reshape(K,M,Q), axis=1) # (K, Q)
return np.sum(max_scores, axis=0) # (Q,)Shape q: (#query_tokens, dim)
Shape d: (batch_size, #document_tokens, dim).
The benchmark I ran where with the following configs on a Dedicated Hetzner Cloud machine (CCX23):
q: (32, 128)
d: (1000, 666, 128)
Avg. over 10 runs: 86ms
Would love to see if we can do better!
Can you contribute to the implementation?
- I can contribute
Is your feature request specific to a certain interface?
It applies to everything
Contact Details
No response
Is there an existing issue for this?
- I have searched the existing issues
Code of Conduct
- I agree to follow this project's Code of Conduct