트랜스포머 — 어텐션을 손으로 계산한다 · 회전 위치 임베딩(RoPE) · 실습

위치를 회전으로 넣는다

목표

회전 위치 임베딩(RoPE)을 표준 라이브러리만으로 직접 만든다. 짝수 차원을 두 개씩 쌍으로 묶어 쌍마다 다른 속도로 자리만큼 돌리고, 그렇게 돌린 질의와 키의 내적이 두 자리의 차이에만 의존한다는 것을 숫자로 확인한다. 더하는 방식(절대 위치 인코딩)을 나란히 만들어 같은 간격에서 값이 흔들리는 것을 견주고, 회전이 길이를 바꾸지 않는다는 것과 거리가 멀어질 때 어느 성분이 남는지를 잰다.

왜 중요한가

자리 벡터를 더하면 내적을 펼쳤을 때 q·P[n] 과 P[m]·k 처럼 한쪽 자리만 든 항이 남는다. 그 항들은 차이로 묶이지 않으므로, 간격이 똑같아도 문장 앞쪽에서 잰 값과 뒤쪽에서 잰 값이 달라진다. 언어에서 중요한 것은 대개 몇 칸 앞의 낱말인가인데 점수에는 절대 자리가 섞여 들어와 있는 것이다.
회전은 그 문제를 연산 자체로 푼다. 두 벡터를 같은 방향으로 같이 돌리면 사이 각이 변하지 않으므로, 각각 제 자리만큼 돌린 뒤 내적하면 값이 차이에만 의존한다. 이것은 근사가 아니라 등식이고, 그래서 숫자로 확인할 수 있다.
이 실습은 모델을 부르지 않는다. 이 파드의 시스템 파이썬에는 numpy 가 없고 인터넷도 없다. math 만으로 만들고, 여기서 잰 숫자만 쓴다. 그래서 "어느 모델이 어떤 밑을 쓴다" 같은 말은 여기서 하지 않는다.
채점기는 여러분이 적어 둔 설명을 믿지 않는다. 여러분의 모듈을 실제로 불러 매번 다른 차원과 자리로 함수를 두드려 보고, 채점기가 따로 계산한 값과 허용 오차 안에서 대조한다. 입력은 실행마다 바뀌므로 값을 외워 넣을 수 없다.

단계

1. /root/work/tf-rope/rope.py 에 DIM = 64·THETA_BASE = 10000.0·demo_vectors()·dot(a, b)·thetas(d, base=THETA_BASE) 를 만드세요. thetas 는 쌍마다 다른 회전 속도를 돌려줍니다.
2. rotate_pair(x0, x1, angle) 를 더해 평면 위의 점 하나를 angle 라디안만큼 반시계로 돌리게 하세요.
3. apply_rope(vec, pos, base=THETA_BASE) 를 더해 벡터를 쌍으로 묶고 각 쌍을 pos * theta_i 만큼 돌린 새 벡터를 돌려주게 하세요.
4. rope_score(q, k, m, n, base=THETA_BASE) 를 만들어 m 번 자리의 질의와 n 번 자리의 키를 각각 돌린 뒤 내적하게 하세요.
5. sin_pos(pos, d, base=THETA_BASE) 와 add_score(q, k, m, n, base=THETA_BASE) 를 만들어 더하는 방식의 같은 점수를 재게 하세요.
6. offset_scan(q, k, offset, starts) 를 만들어 같은 간격을 여러 시작 자리에서 재게 하세요. 돌려주는 값은 (시작자리, 회전 점수, 더하는 점수) 짝의 목록입니다.
7. turns(d, delta, base=THETA_BASE) 와 slow_pairs(d, delta, base=THETA_BASE) 를 만들어 간격마다 쌍이 몇 바퀴를 도는지, 아직 한 바퀴를 안 넘긴 쌍이 몇 개인지 재게 하세요.
8. 잰 값을 /root/work/tf-rope/rope_report.json 과 /root/work/tf-rope/rope_report.md 에 기록하세요.

참고

실행 계약: 채점기는 /root/work/tf-rope/rope.py 를 파이썬 모듈로 불러 DIM·THETA_BASE·demo_vectors·dot·thetas·rotate_pair·apply_rope·rope_score·sin_pos·add_score·offset_scan·turns·slow_pairs 를 직접 씁니다. 스크립트로 실행하지 않으므로 if __name__ == "__main__" 은 없어도 됩니다.
demo_vectors() 는 (q, k) 를 돌려줍니다. 난수가 아니라 식으로 못 박습니다 — q[j] = math.cos(0.7 * j + 0.3), k[j] = math.sin(0.4 * j + 1.1) 이고 둘 다 길이는 DIM 입니다.
thetas(d, base) 는 길이 d // 2 의 목록이고 i 번째 값은 base ** (-2 * i / d) 입니다. 첫 값은 늘 1.0 이고 뒤로 갈수록 작아집니다. 차원마다가 아니라 쌍마다 하나입니다.
rotate_pair(1.0, 0.0, math.pi / 2) 는 (0.0, 1.0) 입니다. 반시계 방향으로 못 박습니다. 돌려주는 값은 두 개짜리 튜플입니다.
apply_rope 는 (vec[0], vec[1]), (vec[2], vec[3]) 처럼 이웃끼리 쌍을 묶습니다. i 번째 쌍의 회전각은 pos * thetas(len(vec), base)[i] 이고, 쌍마다 각이 다릅니다. 넘겨받은 목록을 제자리에서 고치지 말고 새 목록을 만드세요. pos 가 0 이면 원래 벡터와 같습니다.
rope_score(q, k, m, n) 은 dot(apply_rope(q, m), apply_rope(k, n)) 입니다. √d 로 나누지 않습니다.
sin_pos(pos, d, base) 는 길이 d 의 목록이고, j 번째 값은 angle = pos / (base ** ((2 * (j // 2)) / d)) 로 두었을 때 j 가 짝수면 math.sin(angle), 홀수면 math.cos(angle) 입니다.
add_score(q, k, m, n) 은 q 에 sin_pos(m, len(q)) 를, k 에 sin_pos(n, len(k)) 를 각각 더한 뒤 내적합니다. 한쪽에만 더하면 안 됩니다.
offset_scan(q, k, offset, starts) 는 starts 의 자리 s 마다 질의를 s + offset 에, 키를 s 에 둡니다. 순서는 starts 와 같습니다.
turns(d, delta, base) 는 길이 d // 2 의 목록이고 i 번째 값은 delta * theta_i / (2 * math.pi) 입니다. slow_pairs 는 그 값이 1.0 보다 작은 쌍의 개수(정수)입니다.
8단계 보고서는 demo_vectors() 의 q·k 로 잽니다. 간격은 2, 시작 자리는 3·10·100·4000, 노름을 재는 자리는 0·1·7·100·4096, 바퀴를 세는 간격은 1·16·256·4096 입니다. JSON 키는 dim·base·offset·starts·rope_scores·rope_spread·add_scores·add_spread·norm_before·norm_max_gap·deltas·slow_pairs·turns_first·turns_last 입니다. turns_first 는 간격마다 0번 쌍의 바퀴 수, turns_last 는 마지막 쌍의 바퀴 수입니다.
rope_spread 와 add_spread 는 각각 네 점수의 최댓값에서 최솟값을 뺀 값입니다. norm_before 는 q 의 노름, norm_max_gap 은 다섯 자리에서 돌린 뒤 노름과 원래 노름의 차이 중 가장 큰 절댓값입니다.
실수를 == 로 견주지 마세요. 채점기는 abs(a - b) <= 1e-9 + 1e-6 * abs(b) 로 봅니다. 자리가 커지면 각도도 커져 오차가 늘어나므로, 이 실습은 자리를 0 부터 4096 까지만 판정합니다.
이 파드에는 인터넷이 없습니다. pip install 은 되지 않고, numpy 는 /opt/onnx-lab/bin/python 안에만 있어 시스템 파이썬에서는 import numpy 가 되지 않습니다. math 만으로 충분합니다.
공식 문서: [RoFormer — Rotary Position Embedding](https://arxiv.org/abs/2104.09864) · [Attention Is All You Need](https://arxiv.org/abs/1706.03762) · [Python math](https://docs.python.org/3/library/math.html)
흔한 실수: 쌍이 아니라 차원마다 각을 매기기, 모든 쌍에 0번 쌍의 각을 쓰기, 회전 방향을 뒤집기, 질의만 돌리고 키는 그냥 두기, 더하는 방식에서 한쪽에만 자리 벡터를 더하기, 바퀴 수를 셀 때 2π 로 나누지 않기, 넘겨받은 목록을 제자리에서 고치기.

단계 8개

쌍마다 다른 회전 속도
평면 위에서 한 번 돌린다
벡터를 자리만큼 돌린다
돌린 질의와 키의 점수
더하는 방식을 나란히 둔다
같은 간격이면 같은 점수다
멀어지면 어느 쌍이 남는가
잰 것을 기록으로 남긴다