关于双曲空间中计算K和V的相似性问题

[chenying0722]

Hypformer中在应用注意力机制时似乎只计算了(K_T)V，这种做法在双曲空间中是有效的吗？感觉在双曲空间中计算两个点的相似性用双曲距离更或者切空间内积（例如文中的公式（5））合理一些，请问作者对这个问题有好的解释吗？不胜感激！

[marlin-codes]

Lorentzian Inner Product For points i, j on the hyperboloid ($i_t = \sqrt{|\mathbf{i}_s|^2 + 1}$):

$$ \langle \mathbf{i}, \mathbf{j} \rangle_L = -i_t j_t + \mathbf{i}_s \cdot \mathbf{j}_s $$

Squared Lorentzian Distance

$$ |\mathbf{i} - \mathbf{j}|_L^2 = -2 - 2\langle \mathbf{i}, \mathbf{j} \rangle_L $$

Let us rearrange:

$$ \langle \mathbf{i}, \mathbf{j} \rangle_L = -\frac{1}{2}|\mathbf{i} - \mathbf{j}|_L^2 - 1 $$

Spatial Component Relationship
Substitute $i_t = \sqrt{|\mathbf{i}_s|^2 + 1}$ into the inner product:

$$ \langle \mathbf{i}, \mathbf{j} \rangle_L = -\sqrt{|\mathbf{i}_s|^2 + 1}\sqrt{|\mathbf{j}_s|^2 + 1} + \mathbf{i}_s \cdot \mathbf{j}_s $$

Approximation for Nearby Points
For small spatial components ($|\mathbf{i}_s|, |\mathbf{j}_s| \ll 1$):

$$ \sqrt{|\mathbf{i}_s|^2 + 1} \approx 1 + \frac{|\mathbf{i}_s|^2}{2}, \quad \sqrt{|\mathbf{j}_s|^2 + 1} \approx 1 + \frac{|\mathbf{j}_s|^2}{2} $$

$$ \langle \mathbf{i}, \mathbf{j} \rangle_L \approx -\left(1 + \frac{|\mathbf{i}_s|^2 + |\mathbf{j}_s|^2}{2}\right) + \mathbf{i}_s \cdot \mathbf{j}_s $$

Substitute into $|\mathbf{i} - \mathbf{j}|_L^2$:

$$ |\mathbf{i} - \mathbf{j}|_L^2 \approx |\mathbf{i}_s - \mathbf{j}_s|^2 $$

Thus:

$$ \mathbf{i}_s \cdot \mathbf{j}_s = \frac{|\mathbf{i}_s|^2 + |\mathbf{j}_s|^2 - |\mathbf{i}_s - \mathbf{j}_s|^2}{2} \approx -\frac{1}{2}|\mathbf{i} - \mathbf{j}|_L^2 + \text{constant} $$

Attention Mechanism
HRC computes:

$$ \text{Attention} = \text{softmax}\left(\frac{\mathbf{Q}_s \mathbf{K}_s^\top}{\sqrt{d}}\right) \approx \text{softmax}\left(-\frac{|\mathbf{i} - \mathbf{j}|_L^2}{2\sqrt{d}}\right) $$

[marlin-codes]

QK^TV==Q(K^TV) # Matrix Associativity

In this work, we use Q(K^TV) to speed up the computation efficiency and get the attentive outcome. But if we want to explicitly get the attention score, we still need to compute QK^T.

[chenying0722]

十分感谢您详尽的推导和解释，数学真的很奇妙！