Матрица попарного сходства между множеством векторов в PyTorch

Question

Предположим, что у нас есть 3D-тензор PyTorch, где первое измерение представляет batch_size следующим образом:

import torch
import torch.nn as nn
x = torch.randn(32, 100, 25)

То есть для каждого i, x[i] является набором 100 25-мерных векторов. Я хотел бы вычислить сходство (например, косинусное сходство - но в целом любую такую ​​попарно матрицу расстояния / сходства) этих векторов для каждого элемента партии.

То есть для каждого x[i] мне нужно вычислить матрицу [100, 100], которая будет содержать попарные сходства вышеуказанных векторов. Более конкретно, (i, j) -ый элемент этой матрицы должен содержать сходство (или расстояние) между i-й и j-й строками (100x25) x[t] для всех t=1, ..., batch_size.

Если я использую torch.nn.CosineSimilarity(), независимо от того, что dim я использую, результатом будет либо [100, 25] (dim=0), либо [32, 25] (dim=1), где мне нужно тензор размера [32, 100, 100]. Я ожидал бы, что torch.nn.CosineSimilarity() будет работать таким образом (поскольку, по крайней мере для меня, это выглядит более интуитивно понятным), но это не так.

Может ли это быть сделано с помощью чего-то вроде ниже?

torch.matmul(x, x.permute(0, 2, 1))

Я полагаю, что это может дать матрицу расстояний, но что, если мне понадобится произвольная парная операция? Должен ли я построить эту операцию, используя описанную выше?

Или, может быть, я должен повторите x таким образом, чтобы я мог использовать встроенный torch.nn.CosineSimilarity()?

Спасибо.

Answer 1

Если вы внимательно прочитаете документацию nn.CosineSimilarity и nn.PairwiseDistance, вы увидите, что они не вычисляют все попарные сходства / расстояния (как вам требуется) , а скорее ожидает два входа одинаковой формы и вычисляет сходства / расстояния между всеми соответствующими точками * только 1010 *.
То есть, если у вас есть два набора по 100 векторов в В 32 измерениях эти функции будут вычислять сходство / расстояние между i -ым вектором в первом наборе и соответствующим i -ым вектором во втором наборе, в результате чего получится только 100 значений сходства / расстояния.

Если вы хотите вычислить все попарные расстояния, вам нужно будет их вручную вычислить.
Использование torch.matmul кажется шагом в правильном направлении.

Если вы ищете эффективный способ вычисления расстояний L2, вы можете найти метод в этом ответе полезным.