DFM은 청크 단위 메모리를 디코더 중간에 융합해 검색 없이 지식을 주입한다
Deep Fusion Memory: A Chunk-Level Parametric Memory for Large Language Models
무엇인가
이 논문은 LLM이 사전학습 이후의 외부 지식을 효율적으로 습득하고 사용하는 문제를 다룬다. 계속 사전학습은 기존 능력의 파국적 망각을 부를 수 있고, 추론 시 검색(RAG, RETRO 등)은 데이터스토어 검색과 증거 처리 때문에 지연을 더한다. 기존 파라메트릭 메모리는 추론 시 검색을 없애지만, 토큰 단위 목표를 만들기 위한 오프라인 검색 비용이 크고 출력 확률 보간에 그쳐 융합이 얕다는 한계가 있다. DFM은 이 두 비용을 줄이면서 메모리를 디코더 중간 은닉 상태에 넣는 것을 목표로 한다.
어떻게 동작하나
DFM의 훈련 데이터 구성은 토큰 단위가 아니라 청크 단위다. 훈련 시퀀스를 C개 토큰의 겹치지 않는 청크 X0,...,XJ-1로 나누고, j≥1인 각 청크 Xj에 대해 오프라인 검색기가 직전 청크 Xj-1을 인코딩해 코퍼스 데이터스토어에서 상위 K개 이웃을 찾는다. 각 검색 청크는 고정 텍스트 인코더의 밀집 임베딩으로 표현되고, 검색 순위대로 쌓아 목표 메모리 Vj ∈ R^{K×dm}를 만든다. 이 목표는 Xj 안의 여러 토큰 예측이 공유하며, 첫 청크는 직전 청크가 없어 메모리 감독에서 제외된다. 목표 메모리는 오프라인에서 만들어져 훈련 중 고정된다.
무엇과 다른가
메모리 예측기는 직전 문맥으로부터 각 청크의 K개 연속 메모리 벡터를 생성한다. 고정 백본의 토큰 임베딩을 직전 청크 Xj-1에서 평균 풀링하고 필요하면 선형 투영해 청크 임베딩 Cj를 만들고, 인과적 Transformer와 K개의 별도 선형 출력 헤드로 V̂j = Predictorφ(C≤j)를 낸다. 첫 입력 C0은 마스크된 자리표시자다. 예측 메모리는 공유 메모리 프로젝터 Pψ를 거쳐 디코더 은닉 공간의 Mj = Pψ(V̂j) ∈ R^{K×d}가 되며, 각 슬롯에 LayerNorm, GELU가 있는 두 선형층을 독립 적용한다. 이 Mj는 선택된 여러 디코더 층에서 재사용된다.
어떻게 쓰나
융합은 선택된 디코더 층에서 게이트 교차어텐션으로 일어난다. 정규화된 토큰 상태 zt를 쿼리로, Mj를 키와 값으로 사용해 H개 헤드가 K개 슬롯에 소프트맥스를 적용하고, 헤드별 출력 ot,a를 낸다. 게이트 gt = σ(Wg zt + bg)는 토큰과 헤드마다 스칼라를 주어 ut = WO Concat(g_{t,a} ot,a)로 메모리 기여를 조절한다. 메모리 분기는 인과적 자기어텐션과 병렬로 같은 잔차 업데이트 h̃t = ht + SelfAttn(z≤t) + ut에 들어가고, 이후 원래 피드포워드 층을 통과한다. 예측기와 프로젝터는 융합 층들이 공유하지만, 각 융합 층의 메모리 어텐션과 게이트 투영은 층별로 따로 둔다. 다음 토큰 예측과 정렬하기 위해 위치 t의 상태가 x_{t+1}을 예측하고 청크 인덱스 j=floor((t+1)/C)를 쓰며, 청크 경계에서는 직전 청크가 완료된 뒤 풀링 쿼리가 사용된다.
전제와 한계
훈련은 백본을 고정한 채 예측기와 융합 모듈을 다음 토큰 교차엔트로피와 메모리 예측 MSE로 공동 최적화한다. 목적식은 L = L_LM + β L_MSE이고, L_MSE는 첫 청크를 제외한 예측 메모리와 목표 메모리의 프로베니우스 노름 제곱을 (J-1)Kdm으로 나눈 값이다. β=1이며 MSE는 L2 정규화된 목표 임베딩에 직접 계산하고 예측을 다시 정규화하지 않는다. L_LM의 그래디언트는 메모리 융합을 통해 예측기로 흐르고, L_MSE는 예측 메모리 슬롯을 직접 감독한다. 추론 시에는 목표 메모리를 만들지 않고 같은 예측기와 융합 모듈만 사용하며, 데이터스토어 검색이나 검색 텍스트, 외부 임베딩 인코더가 필요 없다. QA에서는 보이는 프롬프트로 예측 메모리를 계산해 답 생성 동안 고정한다.
QA 실험은 2018년 12월 영어 위키백과 약 4B 토큰으로 훈련하고, 단일 홉 NQ와 WebQA, 다중 홉 HotpotQA, 2WikiMultiHopQA, MuSiQue, Bamboogle에서 평가했다. Mistral-7B-v0.3과 Qwen3-8B-Base에서 DFM은 온라인 검색 없이 두 백본 모두 최고 평균을 냈다. Mistral에서는 6개 과제 평균이 xRAG의 24.14에서 27.21로 올라가 최강 베이스라인 대비 12.71% 상대 개선이었고, Qwen3-8B-Base에서는 MLP Memory의 23.80에서 26.46으로 11.19% 상대 개선이었다. 절대 개선은 각각 3.07점과 2.66점이다. Base 대비로는 13.59점과 3.75점을 올려 Mistral은 6개 과제 전부, Qwen3는 5개 과제를 개선했고, Mistral의 모든 과제와 Qwen3의 5개 과제에서 CPT와 LoRA도 앞섰다. 세부적으로 Mistral의 NQ는 33.57, WebQA는 42.47로 Base보다 13.27점, 13.19점 높았고, Qwen3의 NQ는 17.60에서 27.46으로 RAG보다 3.04점 높았다. 반면 Qwen3의 WebQA는 41.78에서 39.47로 떨어져 MLP Memory의 43.26에 못 미쳤고, 2Wiki에서는 Mistral의 25.00이 RAG의 31.02보다 낮았으며 Qwen3 HotpotQA의 26.46도 xRAG의 27.74보다 낮았다.
도메인 적응에서는 Qwen3-4B-Base와 Qwen3-8B-Base로 Code, Math, Law를 실험해 네 벤치마크 모두 Base보다 개선했고, 비교 방법 중 최고 평균 64.69와 67.52를 기록해 각각 6.51점, 4.38점을 올렸다. 4B에서는 HumanEval, HumanEval+, GSM8K가 Base보다 7.47점, 5.48점, 2.65점 올랐고, 8B에서는 HumanEval 73.02, HumanEval+ 59.15에서 64.02를 냈다. LawBench는 4B에서 29.96에서 40.41, 8B에서 38.45에서 46.61로 올라 RAG보다 각각 16.06점, 19.72점 높았지만, 4B LawBench는 CPT보다 1.25점 낮았고 8B HumanEval+와 GSM8K에서는 LoRA가 더 강했다. 효율성에서는 DFM이 0.75B 파라미터를 추가해 Mistral MLP Memory의 1.5B 구성보다 약 50% 적었고, 같은 4B 토큰 코퍼스와 16대 A800에서 MLP Memory 약 64시간, DFM 약 24시간으로 훈련 시간을 약 62.5% 줄였다. 목표 구성은 4B 토큰 코퍼스에서 토큰 단위 4B회 검색 쿼리가 필요한 반면 32토큰 청크는 125M회로 96.88% 줄었고, Top-1024 대 Top-16을 적용하면 검색 항목이 4.096T에서 2B로 99.95% 줄었다.
개발자 관점에서 DFM은 추론 시 검색 인프라를 붙이기 어렵거나 검색 지연이 병목인 서비스에서, 백본을 건드리지 않고 청크 단위 메모리 모듈만 학습해 지식을 주입하는 선택지로 읽힌다. 다만 백본과 검색 인코더, 목표 메모리를 고정하는 전제가 있고, 첫 청크는 메모리 감독에서 빠지며, 메모리 적응이 모든 백본과 벤치마크에서 균일하게 이득이 되지는 않았다. Qwen3 WebQA 하락, 일부 다중 홉 과제에서 RAG 우위, 4B LawBench에서 CPT 우위, 8B HumanEval+와 GSM8K에서 LoRA 우위가 원문에 명시되어 있으므로, 도입 전에는 대상 도메인과 백본에서 재현 실험을 해야 한다. 또한 원문은 코드가 곧 공개될 예정이라고만 밝히고 실제 저장소 URL은 제시하지 않았다.