Atlas: Retriever-Reader 합동 사전 학습으로 110억 개의 파라미터만으로 5,400억 개의 파라미터 LLM을 능가하다

Atlas는 Izacard와 Grave가 자신들의 Fusion-in-Decoder 논문의 후속작으로 내놓은 모델로, 검색기(retriever)와 리더(reader)가 처음부터 함께 학습되는 완전한 합동 학습 시스템으로 FiD를 확장한 것입니다. 제가 지금 이 논문을 읽는 이유는 이것이 원조 RAG 논문부터 FiD를 거쳐 합동 학습된 검색 기술로 이어지는 아키텍처 계보를 완성하기 때문입니다. 이는 모든 장부 질의응답(QA) 시스템이 탐색해야 할 결정적인 영역입니다.

논문

"Atlas: Few-shot Learning with Retrieval Augmented Language Models" (Izacard et al., JMLR 2023)은 검색 증강 모델이 지식 집약적인 퓨샷(few-shot) 태스크에서 거대 파라미터 LLM과 대등한 성능을 낼 수 있는지 질문합니다. 핵심 기여는 Contriever 기반의 밀집 검색기와 T5 기반의 Fusion-in-Decoder 리더를 합동으로 학습시킨 정교한 사전 학습 검색 증강 시스템입니다. 핵심 통찰은 퓨샷 지식 성능을 견인하는 것이 아키텍처가 아니라 합동 사전 학습이라는 점입니다. 이 시스템은 상위 20개의 문서를 검색하여 각각 인코더에서 독립적으로 인코딩한 다음, 저자들의 2021년 논문과 동일한 FiD 설계를 따라 디코더의 크로스 어텐션(cross-attention)에서 이를 융합합니다.

주요 개념

• Atlas-11B는 단 64개의 학습 예시만으로 Natural Questions에서 42.4%의 정확도를 달성하여, 50배 더 적은 파라미터를 사용하면서도 PaLM(5,400억 개의 파라미터)을 약 3포인트 차이로 앞섰습니다.
• TriviaQA(64-shot)에서 Atlas-11B는 필터링된 세트에서 74.5%, 필터링되지 않은 히든 테스트에서 84.7%에 도달하며, 검색 컴포넌트가 제한적인 태스크 감독 학습을 강력하게 보완함을 보여주었습니다.
• 네 가지 검색기 학습 목표가 평가되었습니다: Attention Distillation (ADist), EMDR2 (검색된 문서를 잠재 변수로 처리), Perplexity Distillation (PDist), 그리고 LOOP (leave-one-out). 이들 간의 성능 차이는 미미하며, 계산 효율성을 위해 PDist가 채택되었습니다.
• 라벨이 지정되지 않은 텍스트에 대한 합동 사전 학습이 가장 큰 요인입니다. 모든 검색 증강 사전 학습 구성이 검색 증강 미세 조정(fine-tuning)만 수행한 베이스라인을 압도했습니다.
• 문서 인덱스는 모델 재학습 없이 학습 후에 업데이트할 수 있으며, 이는 동적인 지식 베이스를 위한 아키텍처 측면에서 매우 중요합니다. 시간적으로 일치하지 않는 인덱스는 성능을 눈에 띄게 저하시킵니다.
• MMLU(5-shot)에서 Atlas-11B는 약 16배 적은 파라미터에도 불구하고 GPT-3가 보고한 43.9%를 상회하는 47.9%에 도달했습니다.

유효한 점과 그렇지 않은 점

검색이 적은 파라미터 수로도 퓨샷 지식 성능을 가능하게 한다는 핵심 주장은 설득력 있게 유지됩니다. 64개의 예시로 얻은 42.4%의 NQ 수치는 놀라운 결과이며, PaLM은 당시 최첨단 규모의 벤치마크였으므로 비교도 공정합니다.

그러나 세 가지 의구심이 남습니다. 첫째, 합동 학습 후에도 검색 정확도가 그다지 높지 않습니다. 독립적인 분석에 따르면 Contriever는 약 85%의 사례에서 최소 하나의 정답 문장을 놓치며, QA 검색 정확도는 약 47%에 불과합니다. 합동 학습은 비합동 학습 베이스라인보다 검색 성능을 개선하지만, 리더가 불완전한 검색을 보완하기 위해 엄청난 일을 하고 있는 것입니다. 헤드라인을 장식한 퓨샷 수치는 시스템의 한계를 반영하는 것이지 검색 컴포넌트의 품질을 반영하는 것이 아닙니다. 둘째, 인프라 비용이 상당합니다. 사전 학습 중 문서 인덱스를 갱신하는 것은 약 30%의 추가 계산 오버헤드를 발생시키며, 전체 Wikipedia와 CommonCrawl 인덱스는 fp16 기준으로 587GB가 필요합니다. 이는 연구 환경에서는 감당할 수 있지만 실제 서비스 배포 시에는 실질적인 운영 제약이 됩니다. 셋째, 데이터 누수(data leakage)가 인정되었으나 해결되지 않았습니다. MMLU 질문의 2.8%가 사전 학습에 사용된 CCNet 코퍼스에 그대로 등장하여 MMLU 결과가 확인되지 않은 만큼 부풀려졌을 수 있습니다.

또한 논문에서 충분히 다루지 않은 미묘한 아키텍처적 한계도 있습니다. FiD는 융합 전에 검색된 각 구절을 독립적으로 인코딩하므로 병렬 처리에는 도움이 되지만, 인코더에 구절 간 어텐션(cross-passage attention)이 없음을 의미합니다. 여러 구절에 걸친 정보를 연결해야 하는 긴 멀티홉 추론 체인은 모든 작업을 디코더에서 수행해야 하며, 20개의 검색된 구절이 있을 때 디코더 크로스 어텐션은 매우 무거운 부하를 감당하게 됩니다.

금융 AI에 중요한 이유

Beancount 장부 QA의 경우, Atlas의 가장 관련성 높은 기여는 검색기-리더 합동 학습이 퓨샷 환경에서 성과를 낸다는 실증적 증명과, 그렇지 못한 경우에 대한 정직한 기록입니다. 수년간의 거래 내역을 조회하는 Beancount 에이전트는 정확히 동적 인덱스 문제에 직면합니다. 매일 새로운 항목이 추가되며, 한 달만 지나도 오래된 인덱스는 잘못된 답변을 내놓게 됩니다. Atlas는 재학습 없이 인덱스를 핫스왑(hot-swapping)할 수 있음을 보여주었으며, 이는 아키텍처적으로 고무적인 일입니다.

하지만 검색 정확도 수치는 냉정하게 바라봐야 합니다. 일반 텍스트에 대한 합동 학습 후에도 Contriever가 검색 시도 중 53%에서 관련 장부 항목을 놓친다면, 도메인 특화 상품명, 계정 계층 구조 및 beancount 지시어(directives)를 다루는 금융 도메인 에이전트는 도메인 적응형 검색기 학습이나 구조화된 쿼리 방식(정확한 계정 매칭, 날짜 필터링)에 의한 증강 검색이 필요할 것입니다. RAG 방식의 검색만으로는 합동 학습을 하더라도 고정밀 장부 작업에는 충분하지 않을 것입니다.

PaLM과의 비교는 아키텍처적 트레이드오프를 명확히 해줍니다. 검색을 통해 지식을 더 적은 파라미터로 압축하여 추론 비용을 낮출 수 있습니다. 추론 비용이 규모에 따라 중요해지는 Beancount.io와 같은 서비스에서 Atlas의 설계 철학은 매력적입니다. 그러나 587GB의 인덱스 비용은 부담을 저장 및 검색 인프라로 전가합니다. 이는 벤치마크 수치에는 나타나지 않는 또 다른 종류의 운영 제약입니다.

다음으로 읽을거리

• REALM: Retrieval-Augmented Language Model Pre-Training (Guu et al., arXiv:2002.08909, ICML 2020) — Atlas가 확장한 초기 검색 증강 언어 모델 사전 학습 프레임워크로, Atlas가 실제로 개선한 점과 그대로 둔 점을 이해하는 데 필수적입니다.
• RA-DIT: Retrieval-Augmented Dual Instruction Tuning (Lin et al., arXiv:2310.01352, ICLR 2024) — 합동 사전 학습 대신 인스트럭션 튜닝을 사용하여 Atlas와 대등한 성능을 달성합니다. 이는 인프라 비용 없이 합동 학습과 독립 학습 간의 간극을 좁힐 수 있음을 시사합니다.
• RETRO: Improving Language Models by Retrieving from Trillions of Tokens (Borgeaud et al., arXiv:2112.04426, ICML 2022) — 다른 규모에서 사전 학습 중 검색을 수행하는 DeepMind의 접근 방식입니다. 장부 QA를 위한 아키텍처를 선택하기 전에 검색 증강 사전 학습 접근법의 전체상을 완성해 줍니다.