pytorch余弦相似度矩阵与角度矩阵

向量余弦相似度

$$\text{ cosine\_similarity}\left(\mathbf{x},\mathbf{y}\right)=\cos ⟨\mathbf{x},\mathbf{y}⟩=\frac{{\mathbf{x}}^T\mathbf{y}}{\Vert \mathbf{x}\Vert \Vert \mathbf{y}\Vert }$$
其中$⟨\mathbf{x},\mathbf{y}⟩$表示$\mathbf{x}$和$\mathbf{y}$的角度

余弦相似度矩阵

接下来的问题是考虑矩阵$\mathbf{A}=\left(\begin{array}{c}{\mathbf{a}}_1^T\\ {\mathbf{a}}_2^T\\ ⋮\\ {\mathbf{a}}_m^T\end{array}\right),{\mathbf{a}}_i\in {\mathbb{R}}^k,\mathbf{B}=\left(\begin{array}{c}{\mathbf{b}}_1^T\\ {\mathbf{b}}_2^T\\ ⋮\\ {\mathbf{b}}_n^T\end{array}\right),{\mathbf{b}}_i\in {\mathbb{R}}^k$
计算余弦相似度矩阵$\mathbf{C}\in {\mathbb{R}}^{m\times n}$,
其中$c_{ij}=\text{ cosine\_similarity}\left({\mathbf{a}}_i,{\mathbf{b}}_j\right)$

第一种

先单位化，这样可以不用单位化
这里要注意防止分母为0，所以有一个eps
torch.einsum就是一种矩阵乘法

def cosine_similarity(a, b, eps=1e-7):
    """

    :param a: m*k
    :param b: n*k
    :return: m*n cosine similarity
    """
    a = a / torch.max(torch.norm(a, p=2, dim=-1, keepdim=True), torch.tensor(eps))
    b = b / torch.max(torch.norm(b, p=2, dim=-1, keepdim=True), torch.tensor(eps))
    return torch.einsum('ni,mi->nm', a, b)

第二种

其实就是最后一步直接矩阵乘法

def cosine_similarity(a, b, eps=1e-7):
    """

    :param a: m*k
    :param b: n*k
    :return: m*n cosine similarity
    """
    a = a / torch.max(torch.norm(a, p=2, dim=-1, keepdim=True), torch.tensor(eps))
    b = b / torch.max(torch.norm(b, p=2, dim=-1, keepdim=True), torch.tensor(eps))
    return torch.mm(a, b.T)

第三种（推荐）

from torch.nn.functional import cosine_similarity

cosine_similarity(a.unsqueeze(1), b.unsqueeze(0), dim=-1)

可以查看这里https://pytorch.org/docs/stable/generated/torch.nn.CosineSimilarity.html?highlight=cosinesimilarity#torch.nn.CosineSimilarity

$\mathbf{A}\in {\mathbb{R}}^{m\times k}$，经过.unsqueeze(1)会变成$m\times 1\times k$的矩阵
$\mathbf{B}\in {\mathbb{R}}^{n\times k}$，经过.unsqueeze(0)会变成$1\times n\times k$的矩阵

dim=-1，这意味着最后会变成一个$m\times n$的矩阵
因为形状不一样，所以会广播成$m\times n\times k$的矩阵

举个例子
将$3\times 2$的矩阵广播成$3\times 4\times 2$
a经过unsqueeze(1)会变成$3\times 1\times 2$的矩阵
最后的+zeros是为了触发广播

最后计算余弦相似度的时候
a[0,0]得到的是一个2个元素的向量
就是
a[0,0]和b[0,0]计算余弦相似度
a[0,0]和b[0,1]计算余弦相似度
…
a[0,1]和b[0,0]计算余弦相似度
a[0,1]和b[0,1]计算余弦相似度
…
a[m,n]和b[m,n]计算余弦相似度

角度矩阵

得到余弦相似度矩阵后计算$\arccos$，对应pytorch的torch.acos
需要注意的是，pytorch的acos有点问题https://github.com/pytorch/pytorch/issues/8069
如果数值接近-1或1，则会返回nan
解决方法是$\pm eps$

eps=1e-7
matrix_cos=matrix_cos.clamp(min=-1+eps,max=1+eps)

matrix_angle=torch.acos(matrix_cos)