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

글을 시작하며

작년 말 오픈AI에서 o1을 발표한 이후로, AI 모델의 발전 방향이 ‘빠르게 답변을 제공’하는 것으로부터 ‘문제를 작은 단계로 나눠서 더 신중하게 생각하는’ 방향으로 변했다는 것, 다들 느끼실 겁니다.

CoT (Chain-of-Thoughts) 기법을 구현해서 만든 o1 같은 사고 방식은, 오픈AI의 o1, o3, 그리고 DeepSeek-R1 같은 거대한 추론 모델로 하여금 ‘(생각하는 과정의) 이전 단계로 되돌아가서 다시 생각하고, 추론 과정을 개선할 수 있게’ 해서 까다로운 문제까지도 더 잘 해결할 있게 해 주죠. 이전의 AI 101 에피소드 중 하나로 ‘Test-Time Compute’를 스케일링하는 게 왜 중요한지에 대해 말씀드린 적도 있는데요:

여기 큰 문제가 하나 앞길을 가로막고 있죠 - 바로 메모리 문제예요. 이런 복잡한 추론은, 많은 텍스트, 즉 토큰을 만들어낼 수 밖에 없어서, 메모리를 많이 차지하고 처리 속도도 늦어지고, 비용이 증가하죠. 트랜스포머 구조의 모델에서 이런 현상이 특히 두드러질 수 밖에 없구요. 이런 추론을 위한 거대한 모델이 점점 저 보편적으로 사용될 것으로 예상되기 때문에, 이런 모델들을 어떻게 개선할 건지 충분히 잘 탐색하면서 약점을 줄이는 방법을 만들어나가는 게 중요합니다.

그래서, 오늘은 ‘추론 모델에서 나타나는 메모리 사용량의 증가, 그리고 그에 따른 처리 시간의 지연’이라는 문제에 초점을 맞춰보려고 합니다. 메모리와 관련된 비효율성을 해결할 수 있다면, 모델이 정확도를 높이 유지하면서도 성능과 비용의 밸런스가 좋아지겠죠. 두 가지 주목할 만한 접근법이 이미 제안되어 있는데요:

1) 모델이 스스로의 ‘생각’을 요약하는 법을 학습해서, 짧지만 의미있는 요약 내용을 기반으로 메모리 부담없이 추론 작업을 하도록 해 주는 LightThinker; 그리고

2) DeepSeek가 DeepSeek-V2를 출시했을 때 제안하고 나중에 DeepSeek-V3와 DeepSeek-R1에 구현한 솔루션, Multi-head Latent Attention (MLA)

오늘은 한 번 이 기법들에 대해서 구체적으로 알아보고, 또 이 기법들을 혼합한다면 어떤 좋은 점이 있을까 생각해 보겠습니다..

오늘 에피소드에서는 다음과 같은 내용을 다룹니다:

• LightThinker는 무엇인가?

  • LightThinker의 핵심 아이디어

  • LightThinker의 작동 방식

  • LightThinker의 실제 성능

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

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

  • MLA가 왜 필요했을까?

  • MLA의 작동 방식

  • MLA 기법의 장점과 한계

  • MLA의 성능

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

• 맺으며

• 보너스: 참고자료

LightThinker는 무엇인가?

LightThinker의 핵심 아이디어

앞서 언급한 것처럼, 추론 모델의 확산을 위해서는 ‘메모리 비용을 적절하게 유지하면서 고품질의 추론을 훨씬 더 빠르고 효율적으로 하도록’ 하는 최적화 기법이 필요합니다.

여기서 이야기하려고 하는, 이런 기법 중 하나가 바로 중국 저장 대학과 앤트 그룹이 공동으로 운영하는 ‘Joint Laboratory of Knowledge Graph’에서 개발한 LightThinker입니다. LightThinker는 단순하게 단어나 메모리를 수동적으로 잘라내 버리는게 아니라, 모델이 문제를 해결하는 동안 자기가 하는 ‘생각을 요약’하도록 가르칩니다. 사람들이 이야기를 듣거나 할 때 모든 세부 사항을 적지 않고 핵심적인 요점만 메모해 두는 것과 비슷하다고 할까요?

그럼, 이 LightThinker가 어떻게 작동하는지 한 번 자세히 살펴보도록 하죠.

LightThinker의 작동 방식

LightThinker 길고 상세한 추론 단계를 유지하는 대신에, 추론을 더 짧고 핵심적인 내용만 포함하게끔 압축한 다음에 그걸 기반으로 추론을 계속합니다. 기억해야 할 건, LightThinker가 두 가지 작업을 한다는 점인데요:

• ‘언제’ 추론 단계를 압축할지 결정합니다.

• ‘어떻게’ 그 단계까지 추론한 내용을 압축할지 결정합니다.

이런 작업을 수행하는데, 보통 아래와 같은 기법을 사용합니다.

언제 압축을 할까?

모델의 ‘생각’을 언제 요약, 압축할지 결정하는 두 가지 방법이 있는데요:

• 토큰 레벨 (Token-level) 압축: 모델이 추론 과정에서 일정한 숫자의 단어에 도달한 이후에 그 때까지의 ‘생각’을 압축하는 방법입니다. 간단한 방버이지만, 생각을 잘라낼 때 어색할 수가 있겠죠.

• 사고 레벨 (Thought-level) 압축: 모델이 문장이나 단락 같이 어느 정도 완결적인 아이디어를 만든 후에 그 때까지의 ‘생각’을 압축하는 방법입니다. 이렇게 하면 생각 자체는 더 체계적으로 유지되지만, 생각이 언제 완료되었는지 결정을 해야 하니까 추가적인 처리가 필요하겠죠. 그래도, 연구자들은 이 기법이 ‘의미’를 더 잘 보존한다고 생각해서, LightThinker에서 이 방식의 압축 기법을 선호합니다.

어떻게 압축을 할까?

정보를 요악하는 것도 두 가지 중요한 방법이 있습니다:

• 텍스트 (Text) 압축: 모델이 ‘긴 생각’을 ‘짧은 요약’으로 대체하는 방식입니다. 그런데 이 방법은 추가적인 인코딩 모델과 함께 추가 처리가 필요해서 속도가 느려지는 단점이 있습니다.

• 은닉 상태 (Hidden State) 압축: 텍스트를 다시 작성하는 대신, 모델이 특별한 토큰에 핵심적인 세부 정보를 저장합니다. 이런 토큰들이, AI 모델이 나중에 필요 시 사용할 수 있는 일종의 ‘기억 속에 있는 메모’처럼 작동합니다. 이 기법에서는 추가적인 모델이 필요없어서, 연구자들이 이 기법을 더 선호합니다.

자, 그럼 LightThinker 기법을 사용해서 단계별로 실제로 압축을 구현하는 작업 순서를 알아봅시다:

읽어주셔서 감사합니다. 친구와 동료 분들에게도 뉴스레터 추천해 주세요!