[Paper Review] Generalization Through Memorization: Nearest Neighbor Language Models

Generalization Through Memorization: Nearest Neighbor Language Models

arxiv pdf link for kNN-LM

0. Abstract 🎬

kNN-LM 은 pre-train 된 LM을 선형적으로 kNN 알고리즘과 결합하여 확장한 모델이다. pre-trained LM을 이용해서 input 데이터가 latent space 로 임베딩 되게 되는데, 이 latent space 상의 벡터 간의 거리를 통해서 가장 가까운 k 개의 후보를 정하게 된다. 이는 임의의 데이터셋 (including training data) 을 통해서 가능하다.

이 방식을 Wikitext-103LM 에 적용함으로써, 이 논문에서 소개하는 모델은 SOTA를 달성했으며, 추가적인 training 없이도 15.79의 perplexity로 2.9 point나 줄이는 효과를 보였다. 또한, 이 접근법은 더 큰 훈련 데이터셋, 그리고 다른 domain 으로의 적용 역시 효과적으로 수행할 수 있었다.

질적으로는, 이 모델은 생소한 표현들에 대해서 더욱 효과적인 모습을 보였고, 특히 factual knowledege 에 대해서 효과적이였다. 동시에 이 연구는 LM 의 근본적인 task 인 next token prediction 보다 sequences 간의 similarity 를 학습하는 것이 더 효과적인 접근 방식임을 의미하기도 한다.

1. Introduction ☕️

Language Model 은 일반적으로 아래의 2가지 task를 목표로 한다.

1. 즉 문장의 prefix를 n차원 벡터로 나타낸다. (정확히는 고정 크기의 representation)

2. 이렇게 만들어진 latent space에서의 값을 이용해서 다음 단어를 예측한다.

본 논문에서는 첫번째 task가 두번째 task 보다 쉬운 task라는 가정 하에 접근했다. 예를 들어, “Dickens is the author of” 라는 문장과 “Dickens wrote” 라는 문장을 보았을 때, 그 후에 올 단어를 예측하지 못하더라도 두 문장이 같은 뜻을 내포하고 있음은 누구나 알 수 있다. 실험적으로도, prefix embedding에 대해 kNN을 적용시킨 결과, 성능이 향상됨을 통해 LM이 첫번째 task에 더 효과적이라는 강력한 증거를 제시한다.

3-billion개의 token을 모델의 학습 데이터로 사용하는 것보다, 100-million개의 token을 이용해 학습하고 3-billion개의 token을 가지는 dataset(documents)을 이 모델에 적용하는 것이 더 높은 성능을 보였다. 이는 곧 Large dataset을 사용하는 LM에 대한 새로운 방향성을 제시한다. 비슷하게, 단순히 datastore 에 다른 domain의 데이터를 삽입하는 것만으로도 multiple domain에서도 효과적인 성능을 보였다.

마지막으로, 이 모델은 명시적 기억에 대한 접근 (datastore이라는 명시적인 데이터) 을 통해 long-tail patterns (예를 들어 Factual Knowledge) 에 대해서 더욱 효과적인 것을 발견했다.

2. Nearest Neighbor Language Modeling 🧐

LM은 기본적으로 sequence에 대한 확률을 할당한다. 다시 말해 $c_t = (w_1, \cdots , w_{t-1})$ 라는 context (sequence)가 주어져 있을 때, LM (autoregressive 한)은 $p(w_t|c_t)$를 계산해낸다.

kNN-LM은 pre-trained LM을 이용하여 nearest-neighbors를 검색(retrieval) 하여 augument 하는 과정을 내포하며, 이 datastore에는 key-value 형태의 context-target 쌍들이 저장되어 있다. (see Figure 1)

Datastore

$f()$를 context $c$를 fixed-length vector로 mapping해주는 함수라고 가정해보자. 만약에 i번째 training data인 $(c_i, w_i) \in \mathcal{D}$가 주어졌을 때, 우리는 datastore을 다음과 같은 key-value 집합으로 정의할 수 있다.

\[(\mathcal{K}, \mathcal{V}) = \{(f(c_i), w_i)|(c_i, w_i)\} \in \mathcal{D}\]

Inference

이 모델은 input context $x$가 주어졌을 때, $f(x)$를 통해 $p_{LM}(y|x)$에 대한 확률 분포를 계산하게 된다. 동시에 datastore 에 $f(x)$를 이용해 query를 보내게 되는데, distance function인 $d()$를 통해 k-nearest neighbors에 해당하는 집합 $\mathcal{N}$을 생성한다. (본 논문에서는 distance function을 $L^2$ distance로 정의했다.)

그 후에는 softmax에 negative distance를 넣음으로써, 아래의 확률을 계산해낼 수 있다. (거리가 가까울수록 높은 확률의 정확도를 보인다.)

\[p*{kNN}(y|x) \propto \sum_{(k*i, v_i) \in \mathcal{N}}{\mathbb{1}_{y=v_i} \, exp(-d(k_i, f(x)))}\]

이를 선형적으로 기존 LM에 적용하게 되면, $\lambda$ 변수를 이용해서 다음과 같이 최종 probability를 정의할 수 있다.

\[p(y|x) = \lambda * p_{kNN}(y|x) + (1 - \lambda) * p_{LM}(y|x)\]

Implementation

한 가지 문제점은, Datastore 이 billion 개의 데이터를 내포하고 있어 computationally intensive 하다는 것이다. 이를 극복하기 위해 FAISS라는 open source library 를 사용하여 고차원 상에서의 kNN을 효율적으로 계산하게 된다. 추가적으로, $L^2$ 이외에 inner product distance 라는 distance function 또한 존재하는데, 이 모델에서는 $L^2$ 방식이 더 높은 성능을 보였다.

Related Cache Models

이전의 비슷한 접근 방식에서는, recent past에 대한 caching 을 통해서 최근의 데이터에 대해 더욱 효과적으로 계산하는 방식 또한 존재했다. 하지만 최근의 정보를 copy할 수 있는 self-attention 기법을 가진 Transformer 모델이 등장하고 나서, 이 방식은 인기를 잃게 되며, 얻을 수 있는 이익 또한 줄어들었다. 본 논문에서는 training data에 대한 명시적 기억을 위해 오로지 training data에 대해서만 caching하는 방식을 택하여 비슷한 cases에 대한 효율을 증대시켰다.

3. Experimental Setup 🥽

Data

단순히 dataset에 대한 설명이기에 별도의 설명은 생략한다.

Model Architecture

kNN-LM 모델은 fixed-size context representations을 생성하는 모델이라면 모두 호환이 가능하다. 이 모델은 현재 (당시) SOTA를 기록했던 Decoder-only Transformer을 사용한다. kNN-LM 모델은 모델의 기본 모델인 LM에 대한 훈련을 시행하지 않기 때문에, 기존 아키텍쳐와 최적화 방식을 그대로 사용했다.

16 layer, each with 16 self-attention heads, 1024 dimensional hidden states, 4096 dimensional feedforward layers, 247M parameters 등을 사용하며 추가 정보는 아래와 같다.

Evaluation

이 LM들은 negative log-likelihood 를 loss function으로 사용했으며, 모델의 평가 기준으로써 perplexity를 사용했다. (살짝 첨언하자면, 최근 논문인 Mamba의 reject 원인으로 perplexity가 평가 기준이 되는 것은 정당하지 않다는 글을 본 것 같은데, 이때는 기준이 조금 달랐나보다.)

kNN-LM

Computational Cost

추가적인 Training 없이도 모델을 구현할 수 있지만, key-value 형태의 datastore 을 생성하기 위하여 1 epoch 정도의 시간이 소요된다. 또한 key들이 저장된 후 WIKI-103M의 캐시를 CPU에 적용하는 데에 2시간이 소요되며, 1024개의 NN을 구하는 데에 25분 정도가 소요된다. 물론 데이터셋의 크기에 선형적으로 비례하여 시간이 증가하지만, 이는 쉽게 병렬화가 가능하며, GPU의 사용이 필요하지 않다.

4. Experiemtents 🔬

4.1 Using the Training data as the Datastore

기존의 SOTA 방식과 비교하여, 본 논문의 kNN-LM 이 얼마나 높은 성능을 보였는지를 실험해보았다. 여기서 Training data 그대로 Datastore에 대입했다. 기존의 SOTA와 비교하여 18.65 에서 16.12로 새로운 SOTA를 달성했다.

추가적으로, WIKI가 caching에 유독 좋은 경우를 대비하여 BOOKS corpus를 이용하여 같은 실험을 반복해 보았다. 그 결과는 아래와 같다.

4.2 More data without Training

이번에는 Training dataset과 Datastore을 분리하여, 서로 다른 데이터셋을 이용했을 때에도 효과가 있는지를 확인해보았다.

위 그림과 같이, WIKI-3B와 WIKI-100M dataset으로 실험해 보았는데, 당연하게도 더 큰 데이터셋인 WIKI-3B 을 이용해 학습했을 때 더 높은 성능을 보였다. 하지만, WIKI-100M으로 학습한 뒤에 WIKI-3B를 datastore로 활용했을 때, 그 성능이 기존의 LM을 능가하는 것으로 보아, kNN-LM 방식이 더욱 효율적이고 정확도가 높다는 것을 확인할 수 있었다.

또한, 위와 같이 kNN-LM의 datastore 크기에 대해서도 실험을 해 보았는데, 1.6B의 dataset만 사용했을 때 이미 Vanilla LM의 성능을 능가하였고, 3B의 경우에도 그러했다. 더 나아가 3B의 경우에도 perplexity의 감소도가 saturated(포화) 되지 않는 것을 보아, 더 큰 잠재성이 존재한다.

마찬가지로, 각 datastore 크기에 대해서 optimal 한 $\lambda$를 구해 보았을 때, 위 그림과 같이 결과가 나타났다.

4.3 Domain Adaptation

Domain Adaptation 실험을 위하여, WIKI-3B로 학습된 모델을 BOOK 을 dataset으로 inference 해보았다. 그 결과는 아래와 같다.

순수하게 datastore 없이 추론한 결과 매우 낮은 성능을 보였지만, BOOKS 를 datastore 로 활용하게 되면, perplexity가 15 가까이 떨어지는 것을 확인할 수 있다. 즉 target domain에 대한 datastore이 있다면, 충분히 다른 domain으로 적용이 가능함을 알 수 있다.

5. Tuninig Nearest Neighbor Search 🧪

Key Function

Similarity Search 를 위하여, prior context를 효과적으로 fixed-size representation 으로 변환하는 $f()$는 중요한 요소이다. 이를 실험해보기 위해, Transformer 구조의 다양한 부분을 후보로 실험을 진행하였다. (모두 Transformer의 마지막 layer 부분이다.)

실험에서 보는 것과 같이, FFN input after layer norm 부분이 가장 높은 성능을 나타내었고, 추가적으로 마지막 직전의 layer (second-last) 에도 실험을 해보았으나, 비슷한 경향의 점수지만 살짝 낮은 성능을 보였다.

이를 통해 FFN은 다음 토큰을 예측하는 것에, MHSA는 representation 에 더욱 효과적인 것이라고 추측해볼 수 있다.

Other elements (Number of Neighbors, Interpolation Param, etc)

그림에서 보다시피, k-NN 에서의 k 를 하나의 요소로, interpolation 변수를 하나의 요소로 하여 실험을 진행했다. 그 결과 k 가 늘어날수록 perplexity가 감소하며, 각 상황 (In-domain or Domain Adaptation)에 맞추어 적절함 $\lambda$ 값이 형성되는 것을 볼 수 있다.

또한, Similarity Function의 precision에 대해서도 실험을 진행했는데, $L^2$ distance에 대해서 full precision을 통해 연산함으로써 perplexity 가 16.5에서 16.06으로 감소하는 성능을 보였다.

6. Analysis 🧐

Qualitative Analysis

$p_{kNN}$ 이 왜 $p_{LM}$ 보다 높은 성능을 가지는지 이해하기 위하여 $p_{kNN}$ 이 더 나은 성능을 보였던 예시들을 살펴보면 다음과 같다.

예시에서 보다시피, kNN-LM model은 rare patterns 에 대해서 가장 높은 성능을 보였으며, 이는 곧 factual knowledge를 뜻한다. 특히 training set에 존재하거나 비슷한 context의 경우 더 높은 성능을 보였다.

이를 통해 parameters를 통해 지식을 implicit 하게 학습하는 것보다 explicit 하게, 즉 Nearest Neighbor을 찾아 학습하는 과정이 더욱 효율적이라고 볼 수 있다.

Simple vs Neural Representation

하지만, 이런 rare patterns (long-tail phenomena)는 단순히 N-gram model에서도 충분히 성능을 가질 수도 있다. 따라서 n-gram 모델을 사용하여 학습을 시켜보았는데, 결과는 위와 같이 n-gram model은 현저히 낮은 정확도를 나타내었다. 이를 미루어 보면, kNN-LM은 단순히 local context를 학습하는 것 이상으로, global context를 학습한다고 생각할 수 있다.

Implicit vs Explicit Memory

그렇다면 과연 Explicit Memory 가 Implicit Memory 보다 효과적일까? 본 논문에서의 datastore 을 LM 이 모두 외우는 것이 가능할까?

이를 실험해보기 위하여, Transformer 을 dropout 없이 학습시켜, datastore 의 모든 데이터를 학습하도록 했다. 그 결과 앞선 그림 같이, loss 가 아예 0으로 떨어진 것을 확인할 수 있고, 이는 곧 datastore 의 모든 데이터를 정확하게 학습했다는 의미가 된다.

따라서 이렇게 overfitting 된 LM과 explicit memory인 datastore를 비교하기 위하여 각각을 original LM 에 interpolate 하여 perplexity를 측정했는데, LM의 결과 0.1 향상, datastore 의 결과 1.9 향상으로 explicit memory가 더 뛰어난 성능을 보였다.

이 실험의 결과로, Transformer LM은 datastore의 모든 데이터를 외울 정도로 뛰어난 성능을 보였지만, 이 경우 generalize 성능이 떨어진다는 것을 추측해볼 수 있다.

8. Conclusion & Future Work 🎬

Related Work Section 은 중요한 부분이 아니라고 판단되어 제외하였습니다.

이렇듯 kNN-LM model은 기존의 standard LM을 능가하는 성능을 보였다. 이런 접근 방식은 임의의 NLP task 에 적용될 수 있다. 이 접근 방식의 성공은 인하여 context 간의 similarity 를 학습하는 것이 다음 토큰을 학습하는 것보다 쉬운 task임을 뜻하기도 한다. 추후에는 이런 simlarity function 을 명시적으로 학습하거나, datastore의 크기를 줄이는 것 등의 연구가 필요할 것이다.