[딥러닝, 논문리뷰] Meta - Memory Layers at Scale

https://ai.meta.com/research/publications/memory-layers-at-scale/

 

Memory Layers at Scale | Research - AI at Meta

기존 긴 문맥을 처리하기 위해, RAG와 유사하게 작동하지만 외부 데이터 베이스에서 관련 정보를 찾는 방법과 달리 내부 메모리에서의 Key-Value 구조를 이용하여 긴 문맥을 효율적으로 저장하고 검색하는 기능을 제공하는 방법이라고 생각하면 될 것 같습니다.

1. Abstract

• Memory Layers는 학습 가능한 Key-Value 방식의 조회 메커니즘을 사용하여 사용 계산량(FLOPs라고도 합니다) 증가 시키지 않고 추가 파라미터를 적용 가능함
• Dense FeedFoward Layer(FC라고 보시면 됩니다) 와 상호 보완적이며, 정보를 효율적으로 저장 및 검색하는 전용용량을 제공함
• 대규모 스케일에서 메모리 레이어의 실용성을 입증했으며, 1조개의 토큰으로 훈련된 1,280억개의 메모리 파라미터를 사용하는 모델을 선보임
• 결과적으로 더 높은 계산 비용을 가진 Dense Model과 MoE 모델을 능가했으며, 특히 사실적 태스크에서 두드러진 성능을 확인함
• 완전한 병렬 구조의 메모리 레이어를 구현하는 구조를 제공하며 확장성을 입증하였음.

2. Introduce

• 메모리 파라미터 수의 따른 정확도를 비교한 차트(QA정확도, NLL)
  • 해석 : 메모리 파라미터가 증가할수록 NQ,TQA 데이터셋의 정확도가 상승, TQA는 더 큰 메모리에서 성능이 급격하게 개선되는 모습을 보임 / NLL(Negative Log Likelihood)와 비교하였을 때도 TQA부문이 부정값이 더 낮은 값으로 성능이 개선되는 결과 보고 있음

이 논문에서 말하고자 하는 내용 자체는 추가 계산 비용 없이 파라미터를 확장시킬 수 있는 Key-Value 조회 메커니즘을 이용하여 설계되었고(이 부분은 희소 활성화 패턴이라고도 부른답니다)

 

Dense FeedFoward Layer와 상호 보완적으로도 작동하며, 결국 긴 문맥과 사실에 관련된 지식(해석으로는 사실적 지식이라고 하는데, 제가 판단하기로는 실제 웹상을 통해 문서로 나와 있는 실제 사실들을 얘기하는 것 같아요 → 할루시네이션 대비 진짜 사실 기반으로 답변을 하는 것이 목표이기 때문, 이건 RAG기법과 동일한 목적 같습니다)

 

이 접근 방법은 MoE 모델과 Dense 모델을 능가하는 성느을 보여주면서 사실 기반의 태스크에서 큰 성능을 보임.

핵심 포인트 :

1. Dense Layers (밀집 레이어):
   • 주로 **FLOP-bound (부동소수점 연산량에 의해 제한)**되며, 모든 입력과 출력을 계산해야 하기 때문에 계산 비용이 큽니다.
2. Memory Layers (메모리 레이어):
   • Sparse Activation (희소 활성화) 패턴을 사용함
   • 전체 메모리를 모두 활성화하지 않고, 관련 있는 일부 메모리 슬롯만 접근함 (이 부분이 가장 중요)
   • 이로 인해 계산(FLOPs)보다는 메모리 대역폭에 더 많이 의존함
3. 하드웨어 최적화:
   • 메모리 대역폭에 의존하는 아키텍처는 현재 하드웨어 가속기(예: GPU, TPU)에서 최적화되지 않은 경우가 많음
   • 이는 메모리 레이어가 하드웨어적으로 덜 효율적일 수 있음을 의미하지만, 계산 비용 자체는 증가하지 않음

특징 :

1. FFN 대체:
   • 일부 Transformer 레이어의 **Feed-Forward Network (FFN)**을 Memory Layers로 교체.
   • 다른 레이어(예: Self-Attention 등)는 그대로 유지.
2. 확장성:
   • 기본 모델 크기: 1억 3400만(134M) ~ 80억(8B) 파라미터 범위에서 일관된 성능 개선 확인.
   • 메모리 용량: 최대 1280억(128B) 파라미터로 확장.
     • 이는 기존 연구에서 보고된 메모리 용량 대비 100배 이상 확장된 것임.
3. 성능 개선:
   • 사실적 질문 응답(Factual QA): 정확도가 100% 이상 개선됨.
   • 코딩 태스크(HumanEval, MBPP): 성능이 크게 개선됨.
   • 일반 지식(Hellaswag, MMLU): 성능이 전반적으로 개선됨.
4. Dense 및 MoE 모델과의 비교:
   • Dense Models: 4배 더 많은 계산량을 사용한 Dense 모델과 동등한 성능을 달성.
   • Mixture-of-Experts (MoE): 같은 계산 및 파라미터 규모에서도 더 나은 성능을 보임, 특히 **사실적 태스크(Factual Tasks)**에서 두드러짐.

---

3. Memory Augmented Architectures

1. 어텐션 메커니즘과의 유사성
   • Memory Layers는 Attention 메커니즘과 유사하게 작동함
   • 입력으로는 \(쿼리(Qeury, q \in \mathbb{R}^n, \quad) 키(Key, K \in \mathbb{R}^{N \times n}), \quad 값(Value, V \in \mathbb{R}^{N \times n}\))를 받음
   • 출력은 쿼리와 키의 유사도 값에 따라 가중합 된 값의 조합
2. 메모리 레이어의 수식

수식별 해석 :

 

1단계: Top-k 인덱스 선택

 

\(I = \text{SelectTopkIndices}(K q)\)

 

설명:

• 입력 쿼리(\(q\) )와 키 행렬(\(K\)) 간의 유사도를 계산합니다.
  • 가장 높은 유사도를 가진 Top-k 키(Key)의 인덱스를 선택
  • \(I\): 선택된 Top-k 키의 인덱스 집합.
  • 핵심: 전체 키를 사용하지 않고, Top-k만 활성화하여 계산 비용을 줄임

2단계: Softmax 정규화

 

\(s = \text{Softmax}(K_I q)\)

 

설명:

• 선택된 Top-k 키 \(K_I\)와 입력 쿼리 \(q\)의 유사도를 계산
• Softmax 함수를 통해 유사도를 확률 분포로 변환
• \(s\): Top-k 키에 대한 Softmax 가중치 벡터\(( s \in \mathbb{R}^k )\).

3단계: 가중합 (Weighted Summation)

 

\(y = s V_I\)

 

설명:

• Softmax로 정규화된 가중치\(s\)를 사용해 Top-k 값 \(V_I\)을 가중합
• \(y\): 최종 출력 벡터 \((y \in \mathbb{R}^n)\)
• 핵심: 선택된 Top-k 값들이 가중합되어 최종 결과를 생성

3. 메모리 레이어와 어텐션의 차이점
   • Memory Layers: 키(Key)와 값(Value)이 학습 가능한 파라미터, 활성화 함수를 통해 동적으로 계산하지 않음(고정된 값)
   • Attention Layers: 키와 값은 주로 활성화 함수를 통해 동적으로 계산됨

3-1.Scaling memory layers

• Memory Layers는 계산(FLOPs) 부담은 적고 메모리 대역폭 의존도가 높기 때문에, 확장시 기존 Dense Layer와는 다른 도전 과제를 가지고 있음.
• 이 섹션에서는 이러한 확장성 문제를 자세히 설명하고, 이를 어떻게 해결했는지 다룸.

Scaling Challenge (확장성 문제)

• 계산 비용은 낮지만 메모리 대역폭 사용이 큼:
• 대부분의 연산이 희소 활성화를 통해 이루어지므로, 계산량은 적지만 메모리 접근 속도가 병목이 될 수 있음.
• 하드웨어 최적화 문제
  • 현재 GPU나 TPU 같은 하드웨어 가속기들은 Dense Layers에 최적화되어 있으며, Sparse 메모리 접근 패턴은 하드웨어에서 비효율적일 수 있음.
• Top-k 조회(Top-k Lookup):
  • 수백만 개의 키-값 쌍 중 가장 유사도가 높은 Top-k 쌍을 찾는 작업은 확장할수록 더 어려워짐.

3-1.1. Product-Key Lookup

• Top-k 유사도 기반 검색: 입력 쿼리 \(q\)가 키 행렬 \(K\)과 비교되어 유사도가 높은 Top-k 키(Key)를 선택.
• Sparse Activation (희소 활성화): 모든 키-값 쌍을 탐색하지 않고, 가장 관련성이 높은 일부(Key-Value Pair)만 활성화.
• 효율적인 검색: 수백만 개의 Key-Value 쌍 중에서 필요한 정보만 빠르게 조회.

확장성 문제 해결 

• 메모리 대역폭 문제 완화:
  • 모든 Key-Value 쌍을 동시에 탐색하지 않고, 가장 관련성이 높은 Top-k만 사용하여 메모리 대역폭 사용량을 절감
• 계산 효율성 개선:
  • 희소 활성화(Sparse Activation)를 통해 불필요한 계산을 피하고, 연산량을 크게 줄임
• 확장성 유지:
  • 키(Key)와 값(Value)의 수를 수백만 단위로 확장하면서도, Top-k 조회 방식으로 효율성을 유지

3-1.2. Parallel memory(병렬 메모리화)

• 메모리 블록을 여러 개로 나누고, 각 블록을 독립적으로 병렬 처리
• 각 쿼리(Query)는 여러 메모리 블록에 동시에 접근하여, Top-k Key-Value 쌍을 독립적으로 선택하고 가중합함 

3-1.3.Shared memory(공유 메모리)

• 모든 프로세스가 하나의 동일 메모리 공간을 공유 하는 방법 
• 여러 연산이 하나의 메모리 풀에 접근하여 데이터를 읽고 쓰는 방식 
• 데이터를 중복 저장하지 않고, 여러 연산 간 데이터 일관성을 유지함 

결론

• Parallel Memory: 독립된 여러 메모리 블록을 병렬로 처리하여 확장성과 효율성을 극대화 함
• Shared Memory: 여러 프로세스가 동일한 메모리 공간을 공유하여 데이터 일관성을 보장함

3-1.4.Performance and stability improvements

이 섹션은 Memory Layers의 성능 및 안정성을 개선하기 위한 다양한 기법과 접근법을 다루고 있음

 

기존 메모리 레이어의 문제점 

• 메모리 접근의 불균형: 특정 메모리 슬롯(Key-Value Pair)이 과도하게 활성화되거나 비활성화될 수 있음.
• 학습 안정성 문제: 키-값(Key-Value) 매칭의 불균형이 발생하면 학습이 불안정해질 수 있음.
• 희소 활성화(Sparse Activation)의 한계: 희소화된 Top-k 조회가 때때로 부정확할 수 있음.

이 문제점들을 해결하기 위한 방법이 제기되었는데.

 

1. Key-Value Normalization (정규화)

• Key, Value 값의 분포를 균일하기 만들어 학습을 안정화 하기 위해 사용 
• 키(Key)와 값(Value) 벡터를 LayerNorm 또는 RMSNorm과 같은 정규화 기법을 통해 스케일 조정.
• 쿼리(Query)와 키(Key) 간의 유사도가 특정 슬롯에만 집중되지 않도록 조정.

2. Dynamic Sparse Activation (동적 희소 활성화)

• Top-k 조회가 항상 최적의 키(Key)를 선택하도록 개선하기 위해 사용
• 학습 중 동적으로 Top-K 값을 조정해 유연성을 확보 
• 적절한 슬롯만 활성화 되도록 메커니즘 최적화

3. Memory Slot Regularization (메모리 슬롯 정규화)

• 메모리 슬롯들이 균등하게 활용되도록 하기 위해 사용
• 특정 키(Key)나 값(Value)에 과도한 활성화가 발생하지 않도록 정규화 패널티(Penalty)를 추가.
• 모든 슬롯이 학습 과정에서 고르게 참여하도록 유도.

4. Efficient Query Matching (효율적 쿼리 매칭)

• 쿼리(Query)와 키(Key) 간의 매칭 정확도를 향상함
• 쿼리와 키 간의 유사도 계산 방식을 최적화.
• 희소 조회(Sparse Lookup) 알고리즘 개선으로 더 정확한 Top-k 선택.

개선전 메모리 레이어의 메커니즘과 개선된 메커니즘

 

1. 기존 메모리 레이어 (왼쪽 패널)

 

구성 요소:

1. Query MLP: 입력 쿼리를 학습 가능한 표현으로 변환.
2. Sub Key1, Sub Key2: 두 개의 키를 사용해 Top-k 조회 수행.
3. Top-k & Softmax: 가장 관련성 높은 Top-k 키 선택 후 Softmax로 가중치 계산.
4. Values: 가중합을 통해 최종 출력 생성.
5. RMSNorm: 정규화를 통해 안정적인 학습 보장.
6. Self-Attention: 정규화 후 Self-Attention 단계로 전달.

한계:

• Top-k 조회 및 Softmax 과정에서 메모리 병목 현상 발생 가능.
• Query와 Key 간의 상호작용이 충분히 최적화되지 않을 수 있음.
• 학습 안정성 문제 발생 가능.

2. 개선된 메모리 레이어 (오른쪽 패널)

 

추가된 구성 요소:

1. W1 및 W2: 추가된 선형 변환(Linear Transform) 레이어로, 쿼리와 키 간의 상호작용을 더 깊고 유연하게 만듦.
2. SiLU (Sigmoid Linear Unit): 비선형 활성화 함수로, 정보 흐름을 더 세밀하게 조정.
3. Sub Key1, Sub Key2: 동일하게 Top-k 조회 및 Softmax를 통해 선택된 값을 가중합.
4. Values: Top-k 선택된 값들로 최종 가중합 수행.
5. 잔차 연결(Residual Connection): 출력이 이전 단계와 더해지며, 정보 손실을 최소화.
6. RMSNorm: 정규화를 통해 학습 안정성 유지.
7. Self-Attention: 개선된 메모리 레이어 출력이 Self-Attention으로 전달됨.

개선된 점:

• W1과 W2 추가: Query와 Key 간 상호작용을 더 세밀하게 조정.
• SiLU 활성화 함수: 학습 안정성을 높이고, 비선형성을 추가하여 더 정교한 매칭 수행.
• 잔차 연결(Residual Connection): 그래디언트 손실 문제를 완화하고, 학습 안정성 개선.
• 더 나은 정규화: RMSNorm이 여러 단계에서 사용되어 정규화 효과 극대화.

결론

• 왼쪽 (기존 레이어): 기본적인 Query-Key-Value 매칭을 수행하지만 학습 안정성 및 확장성에 한계가 있음.
• 오른쪽 (개선된 레이어): 추가적인 W1, W2 레이어와 SiLU 활성화, Residual Connection을 통해 더 나은 성능, 안정성, 확장성을 제공함.

 

---

4. Model ablations

이 표는 다양한 모델 아키텍처(Dense, MOE, PEER, Memory, Memory+) 간의 성능을 비교한 표인데 메모리 레이어가 실험에 따라 전반적으로 나은 성능을 보여준 것을 볼 수 있다. 

 

 

 

 

결과 분석

 

모델 유형별 경향성: 

1. Dense:
   • 파라미터 대비 성능이 낮은 편.
2. MOE:
   • 파라미터는 크지만 효율성은 Mixed.
   • OBQA와 같은 특정 태스크에서만 두드러짐.
3. PEER:
   • 대부분의 태스크에서 MOE보다 안정적인 성능을 보여줌.
4. Memory:
   • Dense 및 MOE보다 더 나은 성능을 보여줌.
5. Memory+:
   • 모든 벤치마크에서 가장 뛰어난 성능을 기록함.
   • 특히 NQ, TQA, OBQA 등에서 두드러진 성능 향상.

파라미터 확장에 따른 성능 변화

• Memory+ 모델은 파라미터 크기가 증가할수록 성능이 지속적으로 향상됨.
• 예: Memory+ 4m → 16m → 64m
  • NQ: 9.82 → 20.14 → 20.78
  • OBQA: 27.80 → 26.80 → 30.00

결론

1. 메모리 레이어(Memory, Memory+)는 Dense 및 MOE 모델보다 전반적으로 더 나은 성능을 보여줌.
2. Memory+ 모델은 파라미터를 확장할수록 성능이 선형적으로 향상됨. (섹션 3-1.4에 있는 기존 메모리 레이어를 개선한 모델)
3. 특히 사실적 질문 응답(Factual QA)과 같은 태스크에서 두드러진 성능을 보임.
4. Memory+는 계산량을 크게 증가시키지 않으면서 더 나은 성능을 제공함.

 

5. Conclusion

논문내에는 따로 결론 관련한 내용이 없는 것 같아 전체적인 흐름을 조합해보자면 기존 RAG와는 다른 방법(외부 문서 기반으로 훑는 것이 아닌 학습을 통해 데이터를 확인 하는 것 같음)으로 할루시네이션 방지 및 긴 문장들을 입력 시키려는 방법이지만,

 

학습 이후에서는 모델이 계산에 개입하는 것이 아닌 고정된 상태로 학습이 완료된 Key-Value 기반 테이블로 조회하며 빠른 추론 능력을 보여주는 결과를 얻는 것으로 파악됨