선형 어텐션 원리 정리 — KV 캐시 없이 긴 문맥을 처리하는 트랜스포머의 다음 단계

핵심 메시지

• 표준 소프트맥스 어텐션은 앞선 토큰의 키·값 벡터를 모두 캐시해야 해서, 문맥이 길어지면 메모리 사용량이 그에 비례해 늘어난다.
• 선형 어텐션은 소프트맥스 정규화를 빼서 쿼리를 합 밖으로 끄집어낼 수 있게 만들고, 키와 값의 외적을 누적한 '상태 행렬' 하나만 들고 간다. 그 결과 메모리가 문맥 길이와 무관하게 일정해진다.
• 순차 재귀 구조는 GPU 학습에 불리하므로, 시퀀스를 청크로 잘라 청크 안쪽은 행렬 곱으로 병렬 처리하고 청크 사이에서만 상태를 넘기는 청크 단위 병렬 형태를 쓴다. 근사가 아니라 수학적으로 동일한 계산이다.
• 상태 행렬을 갱신하는 식은 키에서 값을 예측하는 선형 회귀를 확률적 경사하강법으로 한 스텝 학습하는 식과 정확히 일치한다. 즉 추론 도중에도 작은 회귀 모델을 계속 학습시키는 셈이다.
• 단순 누적만으로는 오래된 연관을 지우지 못해 언어 모델 성능이 뒤처진다. 감쇠 게이트와 델타 규칙(옛 값을 지우고 새 값을 써넣는 갱신)이 이 격차를 메우는 핵심 장치다.

쉽게 이해하기

어텐션은 생성형 AI의 핵심 부품이지만, 문맥 길이가 늘어날수록 계산량은 제곱으로, 메모리는 선형으로 불어난다는 제약을 안고 있다. 영상은 먼저 표준 어텐션을 복습한다. 각 토큰 임베딩을 쿼리·키·값 벡터로 사영하고, 쿼리와 키의 내적으로 관련도를 잰 뒤 소프트맥스로 정규화해 값 벡터를 가중합한다. 자기회귀 생성에서는 각 토큰이 자기 자신과 앞선 토큰만 보도록 마스킹한다.

응답을 생성하기 시작하면 흥미로운 지점이 드러난다. 새 토큰을 뽑을 때 앞선 토큰들의 키·값 벡터는 전혀 바뀌지 않는다. 그래서 이를 메모리에 저장해 재사용하는 KV 캐싱이 표준이 됐다. 쿼리는 마스킹으로 어차피 영향이 없어 캐시할 필요가 없다. 문제는 이 캐시가 문맥 길이에 비례해 커진다는 점이고, 그룹 쿼리 어텐션 같은 개선책도 이 선형 증가 자체를 없애지는 못한다.

선형 어텐션은 발상이 단순하다. 소프트맥스를 그냥 빼는 것이다. 정규화가 사라지면 각 항은 쿼리·키·값 세 벡터의 곱이 되고, 내적의 선형성 덕분에 공통 인자인 쿼리를 합 바깥으로 뺄 수 있다. 남은 합, 즉 키와 값의 외적을 차곡차곡 더한 행렬이 '상태 행렬'이다. 상태 행렬은 재귀식으로 갱신되므로 은닉 상태가 벡터가 아니라 행렬인 순환 신경망(선형 RNN)처럼 동작하고, 저장할 것은 이 행렬 하나뿐이라 메모리가 상수로 고정된다.

대신 학습이 까다로워진다. 재귀 구조는 순차적이라 병렬화가 어렵고, 외적은 원소별 연산이라 GPU의 텐서 코어를 살리지 못하며, 매 스텝 상태 행렬을 실체화하면 입출력 부담이 크다. 해법은 절충안인 청크 단위 병렬 형태다. 시퀀스를 겹치지 않는 크기 C의 청크로 나눈 뒤, 청크 안의 외적 합은 행렬 곱 한 번으로 처리하고 청크 경계에서만 상태를 순차로 넘긴다. 출력도 이전 청크에서 넘어온 기여와 현재 청크 내부 기여로 쪼개 병렬 계산한다. 이 방식은 근사가 아니며, 오늘날 효율적 시퀀스 모델 학습의 표준 레시피가 됐다.

그럼에도 선형 어텐션은 언어 모델링에서 소프트맥스 어텐션에 밀린다. 한 가지 원인으로 지목되는 것이 최근성 편향의 부재다. 상태 행렬이 덧셈으로만 쌓이다 보니 최근 토큰을 더 중시하기 어렵다. 그래서 0과 1 사이 감쇠 계수로 과거 기여를 줄이는 게이팅이 도입됐고, 값을 고정하는 대신 임베딩에서 계산하는 데이터 의존 게이트(맘바 2 등), 나아가 차원별로 망각 속도를 조절하는 벡터 게이트로 발전했다.

영상 후반부는 이 갱신 규칙을 '테스트 타임 회귀'로 다시 읽는다. 소프트맥스 어텐션이 키-값 쌍을 통째로 저장하는 연상 기억이라면, 선형 어텐션은 키를 넣으면 값이 나오도록 하는 회귀 함수에 문맥을 압축해 넣는 신경 기억에 가깝다. 손실을 제곱 오차로 두고 경사하강 한 스텝을 밟으면 그 유명한 델타 규칙이 나오는데, 이는 현재 키에 얽힌 옛 값을 지우고 새 값을 써넣는 동작으로 해석된다. 여기서 갱신 스텝을 여러 번 밟거나, 현재 토큰뿐 아니라 이전 토큰들의 예측 오차까지 함께 보거나, 회귀 함수 자체를 비선형으로 키우는 방향(테스트 타임 트레이닝)으로 연구가 뻗어 나가고 있다.

주요 인사이트

• 선형 어텐션의 이점은 '빠르다'보다 '메모리가 문맥 길이와 무관하다'는 쪽이 본질이다. KV 캐시가 상태 행렬 하나로 압축되므로, 긴 문맥에서 추론 지연이 일정하게 유지된다.
• 메모리를 상수로 줄인 대가는 정보 손실이다. 무한한 문맥을 고정 크기 행렬에 욱여넣는 이상, 무엇을 지우고 무엇을 남길지 정하는 장치(게이팅·델타 규칙)가 성능을 좌우한다.
• 상태 갱신이 경사하강 한 스텝과 같다는 관측은 단순한 비유가 아니라 식이 그대로 일치하는 결과다. 이 관점은 모델 파라미터를 학습 때 천천히 갱신되는 '느린 가중치', 문맥에 맞춰 추론 중 빠르게 바뀌는 상태를 '빠른 가중치'로 나눠 보는 오래된 아이디어와 이어진다.
• 청크 단위 병렬 형태는 알고리즘 설계가 하드웨어에 종속된다는 점을 보여 준다. 텐서 코어를 쓰려면 계산을 행렬 곱 모양으로 재배열해야 하고, 그 제약이 곧 모델 구조 선택을 좌우한다.
• 복잡한 아키텍처를 새로 설계하는 대신 트랜스포머 파라미터 일부를 추론 시점에 다음 토큰 예측으로 갱신하는 접근도 나왔다. 문맥 길이에 대한 확장성은 전체 어텐션에 필적하면서 추론 지연은 일정하게 유지된다는 결과다.

자주 묻는 질문

관련 AI 소식