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

왜 중요한가

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

핵심 기여

Retrieval Augmented Search (RAS) 기법 도입

RAS는 LLM이 생성한 자연어 설명을 임베딩 기준으로 유사 예제를 검색하는 contextual retrieval과 빔 탐색 기반의 반복적인 retrieve-optimize-evaluate 루프를 결합한 방법이다. 이 방식은 소스 코드의 표면적 차이보다 알고리즘·자료구조 수준의 유사성을 기준으로 예제를 찾아 LLM에게 더 관련성 높은 in-context 예시를 제공한다. 실험에서 기존의 dynamic retrieval보다 평균 최고 속도 향상 수치가 크게 높게 나타났다.

Atomic Edit Guided Search (Aegis)으로 해석성 향상

Aegis는 학습 데이터의 느린-빠른 코드 쌍을 LLM으로 분해해 원자적 편집 집합을 생성하고 이를 일반화하여 Π_atomic을 구성한다. RAS에서 Π_atomic을 사용하면 LLM이 한 번에 큰 구조적 변경을 하기보다 점진적이고 작은 편집을 적용하게 되어 편집 크기(edit distance)가 감소한다. 실험에서 Aegis는 평균 편집 거리와 편집 크기 측면에서 개선을 보이며, 일부 설정에서 실효성도 유지했다.

맥락 임베딩 기반 검색이 코드 임베딩 기반 검색을 능가함을 입증

문맥 기반 검색은 p의 자연어 설명 s = F_context(p)를 먼저 생성하고 이를 ψ로 임베딩하여 ϕ(p)=ψ(s)를 사용한다. 이 절차는 코드 자체를 바로 임베딩하는 code retrieval보다 유사 예제의 품질을 높였고, ablative 실험에서 이 요소의 제거는 성능을 절반가량 감소시켰다. 따라서 코드 표면적 유사성보다 알고리즘 수준의 유사성을 포착하는 접근이 중요함이 확인되었다.

다양한 모델과 벤치마크에서의 실험적 검증

PIE(C++)와 Mercury(Python) 벤치마크에서 RAS와 Aegis의 성능을 광범위하게 평가했다. RAS는 모델별로 최대 9.18×(DeepSeek 3.2) 평균 최고 속도 향상을 보였고, Aegis는 GPT-4o 환경에서 6.08×를 기록하면서 편집 크기를 줄였다. 추가로 정확도 검증으로 CBMC를 일부 적용하여 자동 생성 코드의 의미론적 동등성도 높은 비율로 유지됨을 확인했다.

핵심 아이디어 이해하기

대형 언어 모델은 프로그램의 실행 성능 정보를 훈련 데이터에서 직접 얻기 어렵기 때문에, 단순히 원본 코드를 입력으로 주고 최적화된 코드를 바로 생성하도록 하는 방식은 한계가 있다. 기존 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 호출이 필요해 학습 단계의 계산 비용이 더 커진다. 또한 논문에서 명시한 바와 같이 객체지향적이고 모듈화된 대형 코드베이스로 확장할 때는 최적화 대상 코드 조각을 분리·선정하는 추가 단계가 필요해 확장성이 제한적일 수 있다.

실무 활용

이 방법들은 블랙박스 환경에서 LLM을 활용해 기존 코드의 실행 성능을 개선할 때 적용 가능하다. 특히 경쟁 프로그래밍이나 알고리즘 중심의 코드베이스에서 손쉽게 고수준 최적화를 자동화할 수 있으며, Π_train과 gem5와 같은 결정론적 측정 파이프라인이 갖춰진 환경에서 재현성이 확보된다. 다만 빔 탐색과 원자적 편집 생성 과정은 추가적인 LLM 호출과 전처리 비용을 수반한다.

경쟁 프로그래밍 출처의 반복 제출 데이터로부터 학습해 코드 제출의 실행 성능을 자동으로 향상시키는 파이프라인 구축
레거시 알고리즘 코드의 알고리즘·자료구조 수준 개선을 자동화해 개발자 검토 전후의 성능 개선 후보를 생성하는 도구
소규모 모델 성능 갭을 줄이고자 하는 환경에서 contextual retrieval을 이용해 저비용 모델의 실행 성능을 개선하는 보조 워크플로
자동화된 최적화 제안과 함께 편집 크기를 줄여 변경 내역의 해석 가능성을 높이는 코드 리뷰 보조 도구

코드 공개 여부: 공개

코드 저장소 보기

키워드

LLM(대형 언어 모델)Retrieval Augmented Search(RAS)Atomic Edits(Aegis)contextual retrieval(문맥 기반 검색)FAISS(벡터 검색 라이브러리)PIE benchmark(PIE 벤치마크)

용어 해설

문맥 기반 검색(Contextual Retrieval): — 프로그램 소스 코드 자체 대신 LLM이 생성한 자연어 설명을 임베딩하여 유사도를 계산하는 검색 방식으로, 구현 세부사항보다 알고리즘·자료구조 같은 추상적 특성을 기준으로 유사 예제를 찾고자 할 때 사용된다.
원자적 편집(Atomic Edit): — 한 번에 하나의 구체적 최적화만 수행하는 점진적 코드 변경 단위로서, 원본-최적화 쌍을 작은 단계로 분해해 LLM이 더 작고 해석 가능한 수정만 적용하도록 유도하는 데이터 표현이다.
FAISS: — 대규모 벡터 집합에서 유사도 검색을 고속으로 수행하는 라이브러리로, 본 논문에서는 임베딩된 자연어 설명이나 코드 임베딩으로부터 top-k 유사 예제를 검색할 때 사용되었다.
gem5: — 하드웨어 시뮬레이터로, C++ 프로그램의 실행 시간을 결정론적으로 측정하여 최적화의 실행 성능을 재현 가능하게 평가하는 데 사용되었다.
빔 탐색(Beam Search): — 반복적 생성 과정에서 여러 후보를 유지하면서 각 단계에서 상위 성능 후보를 선택해 나가는 탐색 전략으로, RAS는 검색-최적화-평가 루프를 m 단계의 빔 탐색 형태로 수행했다.

코드 예제

text

Algorithm 1 Retrieval Augmented Search (RAS)
p0, Π_train, F_opt, F_context, R, ϕ
for i ∈ [1, …, m] do
  Π_i ← top-k{((p,p'), d_ϕ(p_{i-1}, p)) ∣ (p,p') ∈ Π_train}
  p_i^j ∼ F_opt(π_i^j, p_{i-1})    (∀ j ∈ [k])
  p_i ← argmax_{j∈[k]} R(p_i^j)
return p_m

RAS의 핵심 알고리즘을 그대로 옮긴 의사코드로, 각 반복에서 top-k 유사 학습 쌍을 검색하여 F_opt로 후보 최적화를 생성하고 가장 빠른 올바른 프로그램을 선택한다.

text

Algorithm 2 Atomic Edit-Guided Search (Aegis)
Π_train, F_decomp, F_edit, F_gen, F_opt, F_context, R
Π_atomic ← ∅
for (p, p') ∈ Π_train do
  [s_1, …, s_r] ∼ F_decomp(p, p')
  for i ∈ [1, …, n] do
    p_i ∼ F_edit(s_i, p_{i-1})
    e_i ∼ F_gen(s_i, p_i)
    Π_atomic ← Π_atomic ∪ {(e_i, (p_{i-1}, p_i))}
return Π_atomic

Aegis의 원자적 편집 생성 파이프라인을 요약한 의사코드로, 원본-최적화 쌍을 분해해 단계별 편집과 일반화를 통해 Π_atomic을 만든 뒤 RAS에 투입한다.

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

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

용어 해설

코드 예제

관련 토론

댓글

관련 피드

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

관련 토론

댓글

관련 피드