RAG(검색증강생성)를 위한 코사인 유사도(Cosine Similarity) 고속 연산

요구사항

• 유저의 질문 벡터와 가장 유사한 데이터 Top 5를 추출하는 고속 벡터 연산 알고리즘 구현
• 반복문 언롤링 및 프리미티브 배열 최적화를 통해 연산 속도 극대화.

구현 사항

• OpenAi의 text-embedding-3-small 을 사용해 사용자의 질문 및 기본 지식 세팅시 Vector 배열 추출
• 기본 지식을 넣을 API 구축
  • 카테고리 설정으로 모든 정보를 다 뒤지는게 아닌 특정 카테고리 데이터만 서치
  • 외부 API를 사용해서 얻은 지식 벡터 데이터는 postgresql 에 저장
• 사용자의 질문 API 구축
• 사용자의 질문에 대해 유사도 Top5 추출 기능
  • pgvector 기반 DB내 Top5 추출
  • 반복문 언롤링 및 프리미티브 배열 최적화를 통한 고속 벡터 연산 알고리즘 사용해 Top5 추출

 

InMemoryVectorRow

• text-embedding-3-small 를 사용하게 되면 1536 차원의 벡터 데이터를 받게 되는데 이것을 float[] 형태로 저장하고 적절한 분야에서 데이터를 찾기 위해 카테고리 필드를 추가.

InMemoryCacheRegistry

• @PostConstruct 를 사용해 미리 세팅 해 둔 지식 데이터를 메모리에 올리기 위한 준비 클래스

RagSearchService

• 사용자의 질문 Payload를 접수하여, 카테고리로 적절한 분야 기반으로 4칸 루프 언롤링으로 코사인 유사도를 측정
• DB, Memory 구분 

RagSearchController

• 프론트엔드나 포스트맨의 고밀도 JSON 요청을 바디로 깔끔하게 접수하여 결과 보고서를 출고하는 엔드포인트.

---

 

반복문 언롤링

 

우리가 흔히 짜는 for (int i = 0; i < 1536; i++) 이 반복문은 CPU가 내부적으로 명령을 수행 할 때 큰 비효율이 존재한다.

이 루프는 다음과 같은 명령을 수행한다.

1. i < 1536 조건이 참인지 검사
2. 루프 내부 연산 수행
3. i++ 실행 및 다시 루프 시작점으로 이동

컴퓨터 구조 관점에서 CPU는 파이프라인을 통해 여러 명령어를 겹쳐 처리하려고 합니다. 하지만 조건문이나 반복문의 분기 지점에서는 다음에 실행할 명령어의 흐름이 아직 확정되지 않았기 때문에 CPU는 분기 예측을 사용합니다. 만약 예측이 맞으면 파이프라인은 계속 효율적으로 동작하지만, 예측이 틀리면 이미 파이프라인에 들어온 잘못된 명령어들을 버리고 올바른 명령어 흐름으로 다시 채워야 합니다. 이 과정에서 파이프라인 스톨이나 플러시가 발생할 수 있고, 잠깐이지만 ALU와 같은 실행 장치가 제대로 활용되지 못하는 오버헤드가 생깁니다.

 

이런 문제를 해결하기 효율적으로 해결하기 위해 반복문 언롤링 이라는 것을 사용할 수 있습니다. 

for (; i <= len - 4; i += 4) {
    dotProduct += vecA[i] * vecB[i]; normA += vecA[i] * vecA[i]; normB += vecB[i] * vecB[i];
    dotProduct += vecA[i+1] * vecB[i+1]; normA += vecA[i+1] * vecA[i+1]; normB += vecB[i+1] * vecB[i+1];
    dotProduct += vecA[i+2] * vecB[i+2]; normA += vecA[i+2] * vecA[i+2]; normB += vecB[i+2] * vecB[i+2];
    dotProduct += vecA[i+3] * vecB[i+3]; normA += vecA[i+3] * vecA[i+3]; normB += vecB[i+3] * vecB[i+3];
}

 

위 코드는 4-Way 반복문 언롤링으로 한 번의 반복에서 4개의 연산을 처리하도록 만듭니다. 예를 들어 1,536번 반복하던 루프를 384번의 반복으로 줄이는 방식으로. 이 방식의 장점은 반복 횟수가 줄어들면서 루프 증가 연산, 조건 비교, 분기 명령 같은 제어 명령의 실행 횟수도 함께 줄어든다는 점입니다. 따라서 CPU 입장에서는 루프 분기 명령을 처리하는 횟수가 줄어들고, 잘못된 분기 예측이 발생할 수 있는 기회도 줄어듭니다. 또한 한 번의 반복 안에 여러 연산이 배치되기 때문에 CPU가 명령어를 더 연속적으로 처리하기 쉬워져 파이프라인 효율과 명령어 수준 병렬성이 좋아질 수 있습니다.

 

---

 

프리미티브 배열 최적화

 

프리미티브 배열 최적화 부분은 위에서 반복문 언롤링과 이어지게 될 수 있습니다 Java에는

• Primitive type(원시타입)
• Reference type(참조타입)

두가지 타입이 있습니다 간단히 말해 참조타입은 new 으로 생성되기 때문에 메모리 공간에 렌덤적으로 생성되게 됩니다 이런 상황으로 반복문 최적화를 하게 되면 루프를 도는 횟수는 줄어들 수 있지만 내부적으로 엄청난 비효율이 발생하게 됩니다.

 

0번 방 값을 연산하고 나면, 1번 방 값을 읽기 위해 다시 주소값을 타고 메모리의 완전히 다른 구역으로  이동 했다 오는 포인터 추적
오버헤드 가 발생합니다 데이터의 위치가 불 연속적이기 때문에 CPU 캐시를 효율적으로 사용하지 못하고 메모리 까지 내려갔다오는
비 효율이 발생하게 됩니다.

 

이때 프리미티브 타입의 배열을 사용하게 되면 CPU가 메모리에서 데이터를 읽어올 때는 딱 숫자 1개만 가져오지 않고, 효율을 위해 그 주변 데이터까지 세트로 묶어서 캐시 라인(Cache Line, 보통 64바이트 단위) 단위로 가져오게 됩니다. 그렇기 때문에 컴퓨터는 프리미티브 배열 vecA[0]의 숫자를 읽는 순간, CPU 내부의 하드웨어 장치는 내부적으로 vecA[1] vecA[2] ... 데이터도 순차적으로 읽는다는것을 감지 해 CPU 캐시에 올려두게 됩니다. 

 

이렇게 되면 메모리까지 내려갔다 오는 횟수도 최적화가 되고 메인 메모리를 다녀오느라 발생하는 지연 시간(Latency)를 최적화 할 수 있습니다, 또한 참조 타입과 원시타입의 메모리 차지하는 크기도 원시타입이 더 적기 때문에 메모리 공강 효율의 이점을 얻을 수 있습니다.

 

Spring AI 의 Text embedding 모델을 사용해 보았는데 여기서도

float[] embedding = embeddingModel.embed(text);

 

프리미티브 타입의 배열을 반환 하는것을 알 수 있습니다.