🤖 machine learning

LoKiFormer: Locality-aware Attention with Decoupled Knowledge Memory for Efficient Large Language Model Pretraining

이 논문은 국소적 패턴을 포착하기 위한 로컬 퓨전 어텐션(Local Fusion Attention)과 명시적인 전역 지식 검색을 위한 디커플드 지식 메모리 모듈(decoupled Knowledge Memory Module)을 통합함으로써 사전 학습 효율과 수렴 속도를 향상시킨 새로운 대규모 언어 모델 아키텍처인 LoKiFormer를 소개한다.

원저자: Qiuwu Chen, Zimo Liu, Yuchen Li, Ying Sun, Yifan Zhang, Zhijie Qiu, Zeng You, Ryan Dong, Simeng Ma, Yaofo Chen, Mingkui Tan

게시일 2026-08-14

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Qiuwu Chen, Zimo Liu, Yuchen Li, Ying Sun, Yifan Zhang, Zhijie Qiu, Zeng You, Ryan Dong, Simeng Ma, Yaofo Chen, Mingkui Tan

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

당신이 초지능 로봇에게 읽고 쓰는 법을 가르치려 한다고 상상해 보세요. 이 로보는 '거대 언어 모델(LLM)'이라 불리는 인공지능의 일종으로, 이야기 쓰기부터 수학 문제 풀이까지 모든 분야에서 놀라운 실력을 갖추게 되었습니다. 하지만 문제가 하나 있습니다. 이 로봇을 가르치는 것은 마치 티스푼으로 수영장을 채우려는 것과 같습니다. 엄청난 양의 시간과 컴퓨터 연산 능력이 필요하며, 로봇은 때때로 자신이 실제로 무엇을 배우고 있는지에 대해 혼란을 겪기도 합니다.

왜 그런지 이해하려면, 인간이 문장을 읽는 방식을 생각해 보세요. 우리는 단어들이 서로 바로 옆에 붙어 있을 때 대개 연결되어 있다는 것(예: "뜨거운"과 "커피")을 즉각적으로 이해하지만, 동시에 몇 년 전에 배웠던 사실들(예: "커피는 원두로 만든다")도 기억해 내야 합니다. 현재의 AI 모델들은 '어텐션(attention)'이라는 메커니즘을 사용하여 이 두 가지를 동시에 수행하려고 노력하는데, 이는 모든 단어를 살펴보고 그 단어들이 어떻게 연관되는지 확인하는 방식입니다. 문제는 이것이 마치 페이지 위의 모든 글자를 한꺼번에 뚫어지게 쳐다보며 책을 읽는 것과 같아서, 국소적인 연결을 처리하기에는 매우 지치고 비효율적이라는 점입니다. 또한, 로봇이 역사나 과학 같은 일반적인 지식을 저장하려고 할 때, 그 지식을 뇌 깊숙한 곳에 무질서하게 숨겨두기 때문에 필요한 사실을 적시에 꺼내기가 어렵습니다. 과학자들은 더 빨리 배우고, 에너지를 덜 사용하며, 정보를 정확히 어디에서 찾아야 하는지 아는 로봇을 만들고 싶어 합니다.

여기, 로봇의 안경과 기억 저장소를 똑똑하게 업그레이드하는 설계인 LoKiFormer가 등장했습니다. 이 논문의 저자인 Qiuwu Chen과 Zimo Liu가 이끄는 연구진은 기존의 방식이 에너지를 낭비하고 있다는 점에 주목했습니다. 그들은 로봇의 학습 과정을 해결하기 위해 두 가지 새로운 도구를 제안했습니다. 하나는 '국소적인(local)' 것(바로 옆에 있는 단어들)을 처리하기 위한 것이고, 다른 하나는 '전역적인(global)' 것(세상의 거대한 사실들)을 처리하기 위한 것입니다.

첫째, 그들은 국소적인 문제를 **국소 융합 어텐션(Local Fusion Attention, LFA)**이라는 것으로 해결했습니다. 당신이 문장을 읽고 있다고 상상해 보세요. 로봇이 "고양이"라는 단어가 문단 내의 다른 모든 단어와 어떻게 연관되는지 즉시 파악하려고 애쓰는 대신, LFA는 로봇에게 특별한 "국소 망원경"을 쥐여줍니다. 이 망원경은 바로 옆에 있는 몇 개의 단어에만 초점을 맞춰 그것들을 먼저 하나로 섞어줍니다. 이는 마치 로봇이 전체 도시를 이해하기 전에, 주변 동네의 요약본을 빠르게 훑어보는 것과 같습니다. 이를 통해 로봇은 불필요한 작업을 줄일 수 있습니다. 논문에 따르면, 이 간단한 단계를 추가함으로써 로봇은 국소적인 패턴을 훨씬 빠르게 이해하게 되며, 나중에 더 크고 복잡한 관계에 집중할 수 있는 두뇌 전력을 확보할 수 있습니다.

둘째, 그들은 **지식 메모리 모듈(Knowledge Memory-Module, KMM)**을 통해 기억 문제를 해결했습니다. 기존 모델에서 로봇의 지식은 책들이 선반에 본드로 붙어 있고 읽는 방법과 뒤섞여 있는 도서관과 같았습니다. 만약 로봇이 "물리학"에 관한 사실을 찾으려 한다면, 엉망이 된 도서관 전체를 헤집어야 했습니다. LoKiFormer의 KMM은 읽기 지침과는 분리된, 완벽하게 정리되고 라벨이 붙은 서류함과 같습니다. 이것은 사실들을 특정하고 주소 지정이 가능한 슬롯에 저장합니다. 로봇이 "화학"에 대해 알고 싶을 때, 로봇은 주의가 산만해지지 않고 곧장 "화학" 서랍으로 걸어가서 알맞은 파일을 꺼낼 수 있습니다. 이는 지식을 '저장하는 것'과 '사용하는 것'을 분리하여 과정을 훨씬 명확하고 유연하게 만듭니다.

이 두 가지 도구를 결합한 결과는 인상적입니다. 연구진은 새로운 LoKiFormer 모델을 학습시켰으며, 이 모델이 표준 모델보다 1.33배 더 빠르게 학습한다는 것을 발견했습니다. 이를 체감할 수 있게 설명하자면, 표준 모델이 특정 수준의 이해에 도달하기 위해 10,000번의 학습 단계를 거쳐야 했다면, LoKiFormer는 단 7,500번의 단계만으로 동일한 수준에 도달했습니다. 이는 시간과 에너지 측면에서 엄청난 절약입니다. 더욱 흥-미로운 점은, 70억 개의 파라미터를 가진 이 모델의 작은 버전이 다른 기업들의 훨씬 더 큰 유명 모델들보다 더 뛰어난 성능을 보였다는 것입니다. 언어 이해, 추론, 심지어 코딩 테스트에서도 그들을 앞질렀습니다.

이 논문은 새로운 아키텍처가 단순히 로봇을 더 빠르게 만드는 것이 아니라, 자신이 아는 것을 사용하는 데 더 똑똑하게 만든다는 점을 시사합니다. "국소적" 읽기와 "전역적" 메모리를 분리함으로써, 모델은 복잡한 과제를 더 효과적으로 처리할 수 있습니다. 연구진은 또한 이 "서류함(KMM)"이 학습 과정 중에 스스로를 조직화했다는 것을 보여주었는데, 서로 다른 서랍들이 역사, 물리학, 대수학 같은 서로 다른 주제의 전문가가 되도록 자연스럽게 변했다는 것입니다. 이는 로봇이 단순히 암기하는 것이 아니라, 인간이 자신의 지식에 접근하는 방식과 유사하게 정보를 검색하는 법을 배우고 있음을 의미합니다.

요약하자면, LoKiFormer는 AI 모델을 더 좋게 만들기 위해 반드시 모델을 더 크게 만들고 비용을 더 많이 들일 필요는 없다는 것을 보여줍니다. 세부 사항을 보는 더 나은 도구와 큰 사실들을 저장하는 더 깔끔한 방법을 제공함으로써, 우리는 더 빨리 배우고, 훈련 비용이 적게 들며, 세상을 더 명확하게 이해하는 AI를 구축할 수 있습니다. 이는 때때로 초지능 기계를 만드는 가장 좋은 방법은 그 기계가 책상을 정리하도록 돕는 것임을 상기시켜 줍니다.

기술 요약: LoKiFormer

문제 정의

대규모 언어 모델(LLM)이 다양한 응용 분야에서 놀라운 돌파구를 마련했음에도 불구하고, 이들의 아키텍처는 서로 다른 수준의 정보를 모델링하는 데 있어 두 가지 주요한 한계로 인해 사전 학습(pretraining) 과정에서 계산 효율성이 떨어집니다:

비효율적인 국소 모델링(Inefficient Local Modeling): 표준 셀프 어텐션(self-attention) 메커니즘은 국소성(locality)에 대한 명시적인 귀납적 편향(inductive bias)이 부족합니다. 이는 시퀀스 내부의 국소 정보(예: 구문 구조 및 단거리 의존성)를 중복된 전체 쌍별 상호작용(full pairwise interactions)을 통해 모델링합니다. 이로 인해 모델이 국소적 패턴만을 필요로 하는 상황에서도 글로벌 상관관계를 계산해야 하므로, 국소 패턴을 학습하는 것이 계산적으로 비효율적이 됩니다.
암묵적 지식 결합(Implicit Knowledge Coupling): 혼합 전문가(Mixture-of-Experts, MoE) 아키텍처는 지식 저장과 계산 경로를 암묵적으로 결합합니다. 지식이 전문가 네트워크의 가중치에 분산되어 있어, 검색이 간접적이며 시퀀스 외부의 글로벌 지식(예: 상식 및 사실적 정보)에 대한 유연한 접근을 방해합니다. 이러한 결합은 사전 학습 중 지식 내재화 및 검색의 투명성과 효율성을 제한합니다.

방법론

이러한 한계를 해결하기 위해 저자들은 시퀀스 내부의 국소 정보와 시퀀스 외부의 글로벌 정보를 각각 처리하도록 설계된 두 개의 전용 모듈로 표준 디코더 블록을 보강한 새로운 LLM 아키텍처인 LoKiFormer를 제안합니다.

1. 국소 융합 어텐션 (Local Fusion Attention, LFA)

LFA는 단거리 의존성 모델링을 개선하기 위해 명시적인 국소 귀납적 편향을 도입하도록 설계되었습니다.

메커니즘: 어텐션 메커니즘 이전에, LFA는 인접한 토큰 표현을 융합하는 그룹 컨볼루션(group convolution) 연산을 적용합니다. 은닉 상태(hidden states)는 $h$ 개의 그룹(어텐션 헤드 수와 일치)으로 분할되며, 각 그룹 내에서 $k$ 개의 즉각적인 이웃으로부터 정보를 집계하는 커널 크기 $k$ 의 인과적 1D 컨볼루션(causal 1D convolution)을 수행합니다.
효과: 이 컨볼루션 융합은 국소적 패턴을 명시적으로 포착하여 더욱 풍부하고 문맥을 인식하는 표현을 생성합니다. 결과적으로, 후속 멀티 헤드 잠재 어텐션(Multi-Head Latent Attention, MLA) 메커니즘은 더 많은 정보를 가진 표현을 바탕으로 작동할 수 있어, 중복된 국소 상호작용 대신 더 넓고 글로벌한 문맥 관계를 모델링하는 데 집중할 수 있습니다.
효율성: LFA는 경량화된 수정 사항으로, 전체 어텐션 복잡도를 변경하지 않으면서 파라미터를 아주 미미하게 증가시킵니다 (7B 모델 기준 0.41%).

2. 지식 메모리 모듈 (Knowledge Memory Module, KMM)

KMM은 지식 저장을 계산으로부터 분리하여 글로벌 지식의 명시적이고 직접적인 검색을 가능하게 합니다.

메커니즘: 지식을 가중치에 암묵적으로 인코딩하는 MoE와 달리, KMM은 파라미터화된 키-값(key-value) 메모리 구조를 도입합니다. 이는 MLA 모듈의 잠재 표현을 입력으로 받아 이를 쿼리 공간으로 투영하고, 도메인 특화 지식 필드를 인코딩하는 학습 가능한 키(keys)와 매칭합니다.
구조: 이 모듈은 지식을 $F$ 개의 주소 지정 가능한 슬롯(필드)에 저장합니다. 학습 중에 이 필드들은 엔드투엔드(end-to-end) 방식으로 최적화됩니다. 추론 중에 토큰은 이 필드들을 직접 쿼리하여 관련 지식 추상화를 검색합니다.
통합: 검색된 지식 표현은 집계되어 MoE 레이어에 의해 처리되기 전 MLA 레이어의 출력에 더해집니다. 이 설정은 모델이 새로운 패턴의 생성(MoE를 통해)과 사전 학습된 지식의 효율적인 검색(KMM을 통해) 사이에서 균형을 맞출 수 있도록 합니다.
해석 가능성: 저자들은 엔드투엔드 학습을 통해 개별 지식 필드가 특정 도메인(예: 물리학, 역사)의 개념을 나타내도록 자발적으로 진화함을 보여줌으로써, 표적화되고 해석 가능한 지식 검색을 입증했습니다.

주요 기여

본 논문은 세 가지 주요 기여를 설명합니다:

효율적인 어텐션을 위한 국소 귀납적 편향: 컨볼루션 융합을 사용하여 국소 패턴을 명시적으로 포착하는 LFA를 도입했습니다. 이는 중복된 국소 상호작용을 줄이고 어텐션 메커니즘이 복잡한 글로벌 관계에 집중할 수 있게 하여, 셀프 어텐션의 단거리 의존성 모델링 비효율성을 해결합니다.
글로벌 지식을 위한 분리된 파라미터 메모리: 주소 지정 가능한 슬롯에 재사용 가능한 지식 추상화를 명시적으로 저장하는 KMM을 제안합니다. 이는 저장을 계산으로부터 분리하여, 표준 MoE의 암묵적이고 결합된 지식 표현과 대조되는 직접적인 지식 검색과 해석 가능한 추론을 가능하게 합니다.
개선된 사전 학습 효율성: LFA와 KMM의 통합은 사전 학습 중 유의미한 수렴 속도 향상을 가져옵니다.

실험 결과

저자들은 다양한 모델 크기(1B ~ 60B 파라미터)와 벤치마크를 통해 LoKiFormer를 평가했습니다.

사전 학습 수렴: LoKiFormer는 베이스라인 모델보다 1.33배 빠르게 수렴합니다. 구체적으로, 베이스라인이 10K 스텝을 소요할 때 LoKiFormer는 7.5K 스텝 만에 동일한 검증 퍼플렉시티(validation perplexity)에 도달합니다.
다운스트림 성능:
- 유사 규모 모델: LoKiFormer-7B는 언어 이해(MMLU, CMMLU, C-Eval), 추론(HellaSwag, ARC-Challenge), 코드(HumanEval), 수학(GSM8K) 작업 전반에서 규모가 유사한 밀집(dense) 및 MoE 베이스라인(예: Llama-3.1-8B, Qwen2.5-7B, Mixtral-8×7B)을 능가합니다.
- 프런티어 모델: 총 7B 파라미터(활성화 3B)만 보유하고 있음에도 불구하고, LoKiFormer-7B는 MMLU, CMMLU, C-Eval과 같은 벤치마크에서 프런티어 오픈 소스 모델(예: Llama-3.1-70B, Qwen2.5-72B) 및 폐쇄형 시스템(예: Claude-3.5-Sonnet, Gemini-1.5-Pro)을 능가하는 최첨단 성능을 달성합니다. 그러나 추론 및 생성 작업에서는 이러한 더 큰 프런티어 모델들에 비해 성능이 낮으며, 특히 HellaSwag, HumanEval, GSM8K에서 Claude-3.5-Sonnet 및 Gemini-1.5-Pro보다 낮은 점수를 기록했습니다.
확장성: 아키텍처는 1B에서 60B 파라미터까지 매끄럽게 확장되며, 모델 용량이 증가함에 따라 퍼플렉시티와 다운스트림 작업 성능에서 일관된 개선을 보입니다.
효율성: 스텝당 계산 오버헤드는 매우 적으며(처리량 약 2.6% 감소), 이는 학습 단계의 감소가 전체 실제 학습 시간(wall-clock training time)의 상당한 절감으로 직결됨을 의미합니다.
어블레이션 연구(Ablation Studies): 실험을 통해 LFA와 KMM이 모두 독립적이고 상호 보완적으로 성능 향상에 기여함을 확인했습니다. 커널 크기, 그룹 수, 지식 필드 수에 대한 어블레이션은 설계 선택을 검증합니다.

의의 및 주장

본 논문은 LoKiFormer가 LLM 사전 학습의 효율성과 효과성을 높이는 중요한 진전이라고 주장합니다. 국소 의존성 모델링의 비효율성과 암묵적 지식 저장의 경직성을 명시적으로 해결함으로써, 이 아키텍처는 다음과 같은 성과를 달성합니다:

빠른 수렴: 대규모 모델 학습에 필요한 계산 비용과 에너지 소비를 줄입니다.
표현 능력 강화: 국소 패턴과 글로벌 지식을 통합하는 모델의 능력을 향상시켜, 특히 언어 이해와 같은 복잡한 추론 및 지식 집약적 작업에서 더 강력한 성능을 이끌어냅니다.
해석 가능성: 저장과 계산을 명시적으로 분리함으로써 지식 접근에 대한 더 투명한 메커니즘을 제공합니다.

저자들은 LoKiFormer를 기존 아키텍처에 대한 우수한 대안으로 제시하며, 아키텍처 혁신이 특정 도메인(특히 언어 이해)에서 훨씬 더 큰 모델에 필적하거나 이를 능가하는 성능 향상을 가져올 수 있음을 입증함으로써, 보다 자원 효율적이고 지속 가능한 AI 시스템 개발에 기여하고자 합니다. 향후 연구로는 이 아키텍처를 멀티모달 추론 작업으로 확장하는 것이 제안되었습니다.