검색 증강 탐색을 통한 LLM 기반 프로그램 최적화

프로그램 최적화는 기존 컴파일러가 놓치는 알고리즘·자료구조 수준의 변환을 필요로 하며, LLM은 이러한 고수준 최적화를 자동화할 잠재력을 보인다. LLM은 학습 데이터에 성능 정보가 부족해 그대로 사용하면 성능 향상이 제한되며, 본 논문은 블랙박스 방식으로 LLM을 적응시켜 이 문제를 해결했다. 실제 벤치마크에서 실행 시간 개선을 보이며 LLM을 성능 최적화에 실용적으로 적용할 수 있는 길을 제시했다.

대형 언어 모델은 프로그램의 실행 성능 정보를 훈련 데이터에서 직접 얻기 어렵기 때문에, 단순히 원본 코드를 입력으로 주고 최적화된 코드를 바로 생성하도록 하는 방식은 한계가 있다. 기존 dynamic retrieval 방식은 코드 임베딩을 기반으로 관련 예제를 찾지만, 같은 알고리즘을 서로 다른 방식으로 구현한 코드들 사이의 표면적 차이를 제대로 포착하지 못해 관련성 낮은 예제가 선택될 수 있다. 이 결과 LLM의 in-context 학습이 성능 신호를 받아들이는 효율이 떨어진다.

본 논문의 핵심은 두 단계로 맥락을 추상화하고 검색-생성 과정을 반복하는 것이다. 첫째, 주어진 프로그램으로부터 LLM(F_context)을 통해 자연어 설명을 생성하고 그 설명을 임베딩해 검색 공간을 추상화한다. 둘째, 검색된 예제를 바탕으로 F_opt를 호출해 후보 최적화를 생성하고, 이들 후보의 실행 성능 R(p)를 측정해 빔 탐색 형태로 여러 단계에 걸쳐 최적화 경로를 탐색한다. 이 과정은 코드의 구현 세부를 넘어 알고리즘·자료구조 수준의 유사성을 기준으로 예제를 활용하므로 더 관련성 높은 in-context 학습이 가능하다.

Aegis는 학습 데이터 쌍을 더 작은 단위로 분해해 LLM이 적용할 편집의 크기를 줄인다. 구체적으로 F_decomp로 원본-최적화 쌍을 일련의 자연어 편집 s_i로 분해하고, F_edit로 각 편집을 적용해 중간 프로그램을 생성한 뒤 F_gen으로 편집을 일반화해 원자적 편집 e_i를 만든다. 이렇게 구성된 Π_atomic을 사용하면 RAS가 한 번에 큰 변경을 하기보다 여러 단계의 소규모 편집을 조합해 최종 최적화를 도출하게 되고, 이는 편집 거리 감소와 해석성 향상으로 이어진다.

이 접근은 검색 품질과 탐색 전략을 동시에 개선해 단일 호출로 최적화를 시도하는 Instruct Only 방식과, 코드 임베딩 기반의 동적 검색보다 실험적으로 더 우수한 성능을 달성했다. RAS는 모델과 데이터셋에 따라 평균 최고 속도 향상에서 동적 검색 대비 최대 2.06배의 향상을 보였고, Aegis는 편집 크기를 줄이면서도 상당한 속도 향상을 유지했다.

전체 프레임워크는 학습용 느린-빠른 코드 쌍 Π_train과 최적화용 LLM F_opt, 설명을 생성하는 F_context, 임베딩 모델 ψ, 벡터 검색기(FAISS)를 입력으로 받는다. 초기 프로그램 p_0에서 시작해 각 반복에서 현재 프로그램 p_{i-1}의 자연어 설명 s=F_context(p_{i-1})을 생성하고 ψ(s)를 임베딩으로 얻어 FAISS로 top-k 유사 학습 쌍 Π_i를 검색한다. 검색된 각 쌍 π_i^j를 in-context 예제로 F_opt에 주어 p_i^j를 샘플링하고, 모든 샘플 중 테스트 케이스를 통과하면서 R(p_i^j)가 최대인 후보를 p_i로 선택하는 retrieve-optimize-evaluate 루프를 m번 반복한다.

임베딩 파이프라인은 코드 대신 자연어 설명을 임베딩하는 ϕ(p)=ψ(F_context(p))를 사용해 검색의 추상화를 달성한다. 이는 구현 세부 표현의 차이를 무시하고 알고리즘 및 자료구조 수준의 특성을 기반으로 유사도를 계산하기 위한 설계이다. 검색은 FAISS로 수행해 대규모 Π_train에서 top-k를 고속 검색하며, 하이퍼파라미터로 k(검색 수), m(탐색 단계 수), h(샘플 수)가 사용된다.

Aegis 파이프라인은 학습 데이터 전처리 단계로 작동한다. 각 학습 쌍 (p,p')에 대해 F_decomp가 자연어 편집 리스트 [s_1,...,s_r]를 생성하고 F_edit가 순차적으로 이 편집들을 p에 적용해 중간 프로그램 시퀀스를 만든다. 마지막으로 F_gen이 각 편집을 더 일반화된 atomic edit e_i로 바꾸어 Π_atomic을 구성하며, RAS는 이 Π_atomic을 이용해 더 작은 단계의 편집들을 검색하고 적용하도록 LLM 입력을 조정한다.

평가에서는 PIE(C++)와 Mercury(Python) 벤치마크를 사용해 gem5 시뮬레이터로 실행 시간을 측정하고 평균 최고 속도 향상(mean best speedup)과 % Optimized, Beyond@1 같은 지표를 보고했다. 하이퍼파라미터 기본값으로 RAS에서는 k=8, m=4, h=1을 사용했고, baselines와 계산량을 공정하게 비교하도록 LLM 호출 수에 따라 샘플링 전략을 조정했다. 검증을 위해 CBMC를 일부 사례에 적용해 의미론적 동등성 여부를 점검했고, 실행 시간 표본으로 고정된 5개의 테스트 케이스 평균을 사용해 gem5의 결정론적 특성과 상관관계를 확인했다.

PIE 벤치마크(C++ 최적화)에서 RAS는 모델별로 유의한 성능 향상을 보였다. 표에서 GPT-4o 환경에서는 RAS가 평균 최고 속도 향상 8.03을 기록했으며 Qwen-3-Coder에서는 8.70, DeepSeek 3.2에서는 9.18을 기록해 동적 검색(dynamic retrieval)의 4.43, 4.23, 7.03 대비 큰 개선을 보였다. 이러한 결과는 contextual retrieval과 반복적 탐색이 결합되었을 때 블랙박스 적응 성능이 크게 향상됨을 나타낸다.

Aegis는 해석성을 개선하면서도 실무적으로 유의미한 속도 향상을 유지했다. GPT-4o 환경에서 Aegis는 평균 최고 속도 향상 6.08을 보였고, Aegis는 RAS 대비 편집 거리(문자 수준)를 평균 17% 줄였으며 첫 번째 편집에 한정하면 30% 감소를 관찰했다. 이 결과는 원자적 편집을 통한 점진적 변환이 편집 크기를 줄이면서도 상당한 성능 향상을 제공함을 시사한다.

Python 최적화(Mercury) 실험에서는 RAS가 Qwen2.5-7B-Instruct의 mean runtime percentile(Beyond@1)을 10.27 포인트 개선해 작은 모델의 성능 격차를 줄였다. 정확도 측면에서 GPT-4o 실험에서 RAS는 검색 과정 중 잘못된 프로그램을 선택하는 사례가 5/973으로 나타나 정확도는 약 99.5%였고, Aegis는 0/973으로 기록되어 100%로 나타났다. 형식 검증을 위해 선택한 DeepSeek 3.2 결과 10건에 대해 CBMC를 적용한 결과 8건에서 완전 동등성이 확인되어 자동 생성 최적화의 의미론적 안전성도 높은 편이었다.

전체 시스템은 세 가지 LLM 역할(F_context, F_opt, F_decomp/F_edit/F_gen)과 임베딩 모델 ψ, 그리고 FAISS 기반 벡터 검색을 통합한 파이프라인으로 구성된다. F_context는 입력 프로그램을 받아 알고리즘·자료구조 같은 핵심 특성을 기술하는 자연어 설명 s를 생성하고, ψ(s)를 임베딩해서 검색 지표 ϕ(p)=ψ(F_context(p))를 만든다. 검색은 L2 거리 기준으로 FAISS에서 top-k를 선택하고 이들 예제를 in-context로 F_opt에 전달해 후보 최적화를 샘플링한 뒤 R(p)으로 평가해 다음 단계의 시작점을 선택한다.

수학적 관점에서 검색 거리는 d_ϕ(p,q)=||ϕ(p)-ϕ(q)||2로 정의되며, RAS는 각 반복에서 Π_i=top-k{((p,p'), d_ϕ(p{i-1},p))}을 구성한다. F_opt는 주어진 in-context examples π과 대상 프로그램 p에 대해 여러 후보 p'를 샘플링하며, 그 중 테스트 케이스를 모두 통과하는 후보의 실행 시간으로 R(p')를 계산한다. 빔 탐색은 m 단계로 구성되고 각 단계에서 k개의 검색과 h개의 샘플링을 통해 후보를 확장하므로 전체 LLM 호출 수는 m·k·h(설정에 따라 다름)에 의해 결정된다.

Aegis는 학습 데이터의 쌍을 분해하는 단계적 알고리즘을 포함한다. F_decomp는 (p,p') 쌍을 여러 자연어 편집 s_i로 분해하고, F_edit는 각 s_i를 순차적으로 적용해 중간 프로그램 p_i를 생성한다. F_gen은 각 s_i와 p_i를 입력으로 받아 편집을 더 일반화한 e_i를 생성하며, 최종 Π_atomic은 {(e_i,(p_{i-1},p_i))}의 집합으로 구성되어 RAS의 검색 후보로 사용된다.

실험적 설정은 PIE의 gem5 기반 실행 시간 측정과 Mercury의 Beyond@1 지표를 사용했다. PIE 실험에서는 k=8, m=4, h=1을 기본값으로 사용했고 dynamic retrieval baseline에는 k=4, h=32 등의 조정을 통해 LLM 호출 예산을 정규화했다. 실행 시간 평가에서는 각 프로그램에 대해 미리 고정한 5개의 테스트 케이스 평균을 사용했으며, 전체 테스트셋 실행 시간과의 상관계수(Pearson r=0.89, Spearman ρ=0.86)를 확인해 샘플링의 신뢰성을 확보했다.

두 방법 모두 빔 탐색과 반복적 LLM 호출을 포함하므로 실행 비용이 크고 대규모 배포에서 비용·지연 문제가 발생할 수 있다. Aegis는 Π_atomic을 만들기 위한 추가 전처리와 LLM 호출이 필요해 학습 단계의 계산 비용이 더 커진다. 또한 논문에서 명시한 바와 같이 객체지향적이고 모듈화된 대형 코드베이스로 확장할 때는 최적화 대상 코드 조각을 분리·선정하는 추가 단계가 필요해 확장성이 제한적일 수 있다.