Description:
In spd_learn/models/matt.py, the AttentionManifold module contains a logical bug in the attention mechanism. The Softmax is applied to the wrong dimension, and a subsequent permute operation causes the downstream einsum to compute an inverted attention mapping. No runtime error is thrown currently only because the sequence lengths are identical ($L_q = L_k = n_patches$).
Root Cause:
1. Wrong Softmax Dimension: atten_energy has a shape of (batch, n_patches, n_patches) which acts as a $Q \times K$ matrix. Applying F.softmax(..., dim=-2) incorrectly normalizes over the Query dimension instead of the Key dimension (dim=-1).
2. Inverted Attention: atten_prob.permute(0, 2, 1) flips the tensor to $K \times Q$. When passed into log_euclidean_mean, the einsum("...pq,...qij->...pij") matches $p$ to $K$ and $q$ to $Q$. This mathematically forces the Keys to aggregate the Values using the Query indices, which is backward.
Proposed Solution:
Normalize over the Key dimension (dim=-1) and remove the permute operation. This keeps the shape as $(batch, Q, K)$, which perfectly aligns with the einsum logic in log_euclidean_mean.
# spd_learn/models/matt.py -> AttentionManifold.forward
atten_energy = log_euclidean_distance(Q_expand, K_expand)
atten_weights = 1 / (1 + torch.log1p(atten_energy))
# FIX 1: Normalize over the Key dimension
atten_prob = F.softmax(atten_weights, dim=-1)
# FIX 2: Remove the permute step
# atten_prob = atten_prob.permute(0, 2, 1) <-- Delete this line
# The (batch, Q, K) tensor now perfectly matches the einsum "...pq,...qij->...pij"
output = log_euclidean_mean(atten_prob, V)
return output
Description:$L_q = L_k = n_patches$ ).$Q \times K$ matrix. Applying $K \times Q$ . When passed into $p$ to $K$ and $q$ to $Q$ . This mathematically forces the Keys to aggregate the Values using the Query indices, which is backward.$(batch, Q, K)$ , which perfectly aligns with the
In
spd_learn/models/matt.py, theAttentionManifoldmodule contains a logical bug in the attention mechanism. TheSoftmaxis applied to the wrong dimension, and a subsequentpermuteoperation causes the downstreameinsumto compute an inverted attention mapping. No runtime error is thrown currently only because the sequence lengths are identical (Root Cause:
1. Wrong Softmax Dimension:
atten_energyhas a shape of(batch, n_patches, n_patches)which acts as aF.softmax(..., dim=-2)incorrectly normalizes over the Query dimension instead of the Key dimension (dim=-1).2. Inverted Attention:
atten_prob.permute(0, 2, 1)flips the tensor tolog_euclidean_mean, theeinsum("...pq,...qij->...pij")matchesProposed Solution:
Normalize over the Key dimension (
dim=-1) and remove thepermuteoperation. This keeps the shape aseinsumlogic in log_euclidean_mean.