Positional Encodings and PaTH Attention

Llama 부터 Qwen, DeepSeek, Gemma 까지 — 2024-2026 의 LLM 들은 거의 예외 없이 RoPE 를 쓴다. 그래서 "position encoding 은 풀린 문제" 라는 분위기가 굳어졌다. PaTH 는 그 합의에 균열을 내는 논문이고, 이 강의는 그 균열의 입구.

Songlin 이 강의에서 풀고 싶어하는 질문은 두 갈래로 나뉜다.

1. position encoding 이 attention 안에서 정확히 무엇을 하고 있는가 — 단순한 "토큰의 절대 위치 부여" 로 보면 핵심을 놓친다. PE 는 token 쌍 간의 transformation 을 결정하는 함수다.
2. 그 transformation 이 input 에 의존해야 하는가 — RoPE 는 "rotate by mθᵢ" 형태로 입력과 무관하다. 같은 두 토큰이라도 문맥에 따라 관계가 달라야 하지 않나?

이 질문이 답으로 이어지는 사다리는 명확하다. (a) RoPE 의 수학적 형태를 다시 쓴다 → (b) 그 형태가 가지는 표현력 한계를 짚는다 → (c) input-dependent transformation 으로 일반화한다 → (d) 그것이 GPU 위에서 효율적으로 학습 가능한가 답한다. PaTH 는 (c) 와 (d) 를 동시에 푸는 답안.

Self-attention 의 핵심은 q · k 내적이다. 거기에 position 정보를 어떻게 끼워 넣는지가 PE 설계의 전부다. 역사적으로 세 가지 방식이 있었고, 각자가 다른 방식으로 "토큰 쌍 사이의 변환" 을 정의한다.

이 frame 으로 보면 PE 설계의 핵심 자유도가 두 개로 압축된다.

• 변환의 종류 — orthogonal 인가 (RoPE, PaTH), 그냥 bias 인가 (ALiBi), 학습 가능한 임의 행렬인가.
• 변환의 인자 — 절대 위치 i 만 보는가 (absolute), 상대 위치 j-i 만 보는가 (RoPE/ALiBi), 아니면 그 사이를 지나간 토큰들까지 보는가 (PaTH).

RoPE 의 본질은 한 줄 — 2D plane 마다 회전 행렬을 곱한다. d 차원의 query/key 를 d/2 개의 2D plane 으로 자르고, 각 plane k 에 대해 frequency θₖ 의 회전을 i 번 적용한다. 이게 합리적으로 보이는 이유와, PaTH 가 이를 일반화하려는 이유.

# RoPE — 2D plane 으로 자른 형태
def rope(x, position, theta):
    # x: [..., d]  d 는 짝수
    cos = torch.cos(position * theta)
    sin = torch.sin(position * theta)
    x1, x2 = x[..., ::2], x[..., 1::2]
    rot1 = x1 * cos - x2 * sin
    rot2 = x1 * sin + x2 * cos
    return torch.stack([rot1, rot2], dim=-1).flatten(-2)

# 핵심 수학적 사실
# R(α) · R(β) = R(α + β)
# 즉 i 위치의 q 와 j 위치의 k 의 내적은
# q·k = (R(iθ) Wq x_i) · (R(jθ) Wk x_j)
#     = (Wq x_i) · R((j-i)θ) · (Wk x_j)
# → 상대 위치 j-i 만 score 에 영향

이 한계가 PaTH 의 출발점이고, 동시에 state-space model 계열 (Mamba, RWKV) 의 동기와도 겹친다. SSM 들은 "상태가 시퀀스를 따라 흐르며 각 토큰이 그 상태를 변형한다" 는 전제 위에서 짜여 있다 — input-dependent transformation 이 자연스럽게 들어가 있다는 뜻. PaTH 는 self-attention 구조 안에서 같은 idea 를 시도한다.

PaTH 의 핵심 아이디어를 한 문장으로 — token i 의 position transformation 은 i 까지 누적된 Householder reflection 들의 곱이고, 각 Householder 의 방향 벡터 u 는 해당 토큰의 입력 으로 결정된다.

왜 Householder 인가 — 두 가지 이유.

• orthogonal 이 보장된다. H = I − 2uuᵀ 는 임의의 단위벡터 u 에 대해 자동으로 orthogonal. RoPE 가 가진 norm-preservation 성질을 그대로 유지.
• compact representation — Hₖ 자체를 dxd 행렬로 들고 다닐 필요 없이 d 차원 벡터 uₖ 만 저장하면 된다. n 토큰 sequence 면 n·d 메모리만 추가.

그러나 큰 문제 — 곱셈이 비가환이다. RoPE 의 R(α)·R(β) = R(α+β) 같은 trick 이 안 통한다. 두 토큰의 상대 transformation 은 사이의 모든 H 들의 곱에 의존한다. 단순 구현은 attention score 한 쌍마다 j−i 개의 d×d 행렬 곱을 요구한다. 이걸 GPU 위에서 풀어내는 게 § 07 의 문제.

PaTH 가 단순한 "재미있는 PE" 에서 실제 사용 가능한 메커니즘으로 넘어가는 자리. 논문이 말하는 parallel algorithm for products of Householder matrices 의 의미를 풀어 본다.

# compact WY representation — 의사코드
def householder_product(U):
    # U: [d, n]  단위벡터들
    n = U.shape[1]
    T = torch.zeros(n, n)
    for i in range(n):
        T[i, i] = -2
        for j in range(i):
            # Schreiber-Van Loan recursion
            T[j, i] = -2 * (T[j, :j+1]
                            @ (U[:, :j+1].T @ U[:, i]))
    return torch.eye(U.shape[0]) + U @ T @ U.T

# 한 번만 계산해두면, q_i 에 적용은
# q_i = (I + U[:, :i] @ T[:i,:i] @ U[:, :i].T) @ Wq @ x_i
# prefix 가 늘어날수록 U[:, :i] 의 column 이 늘어남

FlashAttention 형 블록화와 결합하면 — chunk 안에서는 dense 한 d×d 행렬을 한 번 만들어두고 q tile 전체에 broadcast, chunk 사이에서는 prefix product 를 누적. 이게 PaTH 가 "FlashAttention-style blockwise variant" 라고 부르는 구현이다.

RoPE 의 약점이 가장 노출되는 자리 — 학습 때 보던 길이보다 긴 sequence 에서 perplexity 가 폭발한다. 이 현상이 NTK-aware scaling, YaRN, LongRoPE 같은 후처리 트릭의 산업을 낳았다. PaTH 는 이 문제를 architectural 으로 다시 본다.

구체적인 수치 비교는 논문의 "synthetic benchmark + moderate-scale LM" 결과에서 PaTH 가 RoPE 와 최근 baseline 들을 모두 outperform 한 것으로 보고된다. 학습 길이 8K 에서 32K, 64K 까지의 perplexity 곡선이 RoPE 보다 훨씬 평평하다는 그림이 강의에서 등장 — 확인 필요.

PaTH 의 수학적 우아함이 GPU 위에서 살아남으려면 — chunk-wise 누적과 FlashAttention 의 tile reduction 이 만나야 한다. 이 섹션은 강의에서 Songlin 이 "여기가 진짜 엔지니어링" 이라고 강조한 부분.

실제 구현은 ThunderKittens 또는 직접 CUDA 로 짜는 것이 자연스럽다 — Triton 으로도 가능하지만 chunk 사이 prefix scan 의 통제가 까다롭다. 강의에서 Songlin 이 PaTH 의 reference 구현이 어디에 있는지 명시했는지는 확인 필요이지만, FlashLinearAttention 같은 Songlin 그룹의 라이브러리 안에 들어있을 가능성이 높다.

PaTH 논문은 두 종류의 실험을 보여준다 — (a) 알고리즘적 reasoning 같은 synthetic 작업에서의 in-context generalization, (b) 일반 LM perplexity 에서의 length extrapolation. 강의에서 다룬 핵심 수치들 — 단, 모든 절대 수치는 확인 필요.

또 하나 강조할 만한 결과 — RoPE 모델을 PaTH 로 "converted continued pretraining" 으로 옮길 수 있다는 점. 즉 처음부터 PaTH 로 학습할 필요가 없고, 기존 Llama 같은 RoPE 모델을 짧은 추가 학습으로 PaTH 모델로 전환할 수 있다. 이 사실이 산업적 채택 비용을 크게 낮춘다 — 논문이 NeurIPS 2025 에 accepted 된 결정적 이유 중 하나.

강의에서 Songlin 이 명시적으로 짚은 한계와, 학습 노트로서 추가하는 caveat 들.

• chunk size 선택의 trade-off — chunk 가 크면 T 의 B² 비용이 커지고, 작으면 chunk 사이 prefix product 가 늘어난다. 최적 B 는 head 차원과 sequence 길이의 함수 — 자동 선택 가능한지 확인 필요.
• backward pass 의 비용 — compact WY 의 forward 는 깔끔하지만 backward 의 atomic accumulation 비용이 RoPE 보다 분명히 크다. 학습 단계의 wallclock impact 가 throughput 표에 다 반영되어 있는지 확인 필요.
• cross-attention / encoder-decoder — 강의는 causal LM 에 집중. cross-attention 처럼 양방향이거나 두 sequence 가 다른 모델에서는 PaTH 의 "지나온 토큰들의 누적" 정의 자체를 다시 짜야 한다.
• multi-modal — image patch 의 2D 위치, vision-language 모델의 mixed sequence 등에서 "지나온 토큰" 의 의미가 자명하지 않다.
• theoretical 보장 부재 — 길이 일반화에 대한 formal proof 는 없다. 실험적으로 잘 되지만 "언제 깨지는가" 의 경계는 확인 필요.

손에 새기기 — 실습 시퀀스

1. RoPE 직접 짜보기 — 64-dim head 위에 cos/sin 두 줄 구현. positional id [0..N) 에서 frequency aliasing 이 어디서 시작하는지 N 을 키우며 직접 본다.
2. compact WY 실험 — n=16 짜리 random unit vector U 로 ∏Hₖ 를 두 방식으로 (직접 곱 / WY form) 비교. forward 결과가 일치하는지, backward gradient 는 어떻게 흘러가는지 확인.
3. PaTH 의 chunk size 영향 — toy attention 위에 chunk B ∈ {32,64,128,256} 을 sweep. wallclock vs perplexity 의 frontier 그림을 그린다.
4. length extrapolation 합성 실험 — random key copying task. 학습 길이 1K, 평가 길이 [1K..32K]. RoPE / RoPE+YaRN / PaTH 세 가지 곡선을 직접 그린다.
5. RoPE → PaTH conversion — 100M 짜리 작은 LM 을 RoPE 로 학습한 뒤 1K step 정도 PaTH 로 continued pretraining. PPL 회복 곡선이 논문 설명과 비슷한 모양인지 확인.
6. 커널 변형 — FlashAttention triton kernel 에 PaTH 의 prefix product 를 끼워 넣은 prototype. ThunderKittens 위에서도 동일하게 시도, 어느 쪽이 short attention 위에서 더 잘 mapping 되는지 본다.

• throughput 손실의 정확한 수치 — RoPE FA 대비 PaTH 의 학습 wallclock 이 얼마나 더 드는지. § 08 의 0.85× 는 추정. 강의에서 직접 수치를 댄 자리가 있는지 확인.
• chunk size B 의 자동 선택 — head 차원 d, sequence 길이 N 에 따른 최적 B. heuristic 또는 autotune 결과.
• RoPE → PaTH conversion 의 step 수 — 강의에서 "짧은 continued pretraining" 이라고 했지만 — 실제로는 몇 K step? data mix 는?
• multi-modal / cross-attention 적용 — vision-language 모델, encoder-decoder 의 cross attention 에서 PaTH 가 동작하는지의 후속 질문.
• PaTH 와 GLA / DeltaNet 의 unification — 모두 input-dependent state evolution 으로 환원 가능한가. Songlin 그룹의 다음 작업이 이걸 다룰 가능성.
• NTK-aware / LongRoPE 와의 직접 비교 — 64K, 128K 길이에서 PaTH (no fine-tune) vs LongRoPE (fine-tuned) 의 fair comparison.