[3D Knowledge] Attention vs Linear Attention

SoftMax Attnetion
     - 보통 Attention이라 부름. ( 둘이 헷갈리면 안됨 )
     - 시퀀스 길이 N에 대해서, 과거 정보를 압축하지 않고 전부 저장
     - Query마다 선택적으로 접근하여 해결

Linear Attnetion
     - Attention을 선형 시간으로 근사/변형한 것
     - 빠르나, Long Sequence에 약함

코드는 다음과 같다

import torch
import torch.nn.functional as F

def softmax_attention(Q, K, V):
    """
    Q: (B, N, D)
    K: (B, N, D)
    V: (B, N, Dv)
    """
    D = Q.shape[-1]
    scores = torch.matmul(Q, K.transpose(-1, -2)) / (D ** 0.5)  # (B, N, N)
    attn = F.softmax(scores, dim=-1)                            # (B, N, N)
    out = torch.matmul(attn, V)                                 # (B, N, Dv)
    return out
    
    
def linear_attention(Q, K, V, eps=1e-6):
    """
    Q: (B, N, D)
    K: (B, N, D)
    V: (B, N, Dv)
    """
    # feature map (예: elu + 1)
    phi = lambda x: torch.elu(x) + 1

    Q_phi = phi(Q)  # (B, N, D)
    K_phi = phi(K)  # (B, N, D)

    # 🔥 과거를 먼저 고정 크기로 압축
    KV = torch.matmul(K_phi.transpose(-1, -2), V)  # (B, D, Dv)

    # query로 읽기
    Z = 1 / (torch.matmul(Q_phi, K_phi.sum(dim=1, keepdim=True).transpose(-1, -2)) + eps)
    out = torch.matmul(Q_phi, KV) * Z               # (B, N, Dv)

    return out