Topic #31: 추론 (Reasoning) 모델의 엄청난 메모리 사용량, 어떻게 줄일 수 있을까?

1. 요약을 하기 전에, LightThinker가 텍스트를 더 작은 섹션으로 나눕니다. ‘요약’한 내용을 표시하기 위해서 섹션들 사이에 특별한 토큰을 삽입합니다:

• (Optional) → 모델에게 ‘이전까지의 생각을 압축하라’고 알려주는 마커 (Marker)

• C (캐시 토큰) → 핵심 포인트를 저장하는, ‘이전 생각’의 요약 토큰

• [o] (출력 토큰) → "이 요약된 내용을 활용해서 추론을 계속하라”고 지시하는 마커

결국, 세 개의 텍스트 섹션 시퀀스가 있다면 이렇게 보일 수 있겠죠: → 섹션 1 → → 요약 (C) → [o] → 섹션 2 → → 요약 (C) → [o] → 섹션 3. 여기까지의 단계를 "데이터 재구성"이라고 합니다. 그 다음으로는:

2. 압축 방법과 요약 내용을 사용하는 법 배우기: 데이터가 구조화되면, 모델은 언제 그리고 어떻게 정보를 압축할지 배웁니다. 사고 기반 (Thought-based)의 어텐션 마스크는 각 단계에서 모델이 ‘볼 수 있는’ 것과 ‘볼 수 없는’ 것을 제어합니다:

• 압축 중에, 모델은 원래 입력, 이전에 압축된 콘텐츠 (C) 및 현재 생각만 처리할 수 있습니다.

• 출력 토큰을 생성하는 동안, 모델은 입력 질문과 이전에 압축된 콘텐츠만 ‘볼 수 있습니다’. 이렇게 해서 추론할 때 모델이 원시 데이터 대신 요약된 내용만 보도록 보장을 하는 거죠. 아래 그림이 세 단계의 추론 작업이 진행되는 동안 LightThinker의 어텐션 마스크 도식입니다:

3. 마지막으로, 모델이 오로지 요약 내용(C)만 사용해서 다음 토큰을 예측하도록 훈련을 합니다. 전체 텍스트를 보면서 ‘치팅하는’ 걸 허용하지 않고, 정보를 효율적으로 저장하고 단계별로 검색하는 방법을 학습하게 됩니다.

정리하면, LightThinker는 긴 생각을 짧은 메모리 토큰으로 압축하고, ‘중요한’ 세부 사항만 가지고 추론을 계속하는 구조입니다. 잊지 말아야 할 특징은, 정확도화 효율성을 높이기 위해서 ‘메모리의 압축’과 ‘추론 과정’을 분리하도록 설계했다는 것입니다.

그럼, LightThinker가 실제로 어느 정도 성능을 보여주는지 확인을 해 봅시다.

LightThinker의 실제 성능

앞서 이야기한 것처럼, LightThinker는 AI 모델이 실제의 추론 작업에 잘 활용할 수 있는 실용적이고 효율적인 기법인데, 인상적인 성능을 보여줍니다:

• 메모리 사용량의 감소
  LightThinker가 최대 토큰 사용량을 70% 정도 줄여줘서, 불필요한 정보를 훨씬 적게 저장해서 메모리 사용량을 절감해 줍니다.

• 더 빠른 처리 속도
  Qwen2.5-7B에서는 추론 시간을 26%, Llama3.1-8B에서는 41% 단축해 주는 걸로 나타납니다. 물론 추론 시간이 감소되는 정도는 응답의 길이에도 의존을 합니다. 예를 들어서, LightThinker가 긴 응답(32K 토큰)을 생성할 때는 추론 시간을 44% 줄여주는 정도로 효과가 크고 더 짧은 텍스트(1K-4K 토큰)의 경우에도 역시 추론 시간의 1%-4% 정도는 절약을 해 줍니다.

• 최소한의 정확도 손실
  정확도는 Qwen에서 1%, Llama에서 6% 정도만 감소하는데, 효율성이 향상되는 정도를 생각하면 합리적인 수준이라고 할 수 있습니다.

• 토큰 생성량의 감소
  LightThinker는 모델이 생성하는 토큰 수를 Qwen에서 15%, Llama에서 13% 정도까지 줄여 줍니다.

• 적응형 (Adaptive) 압축 성능
  수학 문제 같은 간단한 작업에서는 더 공격적으로 압축하고, GPQA (대학원 수준의 증명 Q&A 벤치마크) 같은 복잡한 작업에서는 중요한 정보를 유지하기 위해더 더 신중하게 압축합니다.

• 정확도의 향상
  압축과 추론 단계를 분리한 LightThinker의 설계 구조 덕분에 정확도가 2% 향상되고, 어텐션 마스크 전략과 결합하면 성능이 7% 더 향상됩니다.

• 적절한 캐시 크기의 중요성
  캐시 크기(LightThinker가 메모리에 저장하는 토큰 수)가 작으면 더 자주 압축을 하게 되는 반면, 캐시 크기가 크면 아무래도 압축을 덜 하지만 더 많은 정보를 유지합니다. 따라서 캐시 크기를 늘리면 정확도가 향상되고 추론 시간이 줄어듭니다.

• 수학 문제에는 어려움을 겪는 LightThinker
  수학 문제는 보통 LightThinker가 잘 다루지 못하는 경우가 많은데, 이건 숫자값이 잘못 압축되는 경우가 많기 때문입니다.

LightThinker가 좋은 기법인 이유, 그리고 한계

LightThinker 기법 그 자체처럼, 지금까지 알아본 LightThinker의 좋은 점을 요약해 볼까요?

• 효율성을 높이기 위해서, ‘생각하기’와 ‘요약하기’를 분리한 설계 구조입니다.

• 과거 생각들을 더 잘 추적하도록 해 줍니다.

• 정확성을 유지하면서도 메모리를 절약해 주는 장점이 있습니다.

• 일반적으로 40% 이상 더 빠른 속도로 실행됩니다.

반면에, LightThinker는 아래와 같은 한계점도 가지고 있습니다:

• 수학 과제를 해결하는데는 어려움을 겪고, 큰 도움이 되지 않을 수 있습니다.

• Llama 모델의 경우에는 크게 좋은 기법이 아닐 수 있습니다.

• 다이나믹한 압축 과정 때문에, 가끔 메모리 사용량에 피크를 칠 때가 있습니다.

• 훈련을 위해서 고정된 수의 캐시 토큰을 사용하지만, 실제 작업에서 다양한 토큰 관련 요구사항에 적응이 가능한지 불확실성이 있습니다.

• 훨씬 더 큰 데이터셋으로 LightThinker를 훈련한다고 해도 더 좋아질지는 불분명합니다.

• LightThinker는 리소스를 많이 사용하는 ‘풀 파라미터 (Full Parameter)’ 파인튜닝 기법으로 훈련한 기법으로, 더 적은 수의 파라미터를 사용하는 LoRA라든가 QLoRA 같은, 더 효율적인 튜닝 환경에서는 테스트되지 않았습니다.

자, 그럼 이제 완전하게 다른 방식으로 작동하지만 역시 메모리의 사용을 효과적으로 줄여주는, DeepSeek의 기법에 대해서 알아볼까요?

MLA (Multi-Head Latent Attention)는 무엇인가?

MLA가 왜 필요했을까?

DeepSeek-R1 추론 모델, 엄청난 화제를 불러일으켰죠. DeepSeek의 이 추론 모델은, DeepSeek-V2를 개발하면서 제안한 ‘두 가지의 기술적 혁신’을 활용하는데요 - DeepSeek에 대한 AI 101 에피소드에서도 간략하게 다룬 적이 있어요:

이 두 가지의 기술적 혁신이라는 건:

• 여러 전문가 하위 모델(각 전문가는 또 더 작고 전문화된 부분으로 나뉨)과 항상 활성 상태로 유지되며 공통 지식을 처리하는 공유 전문가로 구성된 특별한 Mixture-of-Experts 시스템

• Multi-Head Latent Attention (MLA)라고 부르는, 전문화된 어텐션 메커니즘.

입니다.

DeepSeek은, 스스로 ‘최고의 AI 모델’을 만드는 과정에서 왜 기존의 Multi-Head Attention (MHA)를 수정해야 한다고 생각했을까요?

트랜스포머 아키텍처에서, MHA 메커니즘은 텍스트를 처리하고 생성하기 위해 입력값의 ‘가장 관련성 높은 부분’에 집중 (어텐션)하는 데 도움을 줍니다. 그렇지만, MHA를 사용하기 때문에 추론 과정 중에 엄청나게 많은 KV 값의 쌍 (Pair)을 저장하게 되어서, 엄청난 양의 메모리를 차지하고 모델의 작업을 느리게 만들기도 하죠.

여기서 MLA가 등장하는데요. DeepSeek의 Multi-Head Latent Attention (MLA)은 KV 캐시를 훨씬 작은 형태로 압축하는, 수정된 어텐션 메커니즘입니다. MLA는 과거의 정보를 효율적으로 압축하면서도 나중에 사용하기 위해서 접근이 용이하도록 유지하는, 일종의 스마트 스토리지 시스템 같은 걸로 이해하면 되겠습니다. 이 작업이 Low-Rank Key-Valye Joint Compression이라는 기술을 활용해서 진행되는데, 이 기술은 ‘모델이 높은 정확도를 손해보지 않으면서도 더 빠르고 적은 메모리로 정보를 처리할 수 있게’ 해 줍니다. 그렇다면, MLA가 어떻게 정확히 KV 캐시를 압축할까요?

MLA의 작동 방식

MLA는 KV의 저장 공간을 줄이면서 Low-Rank Key-Value Joint Compression 기술을 통해서 모델의 강력한 성능은 유지합니다:

1. Key-Value Pair의 압축
   MLA는 각 토큰마다 모든 전체 크기의 KV 쌍을 따로 저장하는 대신, 저장하기 전에 더 작은 저차원의 표현으로 압축합니다. 즉, MLA가 수학적 변환으로 KV 쌍을 더 작은 잠재 공간으로 투영 (Projection)한다는 뜻입니다.

2. 어텐션을 계산하기 위한 압축 해제
   모델이 저장된 KV 쌍을 사용해야 할 때는, 압축된 데이터를 원래 크기로 다시 확장해서 재구성합니다. 압축된 KV 데이터는 추론 과정에서 다시 확장되니까, 모델은 여전히 전체 크기의 KV 쌍을 가지고 있는 것처럼 작동할 수 있습니다. 이렇게 하면, 모델이 이전 토큰에서 중요한 정보를 계속 얻으면서도 메모리 사용량을 줄여서 추론 속도를 높일 수 있습니다.

3. ‘분리된’ 회전 위치 임베딩 (Decoupled Rotary Position Embedding, RoPE)
   RoPE는 단어의 ‘순서’를 기억할 수 있도록 K와 Q를 수정해서 위치 데이터를 포함시키는데요. 하지만 KV 쌍을 압축할 때 위치 정보도 압축된 형태로 얽혀서, 나중에 올바르게 재구성하기가 어려워지죠. MLA는 이 문제도 해결해 주는데, 위치 정보를 주요 KV 압축에서 분리하고, 추가적인 멀티헤드 쿼리와 공유 키를 사용해서 위치 정보를 별도로 처리합니다. 이런 접근 방법은 다음과 같은 이점이 있습니다:

   • 위치 인코딩이 KV의 압축 과정을 방해하지 않도록 방지합니다.

   • 모델이 위치 조정을 위해서 모든 과거 키를 다시 계산할 필요가 없습니다.

   • 모델이 추가적인 처리 없이 과거의 토큰 정보를 빠르게 검색할 수 있습니다.

MLA 기법의 장점과 한계

MLA 기법은 아래와 같은 장점이 있습니다:

• 키와 값을 저차원으로 압축하니까 메모리 사용량을 크게 줄여줍니다.

• 텍스트의 생성 속도를 높여 줍니다: KV의 저장 공간이 적어서 검색과 추론이 더 빨라집니다.

• 전통적인 MHA만큼 잘 작동합니다: 멀티 쿼리 어텐션(Multi-Query Attention, MQA)이나 그룹 쿼리 어텐션(Grouped-Query Attention, GQA) 같은 다른 기법들과 달리, MLA는 출력물의 품질을 낮추지 않습니다.

• 위치 인코딩을 잘 처리합니다: ‘Decoupled’ RoPE 기법으로 단어의 순서가 잘 보존되도록 보장합니다.

물론, 어떤 기법이든 완벽한 건 없으니, MLA에서도 다음과 같은 한계점이 발견되는 건 당연하겠죠:

• 압축으로 인한 일부 정보의 손실: 전혀 손실이 없을 수는 없죠? MLA 기법을 사용할 때, 장거리 의존성이 약간 약화될 수 있습니다.

• 압축과 압축 해제를 위한 추가 계산: 훈련 속도가 조금 느려지고, 추론에도 약간 영향을 줄 수 있습니다.

• RoPE 호환성 문제: MLA는 RoPE를 분리해야 해서 구현의 복잡성이 증가합니다.

• 압축률의 트레이드오프: 메모리를 얼마나 절약할지, 성능은 어느 정도 보존할지 사이의 균형을 맞추기 위해서는, 세심한 조정을 해야 합니다.

• 다양한 모델에 대한 벤치마킹 부족: MLA는 아직까지는 여러 다양한 AI 아키텍처에서 충분히 테스트된 기법은 아닙니다.

• 다른 어텐션 메커니즘과의 통합 문제: 커스텀 어텐션 방식을 사용하는 모델에 대해서는 추가적인 조정이 필요할 수 있습니다.

이러한 한계들이 있지만, MLA는 키-값(KV) 캐시 크기를 잠재 벡터로 압축해서 크게 줄임으로써 효율적인 추론을 보장하는 DeepSeek의 최고 혁신 기술 중의 하나입니다.

DeepSeek-R1을 통해서도 놀라운 성능 결과치를 보기는 했지만, DeepSeek의 다른 모델들을 통해서 한 번 MLA가 가져다주는 성과를 다시 한 번 확인해 보도록 하죠.

MLA의 성능

DeepSeek-V2 모델의 경우, MLA 기법을 통해서 메모리 요구사항을 93.3% 줄이는 놀라운 성과가 있었습니다. 그리고, 저장하고 접근해야 할 데이터가 적어져서 각각의 추론 단계가 더 빨라졌구요. KV 메모리 사용량을 줄이게 되면 더 많은 시퀀스를 한 번에 처리할 수 있게 되어서, MLA를 사용하지 않는 모델과 비교해서 생성 처리량이 5.76배 향상됩니다. 물론 MLA만은 아니고 DeepSeekMoE 아키텍처, 이 두 가지 DeepSeek의 혁신 기술이 성능 수치의 개선에 기여했지만, 메모리 사용량의 압축 전략은 MLA에서 나온 거라고 봐야겠죠.

또, DeepSeek-R1 모델을 통해서, MLA 같은 압축 기술을 사용해도, 느리게 진행하는, 단계적인 추론 능력이 저하되지 않는다는 것을 볼 수 있습니다. 오히려 이 모델은 o1-1217 같은 오픈AI의 고급 추론 모델에 필적하는 추론 성능을 보여주기까지 했습니다.

자, 그럼 여기서 한 가지 생각해 볼 만한 주제가 떠오릅니다.

DeepSeek-R1 같은 모델의 정확도는 유지하면서, 속도와 효율성까지 개선하고 메모리 사용량마저 더 줄이는 게 가능할까요?

LightThinker와 MLA 기법을 함께 적용하면 어떨까?

LightThinker와 MLA, 두 가지 기법 모두 ‘메모리의 사용량을 줄이고’, ‘더 추론을 빠르게’ 해서 LLM의 효율성을 극적으로 높이는 걸 목표로 하는 기술입니다. 그런데 잘 살펴보면, 이 두 가지 기술은, 모델이 작동하는 과정에서 서로 다른 측면을 대상으로 하고 있어요.

LightThinker는 ‘추론 과정’ 자체를 압축하도록 설계된 반면에, MLA는 KV 캐시를 압축해서 저장 공간을 줄이고 추론 속도를 높여서 어텐션의 메모리 사용량을 최적화하는 거죠. 간단히 말하자면, MLA는 메모리를 절약하기 위해서 저장된 데이터를 압축하고 압축 해제하는 거고, LightThinker는 AI 모델이 모든 세부 사항을 기억할 필요가 없이 긴 대화를 요약하는 겁니다.

만약, DeepSeek-R1 같은 모델이 MLA와 함께 LightThinker 같은 기술을 함께 연계해서 사용한다고 상상해 보세요. 이런 조합이 가능하다면, 더 효율적이면서도 강력한 추론 모델을 만들 수 있는 것 아닐까요?

• 먼저, LightThinker는 추론 단계를 압축해서 중복성을 제거하고 핵심 정보만 요약본에 저장합니다.

• 그런 다음 MLA는 모든 토큰을 같은 비중으로 취급하는 대신에 가장 관련성 높은 잠재적인 세부 사항을 우선시하면서, 압축된 요약본에 서로 다른 가중치를 가지고 어텐션을 줄 수 있습니다.

개념적으로는, MLA가 LightThinker의 압축된 ‘생각’을 회상하는 적응형 검색 메커니즘 역할을 할 수 있게 하는 거죠. MLA는 또 LightThinker가 압축한 단계를 검색하고 확장하는 방식을 개선해서, 간결함 vs. 추론의 깊이 간 균형을 맞춰줄 수 있을 겁니다.

게다가, LightThinker와 MLA는 서로의 한계를 보완해 줄 수 있습니다. 예를 들면:

• MLA는 LightThinker가 공격적으로 추론 단계를 요약하는 과정에서 손실된 핵심적인 세부 사항을 복구하거나 중요한 세부 사항의 손실을 방지하는 데 도움을 줄 수 있습니다.

• LightThinker는 수치 값이 잘못 압축될 수 있어서 수학 작업에 어려움을 겪는다고 했는데, MLA는 구조화된 검색을 우선시하기 때문에 압축된 생각에서 수치의 일관성을 강화할 수 있습니다.

• MLA가 LightThinker의 압축된 캐시에 선택적으로 어텐션을 주는 능력은, 실제 작업에서 압축 전략을 더 유연하게 작동하게 해 줄 수 있습니다.

• MLA는 장거리 의존성을 약간 약화시킬 수 있는데, 이 경우에 LightThinker의 단계적 생각 압축으로 MLA가 처리해야 할 초기의 정보 자체를 줄일 수 있습니다.

전반적으로, MLA가 LightThinker의 "메모"(압축된 "생각") 사이에서 어떻게 어텐션을 배분하고, 메모리를 절약하기 위해서 필요할 때만 압축하고 압축 해제하는 방법을 탐구할 수 있다면 얼마나 흥미로울까요?

이 컴비네이션에 관심을 가지는 연구팀이나 스타트업이 있으면 좋겠습니다 ^.^

맺으며

‘메모리 압축’, AI의 발전에 필수적인 최적화 기술입니다. 특히 추론 (Reasoning) 과정을 확장해야 하는, 그러면서 상당한 자원을 사용해야 하는 거대한 추론 모델의 경우에는 더욱 중요하죠.

이 글에서 알아본 LightThinker와 MLA는 효과적으로 메모리의 사용량을 줄이고 추론 처리 속도를 높여주는 것으로 나타납니다. 하지만, 여기서 멈출 수는 없습니다 - AI 모델이 더 높은 정확도로 더 많이 사용되기 위해서는 추론의 단계와 추론 시간을 계속해서 늘려야 할 것이고, 앞으로 o1이나 DeepSeek-R1 같은 모델을 뛰어넘는 모델이 계속 나와야 하기 때문이죠.

그런 관점에서, MLA + LightThinker 같은 ‘하이브리드 접근법’이 어쩌면 더 메모리를 효율적으로 사용하면서 추론도 더 빠르게 하는 중요한 방향이 될 수도 있다고 생각합니다. 개발자나 연구자 여러분은 어떻게 생각하실지 궁금합니다.

보너스: 참고자료

1. LightThinker: Thinking Step-by-Step Compression by Jintian Zhang, Shuofei Qiao, Huajun Chen, Ningyu Zhang et al.

2. DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model by Bo Liu, Damai Dai et al.

3. DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning by Daya Guo, Qihao Zhu et al.

4. Model Compression and Efficient Inference for Large Language Models: A Survey

5. Efficient Transformers: A Survey

6. Keyformer: KV Cache Reduction through Key Tokens Selection for Efficient Generative Inference

7. FastCache: Optimizing Multimodal LLM Serving through Lightweight KV-Cache Compression Framework

8. Beyond RAG: Task-Aware KV Cache Compression for Comprehensive Knowledge Reasoning by Orion Weller