문맥을 두 배로 늘리면 무엇이 네 배가 되나

한 줄 요약

트랜스포머 한 층의 계산은 두 몫으로 갈린다. 어텐션의 점수와 섞기는 문맥 길이의 제곱으로 자라고, 나머지 선형 변환은 길이에 비례해서 자란다. 짧은 문맥에서는 뒤쪽이 지배하고 어느 길이를 넘어서면 앞쪽이 지배한다 — 그 자리를 아는 것이 이 글의 전부다.

왜 이게 필요했나

문맥 한도를 늘려 달라는 요청은 어느 팀에나 온다. 4천에서 8천으로, 8천에서 3만 2천으로. 늘리면 얼마나 비싸지느냐는 물음에 "두 배쯤" 이라고 답했다가 청구서를 보고 놀라는 일이 반복된다.

거꾸로도 있다. 문맥을 두 배로 늘렸는데 비용이 두 배밖에 안 늘어 "제곱이라더니 아니네" 하고 넘어가는 경우다. 둘 다 같은 오해에서 나온다. 어느 한쪽만 보고 있는 것이다.

한 층 안에는 성질이 다른 두 계산이 섞여 있다. 어텐션은 자리마다 문맥 전체를 훑으므로 자리 수가 늘면 훑는 일까지 함께 늘어 제곱이 된다. 반면 질의·열쇠·값을 만드는 투영, 출력 투영, 피드포워드는 자리 하나하나에 똑같이 한 번씩 도는 일이라 자리 수에만 비례한다.

그래서 "문맥 길이에 대해 제곱" 이라는 말은 반쪽만 맞다. 정확히는 제곱으로 자라는 몫과 비례해서 자라는 몫의 합이고, 어느 쪽이 지배하는지는 길이와 모델 폭이 함께 정한다.

무엇을 세면 되나

여기서 시간을 재면 안 된다. 같은 코드도 기계와 부하에 따라 다르고, 컨테이너 안에서는 더 흔들린다. 대신 곱셈 횟수를 센다. 횟수는 어디서 돌려도 같은 정수다.

세는 법은 단순하다. 길이 d 인 두 벡터의 내적은 곱셈이 정확히 d 번이다. 나머지는 그 내적이 몇 번 들어가느냐만 세면 된다.

def dot(a, b, ctr):    total = 0.0    for x, y in zip(a, b):        total += x * y    ctr.add(len(a))      # 곱셈 d 번을 장부에 적는다    return total

[어텐션의 정의](https://arxiv.org/abs/1706.03762)를 이 자로 재 보면 이렇게 나온다.

점수 행렬 QK: 줄도 n개, 칸도 n개, 칸마다 길이 d 의 내적 → n * n * d
섞기 AV: 나오는 것은 n x d 로 작지만, 자리마다 문맥 전체 n 개를 훑는다 → 또 n * n * d
헤드를 h 개로 쪼개도 합은 같다. 헤드마다 폭이 d / h 로 줄고 그런 헤드가 h 개이기 때문이다.

그래서 제곱으로 자라는 몫은 2 * n * n * d 다. 반대쪽은 자리 하나가 치르는 값을 한 번만 세고 자리 수를 곱하면 된다.

질의·열쇠·값 투영 셋과 출력 투영 하나 → 자리마다 4 * d * d
피드포워드 두 층 → 자리마다 2 * d * d_ff

합쳐서 자리마다 4 * d * d + 2 * d * d_ff 이고, 자리가 n 개다.

교차점이라는 것

두 식을 나란히 놓으면 교차점이 손으로 풀린다. 2 * n * n * d 가 n * (4 * d * d + 2 * d * d_ff) 이상이 되는 가장 작은 n 을 찾으면 된다. 양변을 2 * n * d 로 나누면 조건은 n >= 2 * d + d_ff 로 줄어든다.

한 층의 곱셈 횟수를 두 몫으로 갈라 그린 그림. 어텐션의 점수와 섞기는 문맥 길이의 제곱으로 자라는 곡선이고 나머지 선형 변환은 길이에 비례하는 직선이다. 두 선은 n 이 2d 더하기 d_ff 인 자리에서 만나며, 그 앞에서는 선형 몫이 뒤에서는 제곱 몫이 지배한다

숫자가 주는 감각이 중요하다. 폭이 넓은 모델일수록 교차점이 뒤로 밀린다. 자리 하나가 치르는 선형 비용이 폭의 제곱으로 커지는데, 어텐션 쪽은 폭에 비례해서만 커지기 때문이다. 큰 모델에서 "제곱이라며 왜 안 느껴지지" 하는 체감은 여기서 나온다 — 아직 교차점 앞에 있는 것이다.

그리고 교차점을 지나면 이야기가 달라진다. 길이를 두 배로 하면 총합이 네 배에 가까워지기 시작한다.

인과 마스크가 버리는 절반

디코더의 어텐션에는 인과 마스크가 있다. i 번째 자리는 자기 자신까지만 보므로 줄 길이가 1, 2, 3 으로 늘어 n 에서 끝난다. 실제로 쓰이는 점수는 n * (n + 1) / 2 개다.

행렬은 여전히 n 제곱 칸이다. 그러니까 절반 가까이는 계산해 놓고 버린다. 정확히는 (n - 1) / (2 * n) 이 버려지고, 길이가 길어질수록 절반에 다가간다.

순진하게 전체를 계산하고 마스크를 씌우는 구현이 바로 그렇게 한다. 계산을 줄이려면 마스크를 나중에 씌우는 것이 아니라 아예 계산하지 않아야 하고, 그래서 블록 단위로 삼각형만 도는 커널들이 나왔다. [PyTorch 의 scaled_dot_product_attention](https://docs.pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html) 이 is_causal 을 따로 받는 것도 같은 이유다. 마스크를 텐서로 받아서 곱하는 것과 "인과다" 라고 말해 주는 것은 안에서 도는 일이 다르다.

메모리도 같은 모양이다

점수 행렬을 통째로 들고 있으면 값의 개수가 헤드마다 n 제곱이다. 합쳐서 n * n * h 개이고, 여기에 자료형 한 칸의 바이트를 곱하면 바이트가 나온다.

이것이 왜 아픈지는 뒤집어 보면 안다. 예산을 정해 놓고 들어갈 수 있는 최대 길이를 구하면 예산을 네 배로 늘려야 길이가 두 배가 된다. 카드를 두 배로 늘려도 문맥은 1.41배밖에 못 늘린다는 뜻이다.

그래서 점수 행렬을 통째로 만들지 않는 구현이 중요해졌다. 블록으로 잘라 한 조각씩 처리하고 버리면 같은 계산을 하면서도 들고 있는 값의 수가 줄어든다. 계산량은 그대로인데 메모리만 줄어드는 것이다 — 이 둘이 따로 논다는 감각이 없으면 "왜 계산은 같은데 더 긴 문맥이 되지" 가 이해되지 않는다.

넣는 길이와 만드는 길이는 다르게 는다

프롬프트를 한 번에 넣는 일(프리필)과 토큰을 하나씩 만드는 일(디코드)은 자라는 모양이 다르다.

프리필은 n 개의 자리를 한꺼번에 처리하므로 점수가 n 제곱 규모로 생긴다. 반면 이미 n 개가 쌓인 뒤에 한 개를 더 붙이면, 새 질의 하나가 n + 1 개의 열쇠를 볼 뿐이다. 줄 하나다. 그래서 g 개를 만드는 동안 생기는 점수의 총합은 g * n + g * (g + 1) / 2 이 된다.

여기서 재미있는 항등식이 하나 나온다. n 을 0 으로 두면 이 값이 causal_pairs(g) 와 정확히 같다. 한 번에 넣어 마스크로 지우든 하나씩 붙이든, 실제로 필요한 점수의 수는 똑같다. 갈리는 것은 그것을 한 번에 하느냐 나누어 하느냐뿐이다. 그 차이를 무엇으로 메우는지는 [KV 캐시 문서](https://huggingface.co/docs/transformers/en/kv_cache)가 다루는 주제이고 다음 모듈의 몫이다.

현장에서 만나는 모습

첫째, 문맥 한도를 두 배로 늘렸는데 지연이 두 배가 안 늘어 안심한다. 아직 교차점 앞이라 선형 몫이 지배하고 있을 뿐이다. 길이를 더 늘리면 갑자기 기울기가 바뀐다.

둘째, 같은 길이인데 모델을 바꾸니 비용 곡선의 모양이 달라진다. 폭이 다르면 교차점이 다른 곳에 있다. 한 모델에서 재어 얻은 배율을 다른 모델에 그대로 쓰면 어긋난다.

셋째, 긴 문맥에서 메모리가 먼저 터진다. 계산은 버티는데 점수 행렬을 들 자리가 없다. 계산량과 메모리가 같은 n 제곱이어도 부딪히는 벽이 먼저 오는 쪽은 대개 메모리다.

넷째, 마스크를 텐서로 만들어 곱하는 구현이 두 배 느리다. 전체를 계산하고 절반을 버리기 때문이다. 같은 수식이라도 언제 마스크를 쓰느냐로 실제 일의 양이 갈린다.

다섯째, 프리필은 느린데 토큰 생성은 빠르다. 또는 그 반대다. 둘은 다른 모양으로 자라므로 하나의 배율로 묶어 견적을 내면 반드시 한쪽이 틀린다.

실무에서 진짜 중요한 것

"제곱" 이라는 말만 외우지 말고 두 몫으로 갈라 보라. 어느 쪽이 지배하는지는 길이와 폭이 함께 정한다.
시간 대신 횟수를 세라. 횟수는 기계가 바뀌어도 같고, 견적을 남에게 설명할 때 근거가 된다.
교차점을 자기 모델의 숫자로 구해 두라. 그 지점 앞뒤로 용량 계획이 달라진다.
계산량과 메모리를 따로 보라. 같은 n 제곱이어도 줄이는 방법이 다르다.
프리필과 디코드를 분리해 재라. 하나의 배율은 반드시 한쪽을 틀리게 만든다.

다음 실습에서 할 것

/root/work/tf-cost/cost.py 를 한 단계씩 키운다. 표준 라이브러리만 쓴다 — 이 파드의 시스템 파이썬에는 numpy 가 없고(/opt/onnx-lab/bin/python 안에만 있다) torch 도 transformers 도 없다. 대신 곱셈을 세는 계수기를 손으로 만들고, 거기서 나온 정수만 쓴다.

계수기와 내적으로 시작해, 점수 행렬과 섞기를 실제로 돌려 n 제곱이 나오는 것을 세어 확인한다. 그다음 자리 하나가 치르는 선형 변환 여섯 개를 실제로 돌려 그쪽이 길이와 무관하다는 것을 확인하고, 두 몫을 나란히 놓은 표를 만든다.

그 표에서 교차점을 찾는다. 모델 폭을 바꿔 가며 교차점이 어디로 움직이는지도 자기 숫자로 본다. 이어서 인과 마스크가 남기는 점수를 세어 절반이 버려지는 것을 확인하고, 점수 행렬을 통째로 들 때의 바이트와 예산 안에 드는 최대 길이를 구한다.

마지막 단계가 이 실습의 요점이다. 한 토큰을 붙일 때 새로 생기는 점수는 줄 하나뿐이라는 것을 세고, 길이 0 에서 g 개를 만들 때의 총합이 인과 마스크가 남긴 수와 정확히 같다는 항등식을 정수로 확인한다. 채점기는 여러분의 모듈을 실제로 불러 매번 다른 크기로 함수를 두드려 보고, 계수기가 실제로 얼마나 올라갔는지까지 대조한다. 값을 외워 넣을 수 없다.