# Keras - Computing cosine similarity matrix of two 3D tensors

Using TF backend, I need to construct a similarity matrices of two 3D vectors, both with shape (batch_size, N, M), being N and M natural numbers.

The function tf.losses.cosine_distance is only between 1D tensors. I need to build a Tensor matrix batch_sizexNxM such that matrix[k][i][j] will be the cosine similarity of the Tensor1[k][i] and Tensor2[k][j].

How can I do this?

I know of no pairwise distance operations in Keras or tensorflow. But the matrix math can be implemented in TF/Keras backend code and then placed in a Keras layer.

Here's the matrix representation of the cosine similarity of two vectors: $$cos(\theta) = \frac{\mathbf{A}\cdot\mathbf{B}}{\|\mathbf{A}\|_2 \|\mathbf{B}\|_2}$$

I'll show the code and a test that confirms that it works. First, generate non-trivial test data.

import numpy as np
import keras
import keras.backend as K

# set up test data
n_batch = 100
n = 400 # number of points in the first set
m = 500 # number of points in the second set
d = 200 # number of dimensions

A = np.random.rand(n_batch, n, d)
B = np.random.rand(n_batch, m, d)


Define pairwise cosine similarity function.

# convenience l2_norm function
def l2_norm(x, axis=None):
"""
takes an input tensor and returns the l2 norm along specified axis
"""

square_sum = K.sum(K.square(x), axis=axis, keepdims=True)
norm = K.sqrt(K.maximum(square_sum, K.epsilon()))

return norm

def pairwise_cosine_sim(A_B):
"""
A [batch x n x d] tensor of n rows with d dimensions
B [batch x m x d] tensor of n rows with d dimensions

returns:
D [batch x n x m] tensor of cosine similarity scores between each point i<n, j<m
"""

A, B = A_B
A_mag = l2_norm(A, axis=2)
B_mag = l2_norm(B, axis=2)
num = K.batch_dot(A_tensor, K.permute_dimensions(B_tensor, (0,2,1)))
den = (A_mag * K.permute_dimensions(B_mag, (0,2,1)))
dist_mat =  num / den

return dist_mat


Build a shallow Keras model around the function.

# build dummy model
A_tensor = K.constant(A)
B_tensor = K.constant(B)
A_input = keras.Input(tensor=A_tensor)
B_input = keras.Input(tensor=B_tensor)
dist_output = keras.layers.Lambda(pairwise_cosine_sim)([A_input, B_input])
dist_model = keras.Model(inputs=[A_input, B_input], outputs=dist_output)
dist_model.compile("sgd", "mse")


Compare to sklearn implementation

sk_dist = np.zeros( (n_batch, n, m) )
for i in range(n_batch):
sk_dist[i,...] = cosine_similarity(A[i,...], B[i,...])

keras_dist = dist_model.predict(None, steps=1)
np.allclose(sk_dist, keras_dist)
> True

• Thank you very much! There is one little problem though. Lambda don't accept two arguments. You could solve this by making your pairwise_cosine receive the arguments in a list instead of separated. However there is another issue. I need this layer to accept 3D Tensors actually, where the 1st dimension is the batch size. If they were 2D Tensors I could indeed solve it as you said. But in this case I would need something like Lambda(lambda x, y: pairwise_cosine(x, y)). Sorry for not mentioning this previously, was trying to solve one problem at a time. Commented Jan 19, 2018 at 19:03
• Btw, you dont use A_mag and B_mag afterwards. Can you refactor your answer so that A and B are of the form (?, n, m), where ? is the batch size? Commented Jan 21, 2018 at 12:55
• Thanks for the input @Skinish ! I've cleaned up the code and added a test. It's unclear to me what the batch size would be in this case so I refactored it using the number of points in B and A (not necessarily equal) as batch sizes. Hopefully that works Commented Jan 22, 2018 at 17:31
• Thanks a lot, that definitely answers my question! Moreover, are you able to refactor your pairwise_cosine_sim function so that it receive [? x n x d] and [? x m x d] Commented Jan 23, 2018 at 11:33
• @Skinish ahh, thanks. That makes it more clear. I’ll try to edit the answer today to reflect that Commented Jan 23, 2018 at 14:00