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?


1 Answer 1


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

    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
  • 1
    $\begingroup$ 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. $\endgroup$
    – Skinish
    Jan 19, 2018 at 19:03
  • 1
    $\begingroup$ 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? $\endgroup$
    – Skinish
    Jan 21, 2018 at 12:55
  • $\begingroup$ 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 $\endgroup$
    – Nix Searcy
    Jan 22, 2018 at 17:31
  • $\begingroup$ 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] $\endgroup$
    – Skinish
    Jan 23, 2018 at 11:33
  • $\begingroup$ @Skinish ahh, thanks. That makes it more clear. I’ll try to edit the answer today to reflect that $\endgroup$
    – Nix Searcy
    Jan 23, 2018 at 14:00

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.