MSL 멤버 RAG 논문 발표..."컨텍스트 창 16배 확장·효율 강화"

메타의 슈퍼인텔리전스 랩(MSL) 소속 연구원이 참여한 논문이 발표됐다. 기존 검색 증강 생성(RAG)의 효율성을 획기적으로 높인 기술을 공개했다.

MSL은 1일(현지시간) 싱가포르 국립대학교, 라이스대학교 연구진 등과 RAG 효율성을 개선한 디코딩 프레임워크 'REFRAG(REpresentation For RAG)'를 온라인 아카이브에 발표했다.

REFRAG는 대현언어모델(LLM) 컨텍스트 창을 16배 확장하고, 정확도 저하 없이 첫번째 토큰까지의 시간(TTFT)을 최대 30.85배 가속하는 기술이다.

일반적으로 LLM의 어텐션 메커니즘은 입력 길이에 따라 계산량과 시간이 증가한다. 문서 길이가 두배면 계산과 메모리 비용이 4배로 증가할 수 있다. 이는 추론 속도를 떨어뜨릴 뿐만 아니라 키-값(KV) 캐시 크기도 증가시켜 대규모 애플리케이션을 실행하기 어렵게 만든다.

특히 RAG 환경에서는 검색된 대부분의 문서가 최종 답변에 거의 영향을 주지 않지만, 모델은 여전히 ​​이 문서를 처리하는 데 상당한 리소스를 소모한다.

REFRAG는 검색된 내용을 16개의 토큰으로 분할하고, 각 청크를 고밀도 임베딩으로 압축하는 경량 인코더를 도입했다. 따라서 수천 개의 원시 토큰을 입력하는 대신, 디코더는 짧게 처리된 임베딩 시퀀스를 처리한다. 이로 인해 모델 아키텍처를 변경하지 않고도 처리하는 시퀀스 길이가 16배 단축된다.

디코더의 입력 시퀀스 단축은 어텐션 계산과 KV 캐시를 줄이는 결과를 낳는다. 연구진은 경험적으로 k값이 16일 때 TTFT는 16.53배, 32일 때 30.85배 빨라진다고 소개했다. 이는 컨텍스트 창 확장 최신 방식인 CEPE(Context Expansion with Parallel Encoding)를 크게 능가한다는 것이다. 

동시에 강화 학습(RL)을 통해 모델의 정확도를 그대로 유지한다고 설명했다. 정보 밀도가 가장 높은 청크를 식별, 압축에서 제외하고 원시 토큰을 디코더에 직접 전달하는 방식이다. 이를 통해 중요한 정보가 손실되지 않도록 조치한다,

여러 벤치마크에서 REFRAG는 훨씬 낮은 지연 시간으로 작동하며 CEPE에 비해 정보의 복잡도를 유지하거나 향상했다.

소설과 논문 등으로 구성된 '슬림파자마(SlimPajama)' 데이터셋에서 200억개의 토큰을 사전 학습하고 여러 장문 컨텍스트 데이터셋에서 이를 테스트했다. 그 결과 RAG 벤치마크, 멀티턴 대화 작업, 장문 문서 요약에서 뛰어난 성능을 보였다.

'라마(LLaMA-2)'를 대상으로 한 테스트에서는 16배의 컨텍스트 확장은 물론, 4개의 데이터셋에서 CEPE보다 복잡도가 9.3% 향상됐다. 또 동일한 지연 조건에서 더 많은 문서를 처리할 수 있기 때문에 정확도도 향상됐다.

이처럼 REFRAG는 긴 컨텍스트 처리에서 LLM의 효율을 크게 끌어 올렸다고 소개했다. 이를 통해 "전체 보고서 분석, 여러 차례 대화 처리, 엔터프라이즈 RAG 시스템 확장과 같은 대규모 애플리케이션이 정확도 저하 없이 실행 가능할 뿐만 아니라 효율적"이라고 강조했다.

REFRAG의 코드는 깃허브를 통해 공개될 예정이다. 

한편, 이번 연구에 참여한 MSL 연구원은 아마존닷컴 수석 응용 과학자 출신으로 2023년부터 메타에 합류한 비제이 모한이다. 그는 이번 조직 개편으로 MSL에 합류한 기존 '라마' 출신 연구진 중 한명이다.

지난 6월 MSL이 공식 출범한 이후, 멤버가 참여한 논문이 주목받은 것은 이번이 처음이다. 

임대준 기자 ydj@aitimes.com