Generative Representational Instruction Tuning

Abstract

자연어 처리에서 생성(generation)과 임베딩(embedding) 작업을 모두 다룰 수 있는 새로운 방법인 GRIT(Generative Representational Instruction Tuning)에 대해 설명하고 있습니다.

기존의 언어 모델들은 주로 생성이나 임베딩 중 한 가지 작업에서만 좋은 성능을 보였습니다. 하지만 GRIT은 명령어(instruction)를 통해 생성과 임베딩 작업을 구분함으로써, 하나의 대형 언어 모델로 두 가지 작업을 모두 처리할 수 있게 합니다.

GRIT을 적용한 GRITLM 7B 모델은 MTEB(Massive Text Embedding Benchmark)에서 새로운 최고 성능을 달성했고, 다양한 생성 작업에서도 같은 크기의 다른 모델들을 능가했습니다. 더 큰 규모인 GRITLM 8X7B 모델은 시도한 모든 개방형 생성 언어 모델보다 우수한 성능을 보였고, 동시에 임베딩 모델 중에서도 최고 성능을 유지했습니다.

또한, GRIT을 사용하면 생성이나 임베딩 중 한 가지에만 특화된 학습을 해도 성능 손실 없이 두 작업을 통합할 수 있다는 사실을 발견했습니다. 이를 통해 GRIT은 긴 문서에 대한 RAG(Retrieval-Augmented Generation) 작업 속도를 60% 이상 향상시킬 수 있었는데, 이는 별도의 검색 모델과 생성 모델이 더 이상 필요하지 않기 때문입니다.

요약하면, GRIT은 하나의 언어 모델로 생성과 임베딩 작업을 모두 효과적으로 처리할 수 있는 새로운 방법으로, 자연어 처리 분야에 큰 기여를 할 것으로 기대됩니다.

1 Introduction

인공지능 분야에서 오랫동안 추구해온 목표 중 하나는 다양한 작업을 잘 수행하는 단일 모델을 만드는 것입니다. 최근 대형 언어 모델(LLM)이 이를 위한 유망한 방향으로 떠오르고 있지만, 임베딩을 사용하는 작업은 거의 무시되어 왔습니다. 

저자들은 GRIT(Generative Representational Instruction Tuning)이라는 새로운 방법을 소개합니다. GRIT은 명령어를 통해 생성과 임베딩 작업을 구분함으로써, 하나의 LLM이 두 가지 작업을 모두 잘 처리할 수 있도록 합니다. 이는 생성과 임베딩 작업을 통합하는 것으로, 세 가지 이점을 가집니다:

1. 성능: GRIT 모델은 생성 전용 또는 임베딩 전용 모델과 동등한 성능을 보이며, 일부 작업에서는 더 우수한 성능을 보입니다.

2. 효율성: 생성 모델과 임베딩 모델을 함께 사용하는 시나리오(예: RAG)에서, GRIT은 계산을 캐시하여 순전파 횟수를 절반으로 줄일 수 있습니다. 이는 긴 문서에 대해 60% 이상 빠른 RAG 추론을 가능하게 합니다.

3. 단순성: GRIT을 사용하면 API 제공자가 별도의 생성 및 임베딩 엔드포인트를 제공할 필요가 없어져 인프라 요구사항이 크게 단순화됩니다.

GRIT의 주요 단점은 두 가지 목적 함수로 훈련하기 때문에 더 많은 계산이 필요하다는 것이지만, 사전 훈련에 비해 미세 조정이 저렴하기 때문에 이 문제보다는 이점이 훨씬 크다고 합니다. 따라서 저자들은 명령어를 따르는 언어 모델을 구축하는 실무자들이 미세 조정 중에 GRIT을 채택할 것을 권장합니다.

2 GRIT

GRIT(Generative Representational Instruction Tuning)이 어떻게 임베딩 작업과 생성 작업을 하나의 모델로 통합하는지 자세히 설명하겠습니다.

GRIT은 사전 훈련된 대형 언어 모델을 임베딩 명령어 데이터와 생성 명령어 데이터로 미세 조정합니다. 

임베딩 데이터의 경우, 배치 내 부정 샘플을 사용하는 대조 목적 함수로 손실을 계산합니다 (수식 1). 

이 때, 양방향 어텐션과 평균 풀링을 사용하여 입력 샘플의 최종 은닉 상태를 평균내어 표현을 얻습니다.

생성 데이터의 경우, 언어 모델링 목적 함수를 사용하여 모델이 다음 토큰을 예측하도록 합니다 (수식 2). 

중요한 고려사항은 생성 손실을 샘플 수준에서 집계할 것인지 토큰 수준에서 집계할 것인지 입니다. 샘플 수준 집계는 토큰 수에 관계없이 각 샘플에 동일한 가중치를 부여하는 반면, 토큰 수준 집계는 각 토큰에 동일한 가중치를 부여하여 토큰이 많은 샘플이 더 중요해집니다. 전자는 짧은 생성에 치우친 모델로, 후자는 긴 생성을 선호하는 모델로 이어질 수 있습니다.

최종적으로 임베딩 손실과 생성 손실을 선택적 가중치와 함께 합산하여 GRIT 손실을 계산합니다 (수식 3). 

이 때, 임베딩 샘플 수(M)와 생성 샘플/토큰 수(N)가 다를 수 있어 임베딩 배치 크기를 크게 늘리면서 생성 배치 크기는 고정할 수 있습니다. 이는 계산량 증가의 대가로 성능 향상을 가져올 수 있습니다.

3 Experiments

3.1 Setup

GRITLM은 Mistral 7B와 Mixtral 8x7B 모델을 기반으로 E5와 Tülu 2 데이터셋을 사용하여 파인튜닝된 모델입니다. E5에는 과학 데이터를 늘리기 위해 S2ORC가 추가되었고, Tülu 2에서는 모델 출처와 관련된 답변이 포함된 맞춤형 프롬프트가 제거되었습니다. GRITLM 7B는 임베딩 데이터에 배치 크기 2048, 생성 데이터에 256을 사용하고 총 1253 스텝 동안 학습되었으며, GRITLM 8X7B는 컴퓨팅 제한으로 임베딩 배치 크기가 256이었습니다. 학습 중 메모리 요구량을 줄이기 위해 여러 전략이 사용되었습니다.

3.2 Main Results

주요 결과로는 GRITLM 7B가 MTEB에서 모든 이전 오픈 모델을 능가하면서 70억 매개변수 크기의 생성 모델들보다 성능이 좋았으며, GRIT 모델만이 임베딩과 생성 모두에서 최고 성능을 보였습니다. Llama 70B를 임베딩에 사용하면 MTEB 점수가 35.6에 불과한 반면, GRITLM은 거의 두 배인 66.8의 SOTA 성능을 보이면서 생성 작업에서도 Llama 70B보다 20% 이상 우수했습니다. GRITLM 8X7B는 공개 모델 중 생성 성능 평균이 가장 높았지만, 임베딩 성능은 GRITLM 7B보다 약간 낮았습니다. 임베딩만 하거나 생성만 하는 변형 모델은 반대 작업에서는 거의 랜덤 수준의 성능을 보여, GRIT 접근법을 통한 공동 최적화가 임베딩과 생성 모두에서 강력한 성능을 달성하는 데 중요함을 알 수 있었습니다.

세부 실험 결과로는 GRITLM이 임베딩 전용 및 생성 전용 변형 모델과 성능이 일치했으며, GRITLM 7B는 쿼리당 상위 10개 문서를 재순위화하여 자체 임베딩 성능을 개선할 수 있었습니다. 임베딩 작업에서는 few-shot 학습이 도움이 되지 않고 오히려 성능이 약간 저하되었습니다. 마지막으로, GRITLM을 KTO를 사용해 정렬 튜닝하면 임베딩 성능은 약간 하락하지만 생성 성능 평균은 더 강력해져, KTO로 튜닝된 모델은 생성 성능이 더 중요한 사용 사례에 유용할 수 있음을 보여주었습니다.

3.3 Ablations

GRITLM은 사전 학습된 언어 모델을 기반으로 만들어졌습니다. 이 모델은 인과적 어텐션(causal attention)으로 학습되었는데, 이는 정보가 한 방향으로만 전파되도록 제한합니다. 하지만 양방향 어텐션(bidirectional attention)을 사용하면 모델이 미래 토큰으로부터 얻은 정보를 기반으로 이전 토큰의 표현을 조정할 수 있어 더 나은 성능을 보일 수 있습니다. 실험 결과, 양방향 어텐션으로 모델을 적응시키는 것이 임베딩 성능을 가장 높였습니다. 통합 모델(unified model)은 임베딩 전용 변형보다 성능이 훨씬 우수했지만, 생성 전용 변형보다는 성능이 떨어졌습니다. 

GRITLM 접근 방식은 모든 생성 언어 모델에 일반화될 수 있습니다. GPT-J 6B, Llama 2 7B, Mistral 7B 등 다양한 사전 학습 모델로 초기화하여 성능을 비교했을 때, Mistral 7B가 임베딩과 생성 작업 모두에서 가장 좋은 성능을 보였습니다. 

생성 데이터셋으로는 Tülu 2가 가장 좋은 성능을 보였고, 임베딩 데이터셋으로는 E5가 가장 좋았습니다. E5는 GPT-4로 생성된 우수한 하드 네거티브와 다양한 작업들 덕분에 성능이 좋았던 것으로 보입니다. 

모델의 임베딩 차원을 줄이기 위해 임베딩 헤드를 추가하는 실험도 했는데, 이는 일부 사용 사례에 적합할 수 있지만 최종 모델에서는 성능을 최대화하기 위해 사용하지 않았습니다. 

배치 크기를 늘리면 임베딩 성능이 향상되었고, 혼합 정밀도(mixed precision)로 학습하는 것이 더 효율적이었습니다. 임베딩 학습 시 배치 내 네거티브(in-batch negatives)는 동일한 데이터셋에서 가져오는 것이 좋았고, 포맷은 Tülu 2 방식이 Zephyr β보다 우수했습니다.

토큰 수준과 샘플 수준의 생성 손실(generative loss)을 모두 사용하되, 표현 손실(representation loss)에 더 높은 가중치를 주는 것이 성능 향상에 도움이 되었습니다. 특히 토큰 수준 손실을 많은 샘플에 걸쳐 적용하는 것이 긴 생성을 유지하는 데 중요했습니다.

4 RAG with GRIT

GRITLM은 임베딩과 생성을 통합함으로써 검색 강화 생성(Retrieval-Augmented Generation, RAG)을 단순화합니다. 캐싱을 통해 순전파(forward pass) 횟수를 줄일 수 있는데, 크게 다음과 같은 방법이 있습니다.

(a) 쿼리 캐싱(Query Caching): 전통적인 RAG에서는 쿼리를 임베딩 모델과 생성 모델에 모두 통과시켜야 합니다. 쿼리 캐싱에서는 임베딩 순전파에서 얻은 키-값 상태를 캐시하고 생성 단계에서 재사용합니다. 이는 GRITLM이 동일한 모델이라는 점을 활용한 것입니다. 이를 통해 쿼리의 한 번의 순전파에 해당하는 연산을 절약할 수 있습니다.

(b) 문서 캐싱(Doc Caching): 문서(D)를 캐시합니다. 인덱스를 생성할 때 모든 문서의 키-값 상태를 저장하고 인덱스에 추가합니다. 즉, 인덱스는 문서 임베딩과 키-값 상태로 구성됩니다. 추론 시에는 여전히 임베딩 유사도를 기반으로 검색하지만 인덱스는 텍스트 대신 키-값 상태를 반환합니다. 이 키-값 상태를 모델에 제공하여 재계산을 피할 수 있습니다. 이는 추론 시 모든 문서에 대해 순전파를 한 번씩 절약하는 효과가 있지만 저장 공간이 더 필요합니다.

(c) 쿼리-문서 캐싱 / 문서-쿼리 캐싱: 쿼리 캐싱과 문서 캐싱을 결합하여 추론 비용을 더 절약할 수 있습니다. 그러나 결합하면 전통적인 RAG와 차이가 발생합니다. 전통적인 RAG에서는 쿼리나 문서 중 하나가 다른 하나에 의존하지만, 둘 다 캐시되면 서로에 의존하지 않게 됩니다. 쿼리가 문서 앞에 오면 쿼리-문서 캐싱, 문서가 먼저 오면 문서-쿼리 캐싱이라고 합니다.

이러한 캐싱 기법을 통해 GRITLM은 RAG를 더 효율적으로 만들 수 있습니다.

 

캐싱은 긴 시퀀스에 대해 CPU와 GPU 모두에서 RAG보다 상당한 속도 향상을 가져옵니다. 그러나 캐시되는 토큰 수가 적으면(예: 250개) 속도 향상이 미미합니다.

4000개 토큰에 대해 쿼리 캐싱은 CPU에서 54%, GPU에서 33% 더 빠르고, 문서 캐싱은 CPU에서 63%, GPU에서 31% 더 빠릅니다. 4000개 이상의 토큰에 대해서는 속도 향상이 더 클 것입니다. 그러나 반대 샘플에서는 속도 차이가 거의 없습니다. 샘플 B에서는 문서 캐싱이 4000개 토큰을 캐시하지만, 샘플 A에서는 1개 토큰만 캐시하기 때문에 속도 향상이 없습니다.

따라서 문서가 매우 길 것으로 예상되면 문서 캐싱을, 쿼리가 매우 길 것으로 예상되면 쿼리 캐싱을 사용하는 것이 좋습니다. 실제 상황에서는 입력 길이를 확인하여 캐싱 모드를 전환할 수 있습니다.

쿼리-문서 캐싱과 문서-쿼리 캐싱은 문서와 쿼리 모두를 캐시하므로 샘플 A와 B 모두에서 큰 속도 향상을 가져옵니다. 전반적으로 CPU에서의 속도 향상이 더 큽니다. GPU는 전체 시퀀스를 병렬로 처리할 수 있어 캐싱의 이점이 상대적으로 작기 때문입니다.

대부분의 RAG 설정에서는 모든 문서의 임베딩을 미리 계산하여 인덱스로 저장합니다. 전통적인 RAG에서는 인덱스로 문서 ID를 찾은 후 문서 텍스트를 가져와 생성 모델에 전달해야 하므로 여전히 문서 자체를 저장해야 합니다. 그러나 문서 캐싱 변형에서는 문서를 저장할 필요가 없지만 키-값 상태를 인덱스와 함께 저장해야 합니다. 키-값 상태는 많은 저장 공간을 차지하지만, 인덱스와 달리 디스크에 완전히 오프로드할 수 있어 메모리에 보관할 필요가 없습니다.

5 Discussion

GRITLM은 텍스트 임베딩과 생성을 단일 모델로 통합한 최초의 모델로, 모든 텍스트 기반 언어 문제를 강력한 성능으로 해결할 수 있습니다. 하지만 개선되거나 통합되어야 할 인접 방향들이 여전히 많이 남아 있습니다.

(a) 다국어 지원: GRITLM은 비영어 언어의 임베딩과 생성도 가능하지만, 데이터와 아키텍처 변경을 통해 비영어 작업에서 성능을 크게 향상시킬 수 있을 것입니다.

(b) 멀티모달 지원: 많은 임베딩 및 생성 문제는 순수한 텍스트 기반이 아닙니다. 이미지와 텍스트의 공동 임베딩, 생성적 이미지 캡션 생성, 이미지-텍스트 쌍 분류, 모든 텍스트 문제의 음성 버전 등이 그 예입니다. 이러한 문제들이 텍스트 임베딩 및 생성처럼 쉽게 통합될 수 있는지는 더 연구해 볼 필요가 있습니다.

GRIT가 작동하는 이유: GRIT는 임베딩과 생성 작업을 단일 모델로 통합하면서도 어느 쪽에서도 성능 손실이 없는데, 이는 놀라운 결과일 수 있습니다. 우리의 결과는 생성적 언어 모델링과 텍스트 임베딩이 동전의 양면과 같다는 것을 확인해 줍니다. 두 작업 모두 자연어에 대한 깊은 이해를 필요로 하며, 그 이해를 표현하는 방식만 다를 뿐입니다. 아마도 우리의 통합 모델에는 최종 표현을 평균 풀링 및 후속 임베딩 작업에 유용하게 만들거나 언어 모델링 헤드 및 후속 생성 작업에 맞춰 프라이밍하는 스위치 역할을 하는 소수의 매개변수가 포함되어 있을 것입니다.

GRITLM을 사용한 RAG 최적화: GRITLM은 검색 및 생성 성능이 우수하므로 단일 모델에서 검색기와 판독기 역할을 모두 수행할 수 있습니다. 따라서 둘 중 하나를 최적화하면 다른 쪽의 매개변수도 변경됩니다. 이는 검색기와 판독기의 공동 최적화를 크게 단순화할 수 있는 잠재력이 있습니다. 예를 들어, 다음 토큰 목적 함수만 사용하여 검색기가 관련 없는 맥락을 제공하는 것과 판독기가 주어진 맥락을 잘못 사용하는 것을 동시에 페널티를 부과할 수 있습니다.

6 Related Work

텍스트 임베딩과 생성의 역사는 통합의 역사였습니다. 

과거 임베딩 모델은 단어 표현에 초점을 맞추어 문장이나 문단으로 일반화하는 데 어려움을 겪었습니다. 이후 InferSent, SBERT 등의 모델이 등장하여 맥락을 고려함으로써 단어와 문장 모두를 양질로 임베딩할 수 있게 되었습니다. 그러나 강력한 성능을 위해서는 대칭 작업과 비대칭 작업에 별도의 모델이 필요했습니다. 

최근 임베딩 모델은 프롬프트에서 대칭 및 비대칭 작업을 구분함으로써 단일 모델로 통합할 수 있게 되었습니다. 또한 프롬프트에 상세한 지침을 포함시킴으로써 사실상 모든 임베딩 작업을 단일 모델로 통합할 수 있게 되었습니다.

생성 모델은 과거 번역이나 질의응답 등 단일 작업에 맞춰져 있었습니다. 이후 McCann 등은 여러 생성 작업을 질의응답으로 캐스팅하여 단일 모델 내에서 통합했지만, 성능이 제한적이었고 임의의 작업으로 일반화되지 않았습니다. 대규모 자기 지도 학습 기반 사전 학습을 통해 단일 대형 언어 모델(LLM)을 사실상 모든 생성 작업에 사용할 수 있게 되었습니다. 

그러나 신중한 프롬프팅 없이 LLM을 사용하면 성능이 떨어지는 경우가 많았습니다. 지침에 대해 LLM을 미세 조정하는 것이 모델 사용을 크게 용이하게 하여 강력한 결과로 임의의 생성 작업에 적용할 수 있는 방법으로 부상했습니다.

임베딩과 생성 모델의 두 흐름은 각각 자신의 흐름 내에서 모든 작업을 처리하는 단일 모델로 통합되었습니다. 두 흐름을 임베딩과 생성 모두에 대해 모든 작업을 처리하는 단일 모델로 통합하는 것은 일반적인 다중 작업 모델을 향한 자연스러운 다음 단계입니다. 

LLM은 생성뿐만 아니라 텍스트 임베딩에서도 유망한 성과를 보여주었습니다. SGPT는 그 방향으로 초기 작업이었습니다. SGPT는 대형 언어 모델의 매개변수 중 0.01%만 변경하여 성능이 좋은 임베딩을 생성하도록 적응시켰습니다. 따라서 한 흐름에서 다른 흐름으로 전환하려면 이 작은 양의 매개변수만 변경하면 됩니다. 

그러나 SGPT는 여전히 별도의 대칭 및 비대칭 모델을 필요로 했으며 임베딩 작업의 전체 범위를 고려하지 않았습니다. GRITLM은 이러한 단점을 해결합니다. GRITLM은 편향을 전환할 필요가 없고, 지침을 활용하여 대칭 또는 비대칭 사용 사례를 처리하며, 임베딩 및 생성 작업의 전체 범위를 고려합니다.

7 Conclusion

본 연구에서는 텍스트 임베딩과 생성, 그리고 모든 텍스트 기반 언어 문제를 단일 모델인 GRITLM으로 통합하기 위한 GRIT 접근법을 제시합니다. 

GRITLM 7B는 공개 모델 중 Massive Text Embedding Benchmark에서 SOTA 성능을 달성하면서 동시에 7B 크기까지의 모든 생성 모델을 능가합니다. 특히 동일한 임베딩 전용 및 생성 전용 변형 모델의 성능과 일치하여 성능 손실 없이 두 흐름을 통합할 수 있습니다. 

추론 시 5B 매개변수만 추가하면 GRITLM 8X7B는 70B 매개변수를 가진 Llama 2 기반의 훨씬 더 큰 모델을 포함하여 우리가 시도한 많은 모델 중 가장 뛰어난 공개 생성 언어 모델입니다. 다른 생성 모델과 달리 GRITLM 8X7B는 GRIT 접근법 덕분에 매우 강력한 임베딩 성능도 자랑합니다. 

통합된 기능 덕분에 GRITLM은 재순위화 파이프라인에서 Bi-Encoder와 Cross-Encoder 모두로 사용될 수 있어 16개 중 15개의 검색 데이터셋에서 성능 향상을 이끌어냅니다. 

또한 광범위한 절제 실험을 통해 임베딩 및 생성 모델 연구자들을 위한 주요 통찰력을 밝혀냈습니다: 임베딩을 위한 인과적 언어 모델은 양방향 어텐션과 평균 풀링으로 미세 조정되어야 하고, 미세 조정 전 언어 모델의 임베딩 성능은 미세 조정 후 임베딩 성능을 예측하지 않으며, BF16 혼합 정밀도에서 임베딩 모델을 미세 조정하면 FP32 미세 조정과 일치할 수 있고, BF16 미세 조정 중에는 풀링 및 유사성 계산을 FP32 정밀도로 수행해야 하며, 생성 모델은 일종의 토큰 수준 손실로 지침 조정되어야 합니다.

마지막으로 RAG의 예를 사용하여 GRIT가 이 분야를 단순화한다는 것을 보여줍니다. 검색기와 판독기를 단일 모델로 통합함으로써 GRITLM은 GRIT Doc Caching을 사용하여 성능 손실 없이 긴 시퀀스에 대해 60% 이상의 추론 속도 향상을 이끄는 캐싱 작업을 수행할 수 있게 합니다.