초록
최근의 많은 연구들은 언어 모델에 검색 기능을 추가함으로써 입력 컨텍스트에 검색된 데이터를 포함시켜 모델 성능을 향상시키는 방식을 채택하고 있다. 이러한 접근 방식이 성공하기 위해서는 훈련 시간과 테스트 시간 모두에서 검색된 데이터를 입력에 포함시켜야 한다. 또한 입력 길이가 검색된 데이터 크기에 비례하여 선형적으로 증가하기 때문에, 현대의 트랜스포머 모델에서는 계산 및 메모리 비용이 제곱적으로 증가하게 된다. 이러한 복잡성을 피하기 위해, 우리는 테스트 시점에 표준 훈련 설정을 활용해 검색된 데이터 위에서 모델을 단순히 미세조정(fine-tune)하는 방식을 제안한다. 우리는 Pile 데이터셋의 텍스트 임베딩 기반으로 대규모 분산 인덱스를 구축하였다. 각 테스트 입력에 대해, 시스템은 해당 입력과 유사한 이웃들을 검색한 후, 이 이웃들의 텍스트를 기반으로 모델을 미세조정한다. 놀랍게도, 단지 20개의 이웃을 검색하고 각 이웃에 대해 단 한 번의 기울기 반복(gradient iteration)만 수행해도, Pile 내 20개 이상의 언어 모델링 과제에서 성능이 크게 향상된다. 예를 들어, 가장 가까운 이웃을 활용한 테스트 시점 학습은 GPT-2 소형 모델과 크기가 10배 이상 큰 GPT-Neo 모델 간의 성능 격차를 훨씬 더 좁히는 데 기여한다. 그러나 충분한 인덱스 품질과 규모가 필수적이다. 본 연구는 언어 모델링을 위한 테스트 시점 학습에 대한 첫 번째 기준 기준선(baseline)을 제시한다.