🤖 machine learning

Looped Latent Attention: Cross-Loop KV Compression for Looped Transformers

이 논문은 가중치 공유 트랜스포머(weight-tied Transformers) 내 루프 인덱싱된 KV 캐시의 저계수(low-rank) 구조를 활용하여, SVD 초기화 및 증류(distillation)를 통해 근사 무손실 성능을 유지하면서도 극단적인 압축(최대 32배)을 달성하고 배치 용량을 크게 증가시키는 사후 학습 코덱인 Looped Latent Attention (LLA)을 소개한다.

원저자: James O' Neill, Fergal Reid

게시일 2026-07-20

📖 5 분 읽기🧠 심층 분석

CC BY 4.0

원저자: James O' Neill, Fergal Reid

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

AI 두뇌의 메모리 문제

로봇에게 이야기를 쓰는 법을 가르치고 있다고 상상해 보세요. 로봇은 단어들이 어떻게 연결되는지 아는 수백만 개의 작은 스위치(파라미터)로 이루어진 두뇌를 가지고 있습니다. 하지만 문제는, 긴 이야기를 쓰기 위해 로봇은 단순히 자신의 두뇌만 들여다보는 것이 아니라 '연습장'이 필요하다는 점입니다. 새로운 단어를 쓸 때마다, 로봇은 자신이 방금 무엇을 썼는지 기억하여 내용을 반복하거나 줄거리를 놓치지 않아야 합니다. 인공지능의 세계에서 이 연습장은 **KV 캐시(Key-Value cache)**라고 불립니다. 이것은 모든 대화마다 생성되는 개인용 공책이며, 로봇이 말을 할 때마다 점점 더 커집니다.

이제, 매우 효율적으로 설계된 특별한 종류의 로봇을 상상해 보세요. 이 로봇은 다양한 작업을 수행하기 위해 많은 서로 다른 방(레이어)을 가진 거대한 두뇌를 갖는 대신, 반복해서 방문할 수 있는 작고 영리한 방 하나를 가지고 있습니다. 이 로봇은 더 깊게 생각하기 위해 동일한 일련의 지침을 반복해서 사용하며 루프를 돕니다. 이를 **루프드 트랜스포머(Looped Transformer)**라고 합니다. 이는 두뇌 자체의 공간을 절약해주므로 매우 좋습니다. 하지만 여기에는 교활한 문제가 있습니다. 로봇이 동일한 방을 재사용하더라도, 루프를 돌 때마다 연습장에 새로운 페이지를 써야 한다는 점입니다. 만약 로봇이 수학 문제를 풀기 위해 네 번 루프를 돈다면, 이야기의 같은 순간에 대해 네 개의 별도 메모 페이지를 갖게 됩니다. 즉, "루프형" 로봇은 여전히 실행을 위해 엄청난 양의 메모리가 필요하며, 이는 효율성을 추구하는 목적을 무색하게 만듭니다. 과학자들의 큰 질문은 이것입니다: "우리가 로봇이 생각하는 법을 잊어버리게 만들지 않으면서, 저 연습장을 줄일 수 있을까?"

논문의 발견: 메모리 속의 비밀 지름길

이 논문은 그 메모리 문제를 해결하기 위해 **루프드 레이턴트 어텐션(Looped Latent Attention, LLA)**이라는 영리한 기술을 소개합니다. 연구진은 로봇의 메모리가 보기만큼 엉망이 아니라는 사실을 발견했습니다. 로봇이 사고 과정을 루프 through 할 때, 동일한 단어에 대해 작성하는 메모는 급격하게 변하지 않습니다. 대신 그것은 완만한 언덕을 굴러 내려가는 구슬처럼 매끄럽고 예측 가능한 경로를 따릅니다. 연구진은 로봇의 연습장 모든 페이지를 전부 저장하는 대신, 그 경로의 아주 작은 "요약본"과 로봇이 실제로 읽어야 할 때 특정 페이지를 다시 구축하는 방법(지침)만을 저장합니다.

이렇게 생각해 보세요. 주인공이 복도를 걸어가는 영화를 보고 있다고 가정해 봅시다. 일반적인 컴퓨터라면 주인공이 발걸음을 옮길 때마다 고화질 사진 한 장을 모두 저장할 것입니다. 이는 엄청난 저장 공간을 차지합니다. 하지만 LLA를 사용하면, 주인공이 그저 직선으로 걷고 있다는 것을 깨닫게 됩니다. 그래서 수천 장의 사진을 저장하는 대신, 주인공의 사진 한 장과 "2단계에서는 1인치 앞으로, 3단계에서는 2인치 앞으로 이동"이라는 간단한 메모를 저장합니다. 컴퓨터가 3단계에서 주인공을 봐야 할 때, 이 메모를 사용하여 빠르게 그림을 그려냅니다. 이것이 바로 이 논문이 AI를 위해 하는 방식입니다. 즉, "루프" 부분의 메모를 압축하여, 네 페이지의 뭉치를 단 하나의 작은 요약본으로 변환한 뒤, 필요할 때 즉시 전체 페이지로 확장할 수 있게 만드는 것입니다.

연구진은 이 방법을 Ouro-1.4B와 같은 여러 AI 모델에 테스트했습니다. 그 결과, 이 방법이 놀라울 정도로 잘 작동한다는 것을 발견했습니다. 실제로, H200이라는 강력한 컴퓨터 칩에서 이들은 이전에는 32개의 대화만 수용할 수 있었던 것에 비해, 한 번에 768개의 서로 다른 대화를 메모리에 담을 수 있었습니다. 이는 동시에 AI를 사용할 수 있는 사람의 수가 비약적으로 늘어났음을 의미합니다. 또한, 이 기술이 매우 길고 어려운 수학 문제를 풀 때도 작동함을 보여주었지만, 다만 로봇이 선생님을 따라 하는 것이 아니라 자신의 이야기를 직접 쓸 때 혼란을 겪지 않도록 약간의 추가 훈련이 필요합니다.

이 논문이 "아니오"라고 말하는 것들

이 방법이 아닌 것이 무엇인지 아는 것도 중요합니다. 연구진은 먼저 매우 단순한 아이디어를 테스트했습니다: "만약 우리가 처음 세 페이지의 메모를 버리고 마지막 한 페이지만 남겨둔다면 어떻게 될까?" 그들은 로봇이 몇 번의 루프 후에 최종 답안에 안착하므로 초기 메모는 중요하지 않을 것이라고 생각했습니다. 그들은 이를 시도해 보았으나, 결과는 참담했습니다. AI는 수학 문제를 푸는 법을 완전히 잊어버려 테스트에서 0점을 기록했습니다. 이는 메모가 지나가는 "경로"가 중요하다는 것을 증명합니다. 즉, 단순히 도착 지점만 가져와서는 안 됩니다. 그 여정을 포착해야 합니다.

또한 그들은 메모를 압축(양자화, quantization)하거나 뇌의 서로 다른 부분 간에 메모를 공유하는 방식(헤드 축 압축, head-axis compression)과 같은 다른 메모 축소 방식과 비교했습니다. 이러한 방식들도 약간의 도움은 되지만, 논문은 "루프" 축을 압축하는 것이 핵심 비결임을 보여줍니다. 만약 루프 대신 다른 부분을 압축하려고 하면, AI는 추론 능력을 상실합니다. 논문은 명확히 밝히고 있습니다: 루프가 메모에서 가장 중복되는 부분이며, 압축해야 할 곳은 오직 그곳뿐이라는 점입니다.

얼마나 확신하는가?

저자들은 직접 측정한 결과를 바탕으로 이 결과에 대해 상당히 자신감을 갖고 있습니다. 그들은 단순히 추측한 것이 아니라, GSM8K(수학 문제) 및 HumanEval(코드 작성)과 같은 실제 작업에 AI를 실행했습니다. 그들은 4배 압축(메모리를 4배 작게 만듦)을 적용했을 때도 AI가 원래의 압축되지 않은 버전과 거의 비슷하게 성능을 낸다는 것을 보여주었습니다. 극한의 압축(21.3배)까지 밀어붙였을 때도 AI는 많은 작업을 처리할 수 있었지만, 매우 길고 복적인 추론 단계에서는 어려움을 겪기 시작했습니다.

그들은 또한 32번까지 루프를 도는 다른 구조의 AI 모델인 Huginn-3.5B에서도 이 실험을 진행했으며, 그곳에서도 이 방법이 작동하여 이것이 특정 로봇만의 우연한 현상이 아님을 입증했습니다. 논문은 이 "저계수 궤적(low-rank trajectory, 메모리의 매끄러운 경로)"이 루프형 로봇이 생각하는 방식의 근본적인 특성임을 시사합니다. 이들은 메모리 절약과 용량 증대 측면에서 효과적임을 입증했지만, 매우 긴 작업에서 최상의 성능을 내기 위해서는 로봇이 단순히 선생님을 모방하는 것이 아니라 자신의 압축된 메모리를 가지고 연습하는 '온 폴리시(on-policy)' 훈련 단계가 필요하다고 언급했습니다. 이 훈련은 로봇이 자신의 이야기를 직접 쓸 때 발생하는 작은 오류들을 수정해 줍니다.

요약하자면, 이 논문은 루프형 AI 모델의 메모리가 특정한 방식으로 반복된다는 점을 이용해 훨씬 더 효율적으로 만들 수 있음을 보여줍니다. 반복되는 부분의 전체 메모리 대신 아주 작은 요약본을 저장함으로써, 우리는 값비싼 하드웨어 없이도 더 많은 사람이 강력한 AI를 사용할 수 있도록 단일 컴퓨터 칩에 수백 배 더 많은 대화를 담을 수 있습니다.

기술 요약: 루프드 레이턴트 어텐션 (Looped Latent Attention, LLA)

문제 정의

**루프형 가중치 공유 트랜스포머(Looped, weight-tied Transformers)**에서, 단일 블록을 반복적으로 적용(재귀)하는 방식은 파라미터 수를 줄이기 위해 사용됩니다. 이는 모델의 가중치를 효과적으로 압축하지만, 추론 메모리 병목 현상은 여전히 Key/Value (KV) 캐시에 남아 있습니다. 표준적인 루프 추론에서는 매 재귀 단계마다 어텐션을 다시 계산하며, 각 단계, 레이어, 헤드별로 별도의 KV 캐시를 저장합니다. 결과적으로, $T$ 번의 루프와 $D$ 개의 레이어를 가진 모델은 $T \times D$ 레이어 디코더와 동일한 KV 캐시를 보유하게 되어, 가중치 공유를 통해 얻은 메모리 이점을 상쇄하게 됩니다.

본 논문은 이 루프 인덱스 기반의 캐시가 매우 구조화되어 있다고 주장합니다. 특정 토큰, 레이어, 헤드에 대해, 재귀 단계에 따른 K/V 벡터들은 고차원 공간을 탐색하기보다는 **짧고 낮은 랭크(low-rank)를 가진 궤적(trajectory)**을 그립니다. 이러한 중복성은 **루프 축(loop axis)**에 국한된 것이며, 헤드 및 레이어 축의 스펙트럼은 여전히 상당히 평탄한(높은 랭크를 가진) 상태를 유지합니다. 과제는 이 궤적을 하나의 고정된 점으로 붕괴시켜 성능을 저하시키지 않으면서, 재귀 축을 압축하는 것입니다.

방법론: 루프드 레이턴트 어텐션 (LLA)

저자들은 $T$ 개의 루프당 K/V 벡터를 하나의 압축된 레이턴트 표현과 루프별 재구성 맵으로 대체하는 포스트 트레이닝 캐시 코덱인 **루프드 레이턴트 어텐션 (LLA)**을 제안합니다.

핵심 메커니즘

레이턴트 저장 (Latent Storage): 특정 토큰, 레이어, 헤드에 대해 $T$ $T$ 개의 벡터( $x_1, \dots, x_T$ $x_{1}, \dots, x_{T}$ )를 저장하는 대신, LLA는 하나의 중심화된 저랭크 레이턴트 벡터 $c \in \mathbb{R}^r$ $c \in R^{r}$ 를 저장합니다.
- $c = (x - \mu)W_{\text{down}}$
- 여기서 $\mu$ 는 캘리브레이션 세트의 평균이며, $W_{\text{down}}$ 은 다운 프로젝션 행렬입니다.
루프별 재구성 (Loop-Specific Reconstruction): 재귀 단계 $t$ $t$ 에 대한 특정 K/V 벡터는 루프별 업 프로젝션을 사용하여 재구성됩니다.
- $\hat{x}_t = c W_{\text{up},t}^\top + \mu_t$
비대칭 랭크 (Asymmetric Ranks): 이 방법은 Key( $r_k$ )와 Value( $r_v$ )에 대해 서로 다른 랭크를 사용합니다. 경험적 분석 결과, Value 벡터가 Key 벡터보다 루프를 거치며 더 느리게 안정화된다는 것을 확인하였으며, 이에 따라 $r_v > r_k$ 로 설정하는 것이 타당함을 입증했습니다.
LLA-2D 변형: 극단적인 압축을 위해, LLA-2D는 **헤드 축(head axis)**을 레이턴트에 통합하여, 모든 $H$ 개의 헤드를 재구성하기 위해 토큰 및 레이어당 하나의 레이턴트를 저장합니다. 이는 캐시 크기를 $H$ 배만큼 줄이지만 헤드 간의 결합을 초래합니다.

훈련 및 초기화

SVD 초기화: 코덱은 캘리브레이션 세트의 티처(teacher) 활성화 값에 대한 특이값 분해(SVD)를 통해 초기화됩니다. 다운 프로젝션은 상위 $r$ 개의 우특이 벡터(right singular vectors)로 설정됩니다. 이 초기화는 매우 중요하며, 무작위 초기화는 훨씬 높은 KL 발산을 초록합니다.
증류 (Distillation): 코덱은 티처 모델을 고정하고, 티처의 분포와 교체된 모델의 분포 사이의 순방향 KL 발산(forward KL divergence)을 최소화함으로써 정교화되며, 여기에 어텐션 출력 매칭 항이 추가됩니다.
온-폴리시 정교화 (On-Policy Refinement): 긴 자기회귀적 롤아웃(모델이 티처의 접두사가 아닌 자신의 생성된 접두사에 조건화되는 경우)에서의 노출 편향(exposure bias) 문제를 해결하기 위해, 두 번째 단계로 온-폴리시 증류가 수행됩니다. 코덱은 학생(student) 모델 자체에서 샘플링된 접두사를 사용하여 업데이트되어, 긴 생성 성능을 안정화합니다.

주요 기여

LLA의 정의: 가중치가 공유된 트랜스포머를 위한 교차 루프 레이턴트 K/V 코덱을 정의하며, 압축률 $\rho = \frac{2Td_{\text{head}}}{r_k + r_v}$ 를 달성합니다. 이는 토큰당 레이어당 $2Td_{\text{head}}$ 대신 $(r_k + r_v)H$ 개의 스칼라를 저장합니다.
축 특화 압축 (Axis-Specific Compression): 본 논문은 루프형 모델(Ouro-1.4B, Ouro-2.6B-Thinking, Huginn-3.5B)에서 재귀(루프) 축이 가장 압축 가능한 차원임을 입증합니다. 루프 축을 압축하는 것이 헤드 축 압축(MLA), 레이어 축 압축(교차 레이어 공유), 또는 동일한 캐시 예산에서의 양자화보다 뛰어난 성능을 보입니다.
고정점 붕괴 거부 (Rejection of Fixed-Point Collapse): 저자들은 루프 궤적이 단일 상태(최종 루프)로 대체될 수 없음을 보여줍니다. 최종 루프의 K/V만을 재사용하는 제로 파라미터 컨트롤은 생성 성능을 붕괴시키며(GSM8K 정확도가 0.000으로 하락), 이는 "짧은 저랭크 궤적"이 필수적임을 증명합니다.
배포 영향: 정확한 스칼라 감소는 직접적인 서빙 용량 증가로 이어집니다. 단일 H200 GPU에서, LLA는 4k 컨텍스트 기준 Ouro-1.4B의 최대 배치 크기를 32에서 768 시퀀스로 증가시킵니다 (21.3배 압축 시).

실험 결과

압축 대 정확도

GSM8K 성능: Ouro-1.4B에서 가벼운 압축(1.33× ~ 2×)은 티처의 정확도(약 0.79–0.80)와 일치합니다. 4× 압축 시에도 LLA는 0.59의 정확도를 유지하며, 이는 헤드 축 압축(0.52) 및 최종 루프 재사용(0.00)보다 유의미하게 높습니다.
규모 일반화: Ouro-2.6B-Thinking에서 LLA는 4×에서 10.67× 압축 구간에서도 거의 손실 없는 성능(0.81–0.83)을 유지하는 반면, 헤드 축 베이스라인들은 상당한 정확도 저하를 보입니다.
Huginn으로의 전이: 이 방법은 Huginn-3.5B(32 재귀 단계)로 전이됩니다. 훈련이 필요 없는 SVD 코덱은 디코더 독립적 지표에서 32× 압축까지 거의 손실 없는 충실도를 달성합니다.

긴 생성 안정성

MATH-500: 긴 수학적 롤아웃은 접두사 드리프트(prefix drift)에 민-감합니다. 오프-폴리시 훈련(티처 접두사 사용)은 생성 길이가 길어짐에 따라 정확도 저하와 높은 "무응답(no-answer)" 비율을 초래합니다.
온-폴리시 이득: 온-폴리시 정교화를 적용하면 4× 압축 시 MATH-500 정확도가 0.43에서 0.66으로 상승하며, 무응답 생성을 크게 줄여줍니다.

서빙 용량

메모리 감소: 캐시 감소는 정확합니다. 21.3× 압축 시, 1M 컨텍스트에서의 Ouro-1.4B KV 캐시는 768 GiB에서 36 GiB로 감소합니다.
처리량 (Throughput): "레이턴트 저장 경로"(레이턴트만 캐싱)는 연산 집약적(compute-bound)이지만, 이전에 OOM(메모리 부족)으로 불가능했던 배치와 컨텍스트를 가능하게 합니다. "패스트 패스"(K/V를 재구성하여 캐싱)는 처리량을 보존하지만 메모리 절감 효과는 실현하지 못합니다.

의의 및 주장

본 논문은 재귀가 가중치 공유형 루프 트랜스포머에서 가장 압축 가능한 캐시 축이라고 주장합니다. 그 의의는 다음과 같습니다:

기계론적 통찰: 루프 축이 헤드 및 레이어 축보다 스펙트럼 측면에서 구별됨(낮은 랭크)을 식별하여, 모델의 동적 동작을 정적인 고정점으로 붕괴시키지 않고도 효율적인 압축을 가능하게 함을 보여주었습니다.
실질적 배포: 기본 모델 아키텍처를 재학습할 필요 없이 기존 하드웨어에서 배치 크기와 컨텍스트 길이를 직접적으로 늘릴 수 있는 포스트 트레이닝 솔루션을 제공합니다.
미래 방향: 이 결과는 네이티브 재귀 아키텍처가 전체 K/V 상태를 다시 생성하는 대신, 저렴한 루프별 델타(delta)와 공유된 콘텐츠를 계산하는 분리된 레이턴트 캐시의 이점을 얻을 수 있음을 시사합니다.

저자들은 본 방법의 정확도 주장이 흡수된 어텐션(absorbed attention) 구현에 의존하지 않는다는 점을 강조합니다. 핵심 결과는 저랭크 궤적 가정의 타당성과 엄격한 메모리 제약 하에서 티처의 동작을 보존하는 코덱의 효능입니다.