Can anyone help me implement similarity/cosine similarity in Nx?

pejrich · February 10, 2023, 2:33am

I’m trying to port some code from python, to elixir. The python code generates tensors from embedding in BERT, then does some form of similarity comparison between them. From what i’ve found online, it looks like cosine similarity is the calculation I’m looking for, but I can’t quite understand it enough to implement it in Nx.

The formula is listed as A ⋅ B / ||A|| ||B||. I have two tensors with with the shape #Nx.Tensor<f32[1][6][119547]...>. So far this is all i’ve come up with:

for i <- 0..5, j <- 0..5 do
  t1 = tensor1[0][i]
  t2 = tensor2[0][j]
  Nx.dot(t1, t2) / ???
end

I found another formula on the wikipedia page that’s numpy code, which says:

np.sum(a*b)/(np.sqrt(np.sum(a**2)) * np.sqrt(np.sum(b**2)))

I think converted to Nx that’s:

defmodule CosSim do
  import Nx.Defn
  defn cosine_similarity(a, b) do
    left = Nx.sqrt(Nx.sum(a**2))
    right = Nx.sqrt(Nx.sum(b**2))
    Nx.sum(a * b) / (left * right)
  end
end

If I do a quick test:

a = Nx.tensor([1,2,3])
b = Nx.tensor([4,5,6])
CosSim.cosine_similarity(a, b)

#Nx.Tensor<
  f32
  EXLA.Backend<host:0, 0.528063503.4042653716.111775>
  0.9746317863464355
>

If I try to validate it in python:

>>> a = np.matrix([1,2,3])
>>> b = np.matrix([4,5,6])
>>> np.sum(a*b)/(np.sqrt(np.sum(a**2)) * np.sqrt(np.sum(b**2)))

ValueError: shapes (1,3) and (1,3) not aligned: 3 (dim 1) != 1 (dim 0)

So something is off. Admittdly i’ve very new to any of this ML/Nx stuff, so maybe I’m way off, or maybe i’m close. Any tips?

polvalente · February 10, 2023, 2:46am

We have Cosine Distance available in Scholar: scholar/distance.ex at main · elixir-nx/scholar · GitHub

Cosine Similarity would be 1 - CosDistance

pejrich · February 10, 2023, 2:50am

Awesome, thanks. I’ll give that a go. I swear I spent 30 minutes googling anything related to the subject in Elixir/Nx, and Google was useless. It seems to have gotten steadily worse over the past few years, and in the past few months dramatically so. Maybe it’s just me.

pejrich · February 10, 2023, 2:54am

FYI, scholar is not available via hex, or on hex.pm. I had to pull it manually, even though the docs say to add it as a dep in mix.exs.

polvalente · February 10, 2023, 3:09am

Yes, the library hasn’t been released yet. The README probably still has the default autogenerated instructions.

lucaong · February 10, 2023, 11:15am

Pitching in just to say that your function is formally correct, although you could simplify it a bit by using existing Nx functions for the dot product and the norm (possibly benefiting from some optimizations):

defmodule CosSim do
  import Nx.Defn

  defn cosine_similarity(a, b) do
    Nx.dot(a, b) / (Nx.LinAlg.norm(a) * Nx.LinAlg.norm(b))
  end
end

You can verify that it’s correct with some sanity checks:

# Similarity for vectors pointing in the same direction is 1
a = Nx.tensor([1, 2])
b = Nx.tensor([2, 4])
CosSim.cosine_similarity(a, b)
# =>
# #Nx.Tensor<
#   f32
#   1.0
# >

# Similarity for orthogonal vectors is 0
a = Nx.tensor([3, 0])
b = Nx.tensor([0, 3])
CosSim.cosine_similarity(a, b)
# =>
# #Nx.Tensor<
#   f32
#   0.0
# >

# Similarity for vectors pointing in opposite directions is -1:
a = Nx.tensor([5, 0, 2])
b = Nx.tensor([-10, 0, -4])
CosSim.cosine_similarity(a, b)
# =>
# #Nx.Tensor<
#   f32
#   -1.0
# >

# Vectors at 45deg have cosine similarity = cos(Pi/4) ~ 0.7071
a = Nx.tensor([1, 0])
b = Nx.tensor([1, 1])
CosSim.cosine_similarity(a, b)
# =>
# #Nx.Tensor<
#   f32
#   0.7071067690849304
# >

Note though that the definition of the cosine distance in scholar also has additional special cases for when the L2 norm is so small that it would cause numerical stability issues, or when one or both the operands have norm equal to zero. Additionally, it allows to calculate the cosine distance between multiple vectors in batch, by passing rank-2 tensors as arguments.

pejrich · February 10, 2023, 6:10pm

Thanks @lucaong and @polvalente for all the help. I’m making my way through the porting of the code, and learning a lot, but it’s all very new to me. The cosine_similarity function that scikit uses(which happens to be what the code i’m porting is using too), seems to do something a bit different, and I’m wondering what the name for it might be, as I’m not even sure what to start googling for. Here’s an example of what it returns.

from sklearn.metrics.pairwise import cosine_similarity
>>> a = np.matrix([[1,2,3], [4,5,6]])
>>> b = np.matrix([[9,8,7], [6,5,4]])
>>> cosine_similarity(a, b)
array([[0.88265899, 0.85280287],
       [0.96546332, 0.94805195]])

Now, if I use the axes option from the Scholar formula. axes: [1] gives me:

#Nx.Tensor<
  f32[2]
  EXLA.Backend<host:0, 0.2926165521.2327445524.101050>
  [0.1173410415649414, 0.05194807052612305]
>

Which if I subtract from 1 as Paulo mentioned before, I do find two of those numbers, the top left to bottom right diagonal.

iex(52)> 1 - 0.1173410
0.882659
iex(53)> 1 - 0.051948
0.948052

I’m not sure where the other two come from, or how this would be calculated, or what it’s called. Any insight? I’m guessing maybe pairwise was the first clue I should I have followed.

lucaong · February 10, 2023, 6:32pm

I’d guess that when you pass matrices to cosine_similarity in scikit, it considers each matrix a collection of vectors, and returns the pairwise similarities between each vector in the first matrix and each vector in the second, therefore you get four elements as a result.

pejrich · February 10, 2023, 6:35pm

Yup, that was it!

>>> [cosine_similarity(a[i], b[j]) for i in range(2) for j in range(2)]
[array([[0.88265899]]), array([[0.85280287]]), array([[0.96546332]]), array([[0.94805195]])]

lucaong · February 10, 2023, 7:42pm

You could implement the same behavior as scikit (pairwise cosine similarities between vectors in the given matrices) without iteration, by normalizing along the last axis, transposing the second matrix, and performing a matrix multiplication:

defmodule Pairwise do
  import Nx.Defn

  defnp l2_norm(x, opts \\ []) do
    (x * x)
    |> Nx.sum(opts)
    |> Nx.sqrt()
  end

  defnp normalize(x, opts \\ []) do
    x / l2_norm(x, axes: opts[:axes], keep_axes: true)
  end

  defn cosine_similarity(a, b) do
    normalized_a = normalize(a, axes: [-1])
    normalized_b = normalize(b, axes: [-1])

    Nx.dot(normalized_a, Nx.transpose(normalized_b))
  end
end

You can now use it passing two vectors:

a = Nx.tensor([1, 2, 3])
b = Nx.tensor([9, 8, 7])

Pairwise.cosine_similarity(a, b)
# =>
#  #Nx.Tensor<
#    f32
#    0.8826588988304138
#  >

Or passing two matrices, and getting pairwise similarities, like with scikit:

a = Nx.tensor([[1, 2, 3], [4, 5, 6]])
b = Nx.tensor([[9, 8, 7], [6, 5, 4]])

Pairwise.cosine_similarity(a, b)
# =>
# #Nx.Tensor<
#   f32[2][2]
#   [
#     [0.8826588988304138, 0.8528028130531311],
#     [0.9654632806777954, 0.948051929473877]
#   ]
>

pejrich · February 10, 2023, 8:22pm

I wish I could give you more than one heart. I’ve been working on a solution since I last commented 2 hours ago. I got it working without defn, but stupidly I checked the performance against the sklearn one(no idea why, I knew mine would be slower, and doesn’t really matter), but of course I wanted to get it working with defn, and I was trying to get something similar to for i <- 0..w1, j <- 0..w2, do: cosine_similarity(t1[i], t2[j]) working with a nested while in defn, very ugly funneling all those parameters through, only to finally learn I can’t update a tensor in a while loop in defn. At least that’s what I think this is telling me:

** (CompileError) iex:242: the do-block in while must return tensors with the same shape, type, and names as the initial arguments.

Body matches template:

{#Nx.Tensor<
   f32[2]
 >, #Nx.Tensor<
   s64
 >, #Nx.Tensor<
   s64
 >, #Nx.Tensor<
   s64
 >, #Nx.Tensor<
   s64[2][5]
 >, #Nx.Tensor<
   s64[3][5]
 >}

and initial argument has template:

{#Nx.Tensor<
   f32[1]
 >, #Nx.Tensor<
   s64
 >, #Nx.Tensor<
   s64
 >, #Nx.Tensor<
   s64
 >, #Nx.Tensor<
   s64[2][5]
 >, #Nx.Tensor<
   s64[3][5]
 >}

I’ve been using elixir for about 5-6 years, so it feels weird like i’m coding with my hands tied behind my back. I was just about ready to throw in the towel, but you came and saved the day for me. Thank you!

NduatiK · February 12, 2023, 6:22am

This error has to do with the first argument of your tuple. You passed in a f32[1] and are returning a f32[2].

I don’t have your implementation, so best I can do is suggest you pad the input to make it f32[2] or drop one of the numbers from the f32[2] output