기술 요약: 프로그래밍 지식 그래프를 이용한 컨텍스트 증강 코드 생성

문제 정의

대규모 언어 모델(LLM)은 자연어-코드 생성 분야에서 뛰어난 능력을 보여주었으나, 특정 API 관례, 코너 케이스(corner cases), 관용적 패턴과 같은 외부 지식을 필요로 하는 복잡한 프로그래밍 작업에서는 빈번하게 어려움을 겪는다. 검색 증강 생성(RAG)은 외부 소스(예: 문서, 저장소)에 생성의 근거를 두는 해결책으로 부상했다. 그러나 기존의 코드용 RAG 방식은 두 가지 주요 병목 현상을 겪는다:

1. 검색 정밀도(Retrieval Precision): 표준 검색 모델은 종종 중복되거나, 부분적으로만 관련이 있거나, 오해의 소지가 있는 컨텍스트를 반환한다. 코드와 문서를 평면적인 최근접 이웃(nearest-neighbor) 청크로 취급하는 것은 이질적인 프로그래밍 산물(코드 중심 vs. 텍스트 중심)의 특성을 반영하지 못한다.
2. 환각 및 노이즈(Hallucination and Noise): 생성 모델은 무관한 검색 데이터에 의해 방해받거나, 특히 제한된 컨텍스트 창 예산 하에서 노이즈가 섞인 컨텍스트를 통합하도록 강제될 때 환각을 일으킬 수 있다.

현재의 방법들은 검색 단위의 세밀함(granularity)과 컨텍스트 조립에 대한 명시적인 제어가 부족하여, 복잡한 시나리오에서 성능 향상이 아닌 성능 저하를 초래한다.

방법론

저자들은 이질적인 지식을 세밀하고 의미론적으로 일관된 노드로 분리하고 조직화하는 RAG를 위한 구조적 표현 방식인 **프로그래밍 지식 그래프(Programming Knowledge Graphs, PKG)**를 제안한다. 방법론은 세 단계로 구성된다:

1. PKG 구축

시스템은 서로 다른 데이터 모달리티를 처리하기 위해 두 가지 별도의 그래프 변형을 구축한다:

• 코드 중심 PKG (Code-Centric PKG): 소스 코드를 추상 구문 트리(AST)로 파싱하여 유향 계층 구조를 생성한다. 함수 노드는 블록 노드(예: if, for, try)의 시퀀스로 확장되며, 이 블록들은 다시 하위 블록으로 확장된다. 이는 함수의 루트로부터 파생되는 다층 트리 구조를 생성하여 구문적 포함 경계를 보존한다.
• 텍스트 중심 PKG (Text-Centric PKG): 튜토리얼 및 문서 텍스트를 구조화된 JSON으로 변환한다. 노드는 JSON 필드/섹션을 나타내고, 엣지는 구조적 관계(예: 설명을 예제에 연결)를 인코딩하는 유향 비순환 그래프(DAG)가 구축된다. 이를 통해 전체 구절을 검색하는 대신 필드 수준의 검색이 가능해진다.

2. 검색 및 컨텍스트 제어

사용자 쿼리가 주어지면 시스템은 코사인 유사도를 기반으로 관련 노드를 검색한다. 컨텍스트 창을 관리하고 노이즈를 줄이기 위해 저자들은 **트리 프루닝(tree pruning, 트리 가지치기)**을 구현한다:

• 루트 노드(예: 함수 또는 JSON 섹션)를 검색하면, 시스템은 자식 브랜치를 평가한다.
• 무관한 브랜치는 남은 서브그래프를 쿼리 임베딩과 대조하여 점수를 매김으로써 제거(pruning)된다.
• 이를 통해 가장 적절하고 세밀한 콘텐츠만이 프롬프트에 직렬화되어 추가되도록 보장한다.

시스템은 두 가지 세밀도 모드로 작동한다:

• Func-PKG: 함수 수준에서 검색한다 (더 거친 단위).
• Block-PKG: 블록 수준에서 검색한다 (더 세밀한 단위), 높은 정밀도를 목표로 한다.

3. 솔루션 재순위화 (Solution Reranking)

검색이 때때로 오류를 도입할 수 있다는 점을 인식하여, 저자들은 사후 생성 재순위화(post-generation reranking) 메커니즘을 채택한다.

• 다양한 전략(NoRAG, BM25, PKG 변형)을 사용하여 여러 후보 솔루션을 생성한다.
• 후보들은 구문적 타당성과 런타임 건전성(샌드박스 실행)을 기준으로 필터링된다.
• 유사도 기반 재순위화 모델이 다양한 후보 세트 중에서 최적의 후보를 선택하며, 쿼리와 일치하는 솔루션을 우선시하는 동시에 잘못된 검색 컨텍스트의 영향을 완화한다.

주요 기여

1. 새로운 표현 방식 (PKG): 코드의 AST 유래 계층 구조와 텍스트의 JSON 기반 DAG를 모델링하여 세밀한 검색을 가능하게 하는 RAG를 위한 구조적 지식 표현.
2. 세밀도 분석 (Granularity Analysis): 함수 수준(Func-PKG) 대 블록 수준(Block-PKG) 검색을 비교하여 정밀도-재현율(precision-recall) 트레이드오프를 정량화함으로써 검색 단위에 대한 명시적 평가 수행.
3. 구조적 컨텍스트 조립: 튜토리얼을 구조화된 JSON으로 정규화하여 필드 수준의 검색과 구조적 컨텍스트 조립을 가능하게 하는 텍스트 중심 검색 메커니즘.
4. 프루닝 메커니즘: 프롬프트 구성 전 컨텍스트 성장을 제어하고 효용이 낮은 브랜치를 제외하기 위해 검색된 그래프 구조에 적용되는 트리 프루닝 전략.
5. 재순위화 통합: 최종 정확도를 높이고 검색 유도 환각을 완화하기 위해 후보 재순위화 단계를 직교적인 선택 메커니즘으로 적용.

실험 결과

본 접근 방식은 다양한 Code-LLM(CodeLlama, StarCoder2, DeepSeek-Coder) 및 일반 LLM(Llama3.1)을 사용하여 HumanEval 및 MBPP 벤치마크에서 평가되었다.

• 성능 향상:
  • NoRAG 베이스라인과 비교했을 때, PKG 방식은 두 벤치마크 모두에서 pass@1 정확도를 최대 20% 향상시켰다.
  • 희소(BM25) 및 밀집(Voyage-Code-2) 검색 베이스라인과 비교했을 때, PKG는 HumanEval에서 최대 8%, MBPP에서 34%의 개선을 보였다.
  • Block-PKG는 재순위화되지 않은 방법들 중 일반적으로 가장 좋은 평균 성능을 보여, 더 세밀한 검색 단위가 무관한 컨텍스트를 줄인다는 것을 시사했다.
• 재순위화의 영향:
  • 재순위화는 상당한 추가 이득을 제공하여, 재순위화되지 않은 최적의 방법들보다 HumanEval에서 약 4포인트, MBPP에서 12포인트의 pass@1 향상을 가져왔다.
  • 이는 검색이 다양성을 확보하기 위한 메커니즘으로 작용할 때 가장 효과적이며, 선택 과정이 최종 정확도의 주요 동력임을 나타낸다.
• 오류 분석:
  • 복잡한 MBPP 데이터셋에서 PKG는 AssertionError(의미적 불일치)를 크게 줄였으나, 일부 모델에서 더 많은 NameError 및 IndentationError를 유발했다.
  • 텍스트 중심 PKG(JSON-PKG)는 이질적인 효과를 보였다: 일반 LLM과 특정 모델의 HumanEval 성능에는 도움이 되었으나, 코드 생성 작업에 대해서는 코드 중심 PKG보다 일관성이 낮았다.
• 비용:
  • PKG 구축에는 추가적인 전처리 시간(임베딩 전용 241분 대비 약 301분)과 저장 공간이 필요하지만, 더 세밀한 컨텍스트를 검색함으로써 추론 시 토큰 사용량을 줄여 생성 중 계산 비용을 낮춘다.

의의 및 주장

본 논문은 검색이 본질적으로 유익한 것이 아니라, 그 성공 여부는 검색 단위의 구조와 선택 전략에 달려 있다고 주장한다. 저자들은 다음과 같이 주장한다:

• 구조가 중요하다: 코드를 평면적인 청크로 취급하는 것은 불충분하다. 유용한 지식을 최소화하면서 노이즈를 줄이기 위해서는 구문적 및 의미적 경계를 존중하는 구조적 표현(코드를 위한 AST 기반, 텍스트를 위한 JSON 기반)이 필수적이다.
• 세밀도는 트레이드오프 관계이다: 더 세밀한 검색 단위(블록)는 일반적으로 정밀도를 높이지만, 필요한 컨텍스트를 누락하지 않도록 주의 깊은 관리가 필요하다.
• 선택이 핵심이다: 가장 효과적인 시스템은 "최선의 단일 컨텍스트"에 의해 정의되는 것이 아니라, 구조화된 검색을 통해 후보 다양성을 생성하고 견고한 재순위화를 통해 검색 유도 오류를 필터링하는 파이프라인에 의해 정의된다.
• 컨텍스트의 한계: PKG의 이점은 균일하지 않다. 이는 모델 아키텍처, 문제 주제(예: 수학/정렬에서의 강력한 이득, 문자열 조작에서의 잠재적 퇴보)에 따라 달라진다.

저자들은 검색 증강 코드 생성은 적절한 세밀도로 작업 관련 산물을 구축하는 것과 검색 노이즈를 처리하기 위해 후보들 사이에서 견고하게 선택하는 것 사이의 결합된 문제로 다루어져야 한다고 결 결론짓는다.