LLM 검색 파이프라인 개선: 리랭킹 스트리밍, 폴백, 동적 top_k

2편을 시작하며

‍

1편에서는 광고 소재 검색 기능을 실제 서비스로 옮기며 늘어난 데이터와 검색 조건을 어떻게 처리할지 살펴봤습니다. 검색할 내용과 필터 조건을 구분하고, Milvus DiskANN으로 후보를 찾는 구조를 마련했습니다.

이번 편에서는 검색 중 발생하는 대기 시간과 오류에 어떻게 대응했는지 이어서 살펴보겠습니다. 리랭킹할 후보의 수와 결과를 보여주는 시점을 조정한 뒤 검색이 어떻게 달라졌는지도 함께 소개합니다.

‍

1편 보러 가기

‍

4. 빠르게 검색하고, 오류에도 대응하기

‍

소재를 한 번 검색하는 데도 여러 외부 호출이 필요합니다. 질의 이해와 리랭킹을 위해 LLM을 두 번 호출하고, 검색어 임베딩과 DB 조회, 벡터 검색도 수행합니다. 어느 한 곳의 응답이 늦어지거나 요청이 실패하면 검색 전체에 영향을 줄 수 있는 구조였습니다.

그래서 두 가지를 함께 살펴봤습니다. 정상적으로 동작할 때 불필요하게 기다리는 구간은 없는지, 일부 단계가 실패했을 때 다른 방법으로 검색을 이어갈 수 있는지입니다.

‍

4-1. 작업에 맞는 빠른 모델과 예비 모델

처음 고민한 건 어떤 모델을 사용할지였습니다.

답하기 전에 추론 과정을 거치는 모델은 할루시네이션(hallucination: 사실과 다른 내용을 그럴듯하게 만들어 내는 현상)이 적은 편입니다. 대신 이런 모델은 응답이 늦습니다. 가벼운 모델은 빠르지만 할루시네이션이 더 잦습니다. 속도와 정확도 중 하나를 포기해야 하는 것처럼 보였습니다.

그래서 모델보다 각 단계에 맡길 작업을 먼저 살펴봤습니다. 질의 이해는 사용자 문장에서 검색어와 필터 조건을 추출하는 일이고, 리랭킹은 주어진 후보의 순서를 다시 매기는 일입니다. 두 작업 모두 추론 과정을 거칠 필요가 없었습니다. 깊이 생각하는 능력보다 지시를 정확히 따르는 능력이 더 중요했습니다. 그래서 응답이 빠른 모델 가운데 할루시네이션이 적은 모델을 골랐습니다.

모델을 선택한 뒤에도 남는 문제가 있었습니다.

제공사 사정으로 응답이 늦어지거나 오류가 발생하는 상황은 직접 제어하기 어려웠습니다. 이에 대비해 서로 다른 제공사의 모델을 주 모델(Primary)과 예비 모델(Secondary)로 두었습니다. 주 모델이 실패하거나 정해진 대기 시간을 넘기면 예비 모델에 다시 요청하도록 했습니다.

여기서 신경 쓴 것은 두 모델에 시간을 나눠 배정하는 방식입니다.

한 단계에 할당된 시간을 주 모델이 모두 사용하면, 예비 모델을 준비해 두더라도 요청할 시간이 남지 않습니다. 그래서 예비 모델이 응답하는 데 필요한 최소 시간을 먼저 확보하고, 나머지 시간 동안만 주 모델의 응답을 기다리도록 했습니다.

‍

4-2. 실패했을 때 사용할 대체 경로

예비 모델을 두는 것만으로는 충분하지 않았습니다. 두 모델이 모두 실패할 수도 있고, DB 조회나 벡터 검색처럼 LLM을 사용하지 않는 단계에서도 문제가 생길 수 있기 때문입니다. 그래서 각 단계가 실패했을 때 사용할 대체 경로, 즉 폴백(fallback)을 마련했습니다.

질의 이해 단계에서 두 모델이 모두 실패하면 오류를 기록하고, 사용자가 입력한 문장을 그대로 임베딩 단계에 전달합니다. 검색할 내용과 조건을 정교하게 분리하지 못하더라도 원래 문장으로 검색을 시도하는 것입니다.

리랭킹 단계에서 두 모델이 모두 실패하면 벡터 검색에서 정한 순서를 그대로 사용합니다. 최종 순위를 다듬는 과정은 생략하지만, 이미 찾아둔 후보를 결과로 제공할 수 있습니다.

다만 LLM 단계를 건너뛰는 대체 경로는 정상 처리와 같은 품질을 보장하지 않습니다. 질의 이해를 생략하면 필터 조건을 제대로 추출하지 못할 수 있고, 리랭킹을 생략하면 결과 순서가 덜 정교해질 수 있습니다. 일부 기능을 사용하지 못하더라도 검색 전체가 중단되는 상황을 줄이기 위한 선택입니다.

Vector DB 검색이 실패하면 2장에서 설명한 브루트 포스 경로를 사용합니다. LLM 단계의 대체 경로와 달리 정확도는 떨어지지 않고, 처리 시간만 늘어납니다.

Vector DB가 응답하지 않는 상황에서 요청이 들어올 때마다 다시 기다리는 문제도 있었습니다. 이를 줄이기 위해 실패가 발생하면 일정 시간 동안 Vector DB를 건너뛰고 곧바로 브루트 포스 경로를 사용하도록 했습니다. 이 대기 기간인 쿨다운(cooldown)이 끝나면 첫 요청에서 가벼운 검색으로 상태를 확인하고, 응답이 돌아오면 원래 경로로 복귀합니다.

필터 조건을 조회하는 DB에도 대체 경로를 두었습니다. 광고가 현재 활성화되어 있는지처럼 실시간으로 바뀌는 정보는 DB에서 확인합니다. 조회에 실패하면 메모리에 보관한 데이터프레임(DataFrame)에서 확인할 수 있는 조건만 적용합니다. 이 경우 DB에만 있는 조건은 반영하지 못합니다.

현재 폴백이 없는 단계는 임베딩입니다. 저장된 소재 임베딩이 모두 같은 모델로 만들어져 있어, 사용자 질의만 다른 모델로 변환하면 두 벡터를 그대로 비교하기 어렵기 때문입니다. 따라서 모든 오류에 대응하는 구조는 아니며, 대체 가능한 단계에서 검색을 이어가도록 설계했습니다.

‍

4-3. 서로 기다릴 필요가 없는 작업은 동시에

오류에 대비한 다음에는 평소 검색에서 기다리는 시간을 줄일 방법을 살펴봤습니다. 처음에는 작업을 순서대로 실행했는데, 그중에는 서로의 결과를 기다리지 않고 시작할 수 있는 작업이 있었습니다.

검색어 임베딩과 DB 필터 조회가 그런 경우였습니다. 두 작업 모두 외부 응답을 기다려야 하지만, 임베딩이 끝나야 DB를 조회할 수 있는 것은 아닙니다. 그래서 두 요청을 동시에 보내고 결과가 모두 준비되면 벡터 검색을 시작하도록 했습니다.

브루트 포스 경로에서도 유사도 계산과 필터 마스크 생성을 동시에 진행했습니다. GPU가 전체 임베딩과의 유사도를 계산하는 동안 필터 조건을 확인하고, 두 작업이 끝나면 계산 결과에 마스크를 적용합니다.

이때는 여러 요청이 GPU 연산 버퍼를 공유한다는 점을 고려해야 했습니다. 같은 버퍼를 동시에 사용하지 않도록 잠금(lock)이 필요한데, 잠금을 거는 범위가 넓으면 다른 요청까지 불필요하게 기다리게 됩니다.

예를 들어 검색어 임베딩의 응답을 기다리는 구간까지 잠그면, GPU가 계산하고 있지 않은 동안에도 다음 요청이 연산을 시작할 수 없습니다. 그래서 잠금은 실제로 GPU가 유사도를 계산하는 구간에만 적용했습니다. 외부 응답을 기다리는 시간까지 GPU 사용을 막지 않았습니다.

각 작업의 처리 속도를 바꾸지 않더라도, 서로 독립적인 작업을 동시에 실행하면 전체 대기 시간을 줄일 수 있었습니다.

‍

5. 리랭킹 결과를 먼저 보여주기

‍

벡터 검색에서 후보를 빠르게 찾아도 사용자에게 보여주기까지는 리랭킹 과정이 남아 있었습니다. 리랭킹 단계에는 사용자의 질의와 후보 소재의 설명 청크를 입력하고, 적합한 순서로 정렬된 소재 ID를 받습니다.

리랭킹에는 LLM을 사용했는데, 반환할 ID가 많아질수록 생성할 토큰도 늘어났습니다. 그만큼 응답을 모두 받기까지 시간이 걸렸고, 리랭킹이 전체 검색 과정에서 가장 오래 걸리는 작업이 됐습니다.

그래서 결과를 전달하는 방식을 살펴봤습니다. 첫 번째 소재가 준비되어 있는데도 전체 목록이 완성될 때까지 기다려야 할까요? 먼저 생성된 결과부터 전달하면, 사용자는 나머지 결과가 준비되는 동안 첫 소재를 살펴볼 수 있습니다.

이를 위해 LLM의 스트리밍 기능을 사용했습니다. 스트리밍에서는 구조화된 출력도 처음부터 완성된 객체로 도착하지 않습니다. JSON 텍스트가 조각으로 나뉘어 들어오기 때문에, 중괄호, 대괄호, 쉼표 같은 구분자의 위치를 이용해 소재 ID를 하나씩 추출하고 순서대로 전달하도록 했습니다.

‍

‍

스트리밍을 적용해도 전체 결과가 도착하는 시간은 같았습니다. 달라진 것은 첫 ID를 받는 시점이었습니다. 첫 ID는 약 1초 만에 도착했고, 이후 ID도 0.1초 미만 간격으로 이어졌습니다.

이렇게 사용자가 첫 소재를 확인하는 시점을 약 3~4초 앞당겼습니다. 전체 처리가 끝나는 시간과 사용자가 결과를 보기 시작하는 시간을 나누어 살펴본 것이 도움이 됐습니다.

‍

6. 검색어마다 리랭킹 후보 수를 다르게 정하기

‍

리랭킹에 걸리는 시간과 비용을 줄이려면 전달 방식뿐 아니라, 처리할 후보 수도 살펴봐야 했습니다. 벡터 검색 결과에서 상위 몇 개를 선택할지 정하는 값이 top_k입니다.

후보를 많이 넘기면 관련 소재를 더 넓게 검토할 수 있지만, 관련성이 낮은 소재까지 처리하면서 비용과 시간이 늘어납니다. 반대로 너무 적게 넘기면 필요한 소재가 리랭킹에 도달하기 전에 제외될 수 있습니다. 앞 단계에서 빠진 소재를 리랭킹이 다시 찾아줄 수는 없습니다.

문제는 검색어마다 관련 소재가 모여 있는 범위가 다르다는 점이었습니다. 관련 소재가 최상위에 모여 있는 검색도 있고, 더 넓은 범위에 걸쳐 있는 검색도 있었습니다. 모든 검색에 같은 개수를 적용하기 어려웠고, 상위 일정 비율만 선택하는 방법도 검색어별 차이를 충분히 반영하지 못했습니다.

그래서 유사도 점수를 기준으로 후보 수를 정했습니다. 검색어별 최고 코사인 유사도에 일정 비율 α를 곱한 값을 임계값으로 사용했습니다.

다만 임계값만으로 후보 수를 정하기에는 데이터 규모가 컸습니다. 같은 규칙을 적용해도 임계값을 넘는 후보 수가 검색어에 따라 수천 개에서 수십만 개까지 달라졌습니다. 이를 모두 리랭킹하면 처리 시간과 비용을 관리하기 어려웠습니다.

그래서 최소 후보 수 K_min과 최대 후보 수 K_max를 함께 두었습니다. 임계값을 넘는 후보 수를 센 뒤, 최종 후보 수가 이 범위 안에 들어오도록 조정하는 방식입니다.

α를 정할 때도 균형이 필요했습니다. α를 낮추면 더 많은 후보를 포함할 수 있어 관련 소재를 놓칠 가능성을 줄일 수 있습니다. 하지만 너무 낮추면 대부분의 검색이 K_max까지 후보를 채우게 됩니다. 그러면 검색어마다 후보 수를 다르게 정하더라도 실제로는 고정된 개수를 넘기는 것과 비슷해집니다.

그래서 검색어별 유사도 분포에 따라 후보 수가 달라지도록 α를 정하고, K_max로 처리량을 제한해 비용과 응답 시간을 관리했습니다. 이렇게 필요 이상으로 많은 후보를 리랭킹하는 경우를 줄이면서, 후보가 적은 검색에서도 최소한의 후보를 검토할 수 있도록 했습니다.

‍

7. 마무리

‍

실험을 마친 검색 기능을 서비스로 옮기려면 더 많은 데이터를 처리하는 것 외에도 풀어야 할 문제가 있었습니다. 검색 조건을 반영하고, 여러 외부 서비스의 응답을 기다리고, 운영 중 발생할 오류에도 대비해야 했습니다. 원하는 소재를 찾는 데서 나아가, 사용자가 기다릴 수 있는 시간 안에 결과를 제공해야 했습니다.

이를 위해 벡터 검색에 Milvus DiskANN을 도입했습니다. 일부 단계에는 실패 시 사용할 대체 경로를 마련하고, 서로의 결과를 기다릴 필요가 없는 작업은 동시에 실행했습니다. 리랭킹할 후보 수는 검색어별로 조절하고, 결과는 스트리밍으로 전달했습니다.

‍

‍

사용자가 먼저 체감할 변화는 첫 결과를 보기까지 걸리는 시간입니다. 기존에는 13초를 기다려야 했지만, 개선 후에는 약 5초 만에 첫 소재를 확인할 수 있었습니다. 후보 검색 시간을 줄인 데 더해 리랭킹이 끝나기 전부터 결과를 전달한 결과입니다.

반환되는 소재 수도 달라졌습니다. 기존에는 모든 검색에서 10,000개를 반환했습니다. 개선 후에는 평균 1,785개, 표준편차 1,097개입니다. 표준편차가 평균의 절반을 넘는다는 것은 검색어마다 개수가 크게 달랐다는 뜻입니다. 검색어별로 후보 수를 다르게 정한다는 6장의 의도대로 동작했습니다.

남은 과제도 있습니다. 데이터가 계속 쌓이는 만큼, 현재의 응답 속도와 품질을 유지하려면 데이터 증가에 맞춰 검색과 운영 방식을 점검해야 합니다. 오류가 발생해 대체 경로를 사용할 때 일부 조건을 반영하지 못하는 한계도 있습니다. 이런 상황에서도 결과의 품질을 유지할 방법을 계속 살펴봐야 합니다.

사용자가 원하는 소재를 더 쉽게 찾을 수 있도록 검색 편의성을 높이고, 더 많은 데이터를 효율적으로 처리할 방법도 찾아갈 계획입니다.

‍

References

_{1. Subramanya, S. J., Devvrit, F., Kadekodi, R., Krishnaswamy, R., & Simhadri, H. V. (2019). DiskANN: Fast Accurate Billion-point Nearest Neighbor Search on a Single Node. NeurIPS 2019.}

_{2. Malkov, Y. A., & Yashunin, D. A. (2018). Efficient and Robust Approximate Nearest Neighbor Search Using Hierarchical Navigable Small World Graphs. IEEE TPAMI. arXiv:1603.09320.}

_{3. Milvus Documentation. MMap-enabled Data Storage.}

_{4. Milvus Documentation. DISKANN.}