128K 컨텍스트 창의 진짜 병목 — Attention의 이차 복잡도와 그 너머

128K 토큰짜리 컨텍스트 창이 이제는 흔한 사양이 됐다. Claude, GPT-4o, Gemini 1.5는 이미 그 선을 넘었고, 일부 모델은 백만 토큰을 공언하기도 한다. 그런데 이 숫자들이 실제로 얼마나 잘 “작동”하는지는 별개의 이야기다.

컨텍스트 창 확장의 병목은 두 갈래다. 하나는 이미 많이 알려진 KV 캐시(Key-Value Cache) 메모리 비용이고, 다른 하나는 Attention 연산 자체의 계산 복잡도다. 전자는 얼마나 많은 메모리가 필요한가의 문제이고, 후자는 얼마나 많은 연산이 필요한가의 문제다. 오늘 이야기는 후자다. 메모리가 충분해도, 또는 메모리 문제를 기술적으로 우회해도, 연산량이라는 근본적인 장벽이 따로 존재한다.

Self-Attention이 이차 함수인 이유

트랜스포머(Transformer)의 Self-Attention은 본질적으로 시퀀스 안의 모든 토큰 쌍을 서로 비교하는 연산이다. 시퀀스 길이를 n이라 하면, 비교해야 하는 쌍의 수는 n × n, 즉 n²이다. 이것이 O(n²) 복잡도라고 표현하는 이유다.

실제 숫자로 들여다보면 체감이 다르다. 시퀀스 길이가 1K(1,024 토큰)일 때 Attention 행렬의 원소 수는 약 100만(10⁶)이다. 4K로 늘리면 16배가 된다. 128K에서는 1K 기준 대비 16,384배다. 계산량이 선형으로 늘어나는 것이 아니라 제곱으로 폭발하기 때문에, 시퀀스를 두 배 늘릴 때마다 Attention에 소요되는 FLOPS(초당 부동소수점 연산 수)는 네 배가 된다. 4K → 128K 구간만 봐도 이미 1,024배의 연산량 증가가 발생한다.

여기서 중요한 뉘앙스가 있다. 모델 전체의 추론 비용은 단순히 Attention만으로 결정되지 않는다. 피드포워드 레이어(FFN)의 비용은 O(n)이다. 시퀀스가 짧을 때는 FFN이 전체 연산의 지배적인 부분을 차지하고, Attention은 상대적으로 작다. 그런데 시퀀스 길이가 늘어날수록 이차 항의 무게가 커지고, 어느 지점을 넘으면 Attention이 전체 비용을 압도하기 시작한다. 대략 모델 차원 크기(d_model)와 비슷한 수준의 시퀀스 길이를 넘어서면 Attention이 병목이 된다고 알려져 있다. 대형 모델의 d_model이 보통 수천에서 1만 수준이니, 수천~수만 토큰 단위가 바로 그 임계점이다.

위치 정보를 어떻게 넣느냐의 문제

계산량과 별개로, 긴 시퀀스에서 모델이 “어느 위치의 토큰인지”를 제대로 인식하는 것 자체도 쉽지 않은 문제다. 원래 Transformer 논문(2017년 Vaswani 외)은 절대 위치 인코딩(Absolute Positional Encoding)을 사용했다. 각 위치에 고정된 사인·코사인 패턴을 더해주는 방식이다. 이 방식은 학습 때 봤던 최대 위치까지는 작동하지만, 그 이상에서는 완전히 다른 값이 나온다. 학습 때 본 적 없는 위치 번호가 등장하면 모델은 그것을 어떤 위치인지 알 수 없다.

학습 가능한 절대 위치 임베딩(GPT 계열)도 비슷한 한계가 있다. 학습 최대 길이를 512나 2K로 잡았다면, 그 이상 시퀀스를 추론하는 순간 외삽(extrapolation) 문제가 생긴다. 모델이 훈련 분포 밖의 입력을 받게 되는 것이다.

이를 근본적으로 다르게 접근한 것이 RoPE(Rotary Position Embedding, 회전 위치 임베딩)다. Jianlin Su가 2021년 제안하고 EleutherAI, Meta의 Llama 계열 전체가 채택한 방식이다. RoPE의 아이디어는 단순하다. 절대 위치 정보를 별도로 더하는 게 아니라, Query와 Key 벡터 자체를 위치에 해당하는 각도만큼 복소수 평면에서 회전시킨다. 그러면 두 위치의 Query-Key 내적값이 자연스럽게 두 위치의 상대적 거리만의 함수가 된다. 절대 위치 몇 번인지가 아니라, 두 토큰이 얼마나 떨어져 있는지가 Attention 스코어에 반영된다.

이 덕분에 RoPE는 상대 위치 의존성을 자연스럽게 인코딩하고, 학습 길이를 넘어서도 어느 정도 외삽이 가능하다. Llama 2, Llama 3, Mistral, Qwen, Phi 등 현재 주류 오픈소스 모델이 사실상 모두 RoPE를 채택한 이유다.

그런데 RoPE도 한계가 있다. 회전 각도는 위치 번호와 차원별 주파수의 곱이기 때문에, 학습 때 최대 위치보다 훨씬 먼 거리에서는 회전 각도가 너무 커져 정보가 섞인다. 8K로 학습한 모델을 128K에 그대로 쓰면 Attention 스코어가 급격히 불안정해지는 것이 바로 이 이유다.

RoPE를 학습 없이 늘리는 방법들

모델을 처음부터 긴 컨텍스트로 학습하는 것은 비싸다. 그래서 등장한 것이 RoPE의 주파수 구조를 조작해 더 긴 컨텍스트로 “늘리는” 기법들이다.

NTK-aware interpolation은 회전 각도의 주파수 스케일을 전반적으로 낮춰, 원래 8K 학습 범위의 회전 패턴이 더 넓은 거리 범위에 걸쳐 분포하도록 만든다. 직관적으로 설명하면 “원래 8K 안에 빽빽하게 들어 있던 위치 정보를 더 넓게 펼친다”는 개념이다. 적은 파인튜닝이나 심지어 제로샷으로도 컨텍스트를 수 배 늘릴 수 있어 실용적으로 많이 쓰인다.

YaRN(Yet another RoPE extensioN)은 여기서 한 발 더 나간다. RoPE의 각 주파수 차원이 고주파(짧은 거리 정보)와 저주파(긴 거리 정보)로 자연스럽게 분류된다는 점에 착안해, 고주파 차원은 거의 손대지 않고 저주파 차원만 선택적으로 내삽한다. 전체를 균일하게 늘리는 방식보다 품질 저하가 적다. LongLlama, Mistral의 일부 긴 컨텍스트 변형 등에서 사용된다.

그러나 이 기법들이 “공짜 점심”은 아니다. 원래 학습 길이를 두세 배 수준으로 늘릴 때는 효과가 있지만, 열 배 이상으로 외삽하면 품질이 급락한다. 8K로 학습한 모델을 YaRN으로 128K까지 늘리는 것은 이론적으로 가능하나, 그 결과가 실제 128K 학습 모델과 동등하지 않다. 게다가 연산량의 이차 복잡도 자체는 이 기법들이 건드리지 않는다. 위치 인코딩을 어떻게 늘리든, 128K 토큰의 전체 Attention 행렬을 계산하는 비용은 그대로다.

지역성 가정으로 우회하기 — Sliding Window Attention

Attention의 이차 복잡도를 줄이는 가장 직접적인 방법 중 하나는 각 토큰이 참조할 수 있는 범위를 제한하는 것이다. Mistral 7B가 도입한 Sliding Window Attention(SWA)이 그 대표적인 구현이다. 각 토큰은 바로 앞 w개의 토큰과만 Attention을 계산한다. w가 예컨대 4K라면, 시퀀스 전체가 100K여도 각 토큰의 Attention 범위는 4K로 고정된다. 이렇게 하면 복잡도가 O(n·w), 즉 n에 선형으로 줄어든다.

이 방식의 전제는 “언어적 맥락의 핵심은 가까운 거리에 있다”는 지역성(locality) 가정이다. 많은 경우 이건 사실이다. 대화의 직전 문장, 문단 내 선행 구조, 코드 함수 내 변수 선언 등은 가까운 위치에 있다. Mistral 7B가 작은 크기에 비해 긴 컨텍스트 성능이 준수했던 것도 이 덕분이다.

그러나 지역성 가정이 무너지는 경우가 있다. 긴 소설에서 앞에 등장한 인물이 후반부에 다시 나올 때, 법적 문서에서 첫 조항이 마지막 조항의 맥락이 될 때, 또는 코드베이스에서 수천 줄 위에 정의된 인터페이스를 참조해야 할 때, SWA 모델은 윈도우 밖의 토큰을 사실상 “잊는다”. Mistral 계열 모델이 롤플레이나 긴 문서 추론에서 때로 두드러진 단절을 보이는 것도 이 한계에서 비롯된다.

실제로 Mistral 7B 모델을 롤링 버퍼 캐시(rolling buffer KV cache)와 함께 쓰면, 오래된 토큰의 KV값이 아예 사라진다. 슬라이딩 윈도우의 특성상 이건 의도된 동작이지만, 사용자가 체감하는 것은 “모델이 앞에서 한 말을 갑자기 기억 못 한다”는 이상한 현상이다. 지역성 가정은 단순히 성능 저하가 아니라, 경우에 따라 시스템적 실패로 이어진다.

데이터 처리 복잡도와 연산 구조

이미지 출처: Unsplash

이차를 선형으로 — SSM 기반 대안 아키텍처

Attention 자체를 대체해 복잡도를 O(n)으로 낮추려는 시도도 있다. 대표적인 것이 Mamba와 RWKV다.

Mamba는 SSM(State Space Model, 상태공간 모델)을 기반으로 한다. 핵심 아이디어는 시퀀스를 전체 Attention 행렬 없이, 압축된 숨겨진 상태(hidden state)를 통해 순차적으로 처리하는 것이다. 각 스텝에서 이전 상태를 업데이트하고 새 토큰의 영향을 반영하는 방식이 RNN과 유사하지만, Mamba는 입력에 따라 선택적으로 상태를 업데이트하는 선택적 SSM 구조를 사용해 표현력을 높였다. 추론 시 계산 복잡도는 O(n)이고, 상태 크기는 시퀀스 길이와 무관하게 일정하다.

RWKV는 “Receptance Weighted Key Value”의 약자로, Transformer의 Attention 구조를 선형 Attention으로 대체한 형태다. 학습 때는 행렬 연산 방식으로 병렬 처리가 가능하고, 추론 때는 RNN처럼 순차 처리가 가능해 두 장점을 함께 취한다. 컨텍스트 길이가 늘어도 추론 비용이 선형이라는 점이 핵심 장점이다.

그런데 이 아키텍처들이 Full Attention Transformer를 완전히 대체하지 못하는 이유가 있다. 핵심은 “무엇이든 기억할 수 있는 능력의 범위”다. Full Self-Attention은 시퀀스의 임의의 두 위치 사이의 의존성을 직접 표현할 수 있다. 반면 SSM 계열은 압축된 숨겨진 상태에 모든 과거 정보를 담아야 한다. 상태 벡터의 크기가 고정돼 있으므로, 아무리 긴 시퀀스여도 기억할 수 있는 정보의 양에 구조적 한계가 있다. 특히 Recall-Intensive Task, 즉 시퀀스의 특정 위치에 있는 정보를 정확하게 끄집어내야 하는 작업에서 SSM 계열은 Full Attention 대비 일관되게 낮은 성능을 보인다. 단순한 “맥락 이해”가 아니라 “정확한 검색”이 필요한 순간에 차이가 드러난다.

최근에는 Full Attention과 SSM을 레이어별로 섞는 하이브리드 아키텍처(Jamba, Zamba 등)가 등장했다. 긴 범위 의존성이 중요한 레이어에는 Full Attention을, 나머지는 SSM으로 처리해 품질과 효율을 절충하는 방식이다.

128K를 채워도 모델이 “보지 못하는” 현상

연산 복잡도와 위치 인코딩 문제를 다 해결해도, 긴 컨텍스트 모델에는 또 다른 수수께끼가 남는다. 바로 “Lost in the Middle” 현상이다. 2023년 Stanford 연구팀이 이름 붙인 이 문제는, 다중 문서를 컨텍스트에 넣고 특정 정보를 찾게 했을 때, 그 정보가 컨텍스트의 처음이나 끝에 있을 때는 잘 찾지만 중간에 있을 때는 현저히 성능이 떨어지는 현상이다.

이 현상의 정확한 원인은 아직 활발히 연구 중이지만, 몇 가지 가설이 있다. 학습 과정에서 시작과 끝 부분이 손실 계산에 더 자주 관여하기 때문에 상대적으로 잘 학습된다는 설명이 있다. Attention 패턴 분석에서도 모델이 시퀀스의 시작과 최근 토큰에 Attention을 더 강하게 쏟는 경향이 관찰된다. 구조적으로 볼 때, 정보가 중간에 있을수록 시작 토큰들의 강한 어텐션 편향과 끝 토큰들의 최신성 편향 모두에서 벗어나 있어 상대적으로 주목받지 못하게 된다는 설명도 설득력이 있다.

Needle-in-a-Haystack 테스트는 이 문제를 정량적으로 드러낸다. 긴 더미 텍스트(haystack) 속 특정 위치에 한 줄짜리 핵심 정보(needle)를 숨기고, 모델이 찾을 수 있는지 보는 방식이다. 많은 128K 모델들이 이 테스트에서 U자형 성능 곡선을 보인다. 처음 1~2만 토큰과 마지막 1~2만 토큰 구간에서는 잘 찾지만, 중간 구간에서 정확도가 크게 떨어진다. “컨텍스트 창 128K”라는 사양이 “128K 전 구간을 균일하게 잘 활용한다”를 의미하지 않는다는 것이다.

이 현상은 실용적으로 중요한 함의를 갖는다. 긴 문서를 컨텍스트에 넣을 때, 중요한 정보를 앞이나 뒤에 배치하는 것이 중간에 두는 것보다 실질적으로 더 나은 결과를 낸다. RAG(Retrieval-Augmented Generation, 검색 증강 생성) 시스템 설계에서 이 점을 감안해 검색된 청크의 배치 순서를 조정하는 것이 의미 있다는 연구도 있다.

연산 자체를 최적화하는 접근 — Flash Attention에서 Ring Attention까지

컨텍스트가 길어질 때 이차 복잡도라는 이론적 장벽 외에, 실제 하드웨어에서의 구현 병목도 따로 존재한다. Attention 행렬 전체를 한 번에 GPU 메모리에 올리는 방식은 128K 이상에서 HBM 용량을 초과한다.

Flash Attention(Tri Dao, 2022)은 이 문제를 IO-Aware 알고리즘으로 해결했다. Attention 행렬을 HBM에 완전히 쓰지 않고, 타일(tile) 단위로 나눠 SRAM 안에서 처리한 뒤 축적 결과만 HBM에 쓴다. 이렇게 하면 메모리 이동(IO) 횟수를 크게 줄이고, 연산량 자체가 아니라 메모리 대역폭 병목을 해소한다. Flash Attention 2, 3으로 세대가 거듭되면서 추가 최적화가 쌓였고, Flash Attention 3은 H100의 비동기 텐서 코어와 FP8 연산을 적극 활용해 이전 세대 대비 1.5~2배 속도 향상을 기록했다. 이제는 대부분의 주요 추론 프레임워크가 기본값으로 Flash Attention을 탑재한다.

Ring Attention은 단일 GPU의 한계를 벗어나기 위한 접근이다. 시퀀스를 여러 GPU에 분산하고, 각 GPU가 자신의 시퀀스 청크에 대해 Attention을 계산하는 동안 다른 GPU로부터 KV 청크를 링 형태로 순환 수신한다. 수백만 토큰 컨텍스트를 다수의 GPU에 분산 처리하는 것이 이론적으로 가능해진다. 실제로 Gemini 1.5가 100만 토큰 컨텍스트를 구현하는 데 이와 유사한 분산 Attention 전략이 사용됐다.

사실 이 방향에서 기억해야 할 점은, Flash Attention류 최적화는 연산의 이차 복잡도를 없애는 것이 아니라 그 비용을 실제 하드웨어에서 더 효율적으로 집행하는 방법이라는 것이다. 128K 토큰의 Attention이 필요한 FLOPs 수는 그대로이고, Flash Attention은 그것을 더 빠르게 처리하는 구현이다. 이차 복잡도라는 수학적 사실은 변하지 않는다.

결국 지금의 긴 컨텍스트 경쟁은 여러 층의 절충 위에 서 있다. RoPE 확장으로 위치 인코딩을 늘리고, GQA나 MLA로 메모리 비용을 줄이고, Flash Attention으로 연산을 효율화하고, Lost in the Middle 문제는 아직 완전한 해법이 없는 채로 남아 있다. 이차 복잡도의 수학은 여전히 유효하고, 모델이 백만 토큰을 받아들일 수 있다고 해서 그 전체를 균등하게 잘 처리한다는 의미는 아니다. 선형 복잡도 아키텍처가 Full Attention의 Recall 품질을 완전히 따라잡는 날이 오거나, 혹은 하이브리드 구조가 두 세계의 장점을 안정적으로 결합하는 수준에 이르기 전까지, “컨텍스트 창 크기”라는 스펙 숫자와 그 숫자가 실제로 의미하는 이해 능력 사이의 간극은 계속 존재할 것이다.

출처