코드 검색은 주어진 쿼리로 부터 특정 기능을 구현하는 코드를 검색하는 작업입니다. 이 작업은 인간 개발자의 생산성의 향상과 RAG를 기반으로한 LLM의 hallucination 해결에 중요한 역할을 담당하고 있습니다. 최근 Programming Language Model이 코드 검색 성능을 비약적으로 향상시켰음에도 불구하고 모델들은 여전히 다양한 유형의 shift에 취약한 모습을 보이고 있습니다.
이를 위해 우리는 Union find based Recursive Clustering Algorithm을 제안합니다. Union find based Recursive Clustering Algorithm은 코드를 구성하는 fragment들 사이의 관계를 바탕으로 fragment들을 clustering하고 각 cluster 사이의 관계를 바탕으로 code 사이의 관계를 조정함으로써 모델이 적은 수의 데이터 만으로도 코드 검색에 대하여 발생하는 다양한 shift가 존재하는 상황에서 성능을 향상시킵니다.
일반적으로 코드 검색에서 adaptation은 minimum entropy problem으로 정의됩니다. 하지만 minimum entropy problem은 주어진 sample 사이의 관계는 반영할 수 있어도 그 sample을 구성하는 fragment들 사이의 관계는 반영하지 못한다는 한계를 가지고 있습니다. 예를들어 위의 그림에서 minimum entropy problem은 함수 avg_moments()와 함수 dict_gather() 사이의 관계에 대해서는 학습을 수행할 수 있어도 Code Fragment (A)와 Code Fragment (C) 사이의 관계와 같이 fragment들 사이의 관계는 학습을 수행할 수 없습니다.
우리는 fragment 사이의 관계를 학습에 반영하기 위해 Union find based Evidence Clustering Algorithm(이하 URECA)를 제안합니다. URECA는 우선 sample들을 fragment들로 나누고 이들에 대하여 clustering을 수행합니다. 일반적인 clustering 알고리즘과 다르게 URECA는 실질적으로 clustering을 하는 게 아니라 clustering에 대한 simulation을 수행합니다. code를 구성하는 fragment 사이의 경계가 비교 대상이 되는 code에 따라 달라지기 때문에 매 iteration 마다 code를 fragment로 달리 나누는 것은 많은 계산을 필요로 합니다. 그리고 매 iteration마다 이러한 fragment들로 clustering을 수행하는 것 또한 많은 계산을 필요로 합니다. 하지만 simulation을 통해 이러한 작업들을 단순한 수치 계산으로 갈음함으로써 적은 양의 계산 비용 만으로도 이러한 작업을 수행할 수 있습니다.
해당 작업을 수행하는 과정에서 URECA는 cluster 사이의 관계를 기반으로 특정 cluster에서 다른 cluster로 옮겨져야 할 부분에 대해 추정합니다. 다만 이러한 추정 과정에서 error가 발생할 수 밖에 없습니다. 이러한 error를 최소화하기 위해서 우리는 Thresholdly Updatable Stationary Assumption을 도입합니다. 이 Assumption은 cluster 사이의 관계가 특정 기간 동안은 변함없이 유지되다가 임계값을 넘어가면 관계가 갱신된다는 가정을 의미합니다. 이를 통해 cluster 사이의 관계를 더욱 정확하게 추정함으로써 error를 최소화합니다.
최종적으로 이렇게 구성된 cluster를 바탕으로 code 사이의 관계를 재조정하도록 학습하여 fragment들 사이의 관계에 대하여 모델이 학습하게 합니다. 그리고 이를 바탕으로 모델이 적은 수의 데이터만으로도 다양한 shift가 존재하는 상황에서 잘 적응하게 만듬으로써 코드 검색 성능을 향상시킵니다.