기술 요약: LLM의 코드 생성을 위한 프로토타입 기반 해석 가능성 (Prototype-Driven Interpretability for Code Generation in LLMs)

문제 정의

대규모 언어 모델(LLM)이 코드 생성 능력을 크게 발전시켰음에도 불구하고, 자동화된 솔루션에 대한 의존은 최적화되지 않거나 보안에 취약한 코드를 생성할 위험을 수반합니다. 특히 코드 생성 작업에 특화된 모델의 해석 가능성 측면에서 중요한 격차가 존재합니다. 기존의 해석 가능성 방법론들은 대개 일반적이거나 코드에 적용하기에는 계산 비용이 너무 많이 듭니다. 예를 들어, ASTrust와 같은 방법은 매 생성 단계마다 전체 어휘 집합(vocabulary)에 대한 확률 분포를 저장해야 하므로 출력 길이에 따라 확장성이 떨어집니다. 반대로, Code-Q와 같은 방법은 반복적인 샘플링을 필요로 하여 추론 오버헤드를 발생시킵니다. 또한, 인컨텍스트 러닝(In-Context Learning, ICL) 데몬스트레이션의 선택이 모델 성능에 결정적인 영향을 미친다는 사실이 알려져 있습니다. 잘못 선택된 예시는 베이스 모델보다 성능을 저하시킬 수 있지만, 효율적이고 원칙적인 샘플링 전략은 여전히 충분히 탐구되지 않은 상태입니다.

방법론

저자들은 모델의 성능과 해석 가능성을 모두 향상시키기 위해 프로토타입 기반 ICL 샘플링과 추상 구문 트리(Abstract Syntax Tree, AST)에 기반한 속성 부여(attribution)를 결합한 프레임워크인 Protocode를 제안합니다. 이 방법론은 두 가지 주요 단계로 구성됩니다.

1. 공동 매니폴드 및 메트릭 학습을 통한 프로토타입 샘플링

데이터의 고유한 기하학적 구조를 포착하지 못하는 통계적 요약(예: 클래스 평균) 대신, 저자들은 매니폴드 관점을 채택합니다. 이들은 고차원 표현이 저차원 매니폴드 위에 존재한다고 가정합니다.

• 매니폴드 구축: 이 방법은 **조각별 선형 매니폴드 학습(piecewise-linear manifold learning)**을 사용합니다. 이는 앵커 포인트 주변에 $m$-차원 선형 서브매니폴드를 구축함으로써 비선형 매니폴드를 국소적으로 선형인 영역으로 분해합니다. 이는 근접 이웃을 반복적으로 추가하고 주성분 분석(PCA)을 통해 재구성 품질을 검증함으로써 수행됩니다.
• 메트릭 학습: 프로토타입이 클래스 간에 변별력을 갖도록 하기 위해, 저자들은 Proxy-Anchor loss를 통합합니다. 이는 학습 가능한 클래스 수준의 프록시 벡터를 도입하여 클래스 내 응집도와 클래스 간 분리도를 최적화합니다.
• 공동 최적화: 신경망 $h_\theta$는 매니폴드 학습 목적 함수(국소적 기하 구조 보존)와 Proxy-Anchor loss(클래스 분리 강제)를 공동으로 최적화하도록 훈련됩니다. 훈련 후, 모멘텀 업데이트된 프록시 벡터는 최종 ICL 데몬스트레이션(프로토타입) 역할을 하기 위해 가장 가까운 훈련 인스턴스로 매핑됩니다.

2. AST 기반 해석 가능성을 위한 프로토타입-그래디언트 속성 부여

토큰 확률 기반 방법의 메모리 오버헤드를 해결하기 위해, 저자들은 전체 어휘 분포를 저장하지 않는 속성 부여 메커니즘을 제안합니다.

• 속성 부여 메커니즘: 영향력은 샘플링된 프로토타입 임베딩과 생성된 코드의 토큰 임베딩 사이의 유사도 그래디언트를 사용하여 추정됩니다. 구체적으로, 평균 프로토타입 벡터와 토큰의 잠재 표현 사이의 내적(dot product)을 계산하고, 토큰 임베딩에 대한 이의 그래디언트를 신뢰도 점수로 사용합니다.
• AST 분석: 이러한 토큰 수준의 신뢰도 점수는 추상 구문 트리(AST)를 통해 전파됩니다. 이 과정은 토큰을 터미널 AST 노드에 정렬한 다음, 이를 비터미널 노드로 계층적으로 집계하고, 최종적으로 8개의 인간이 해석 가능한 구문 카테고리(예: 결정(Decisions), 데이터 구조(Data Structures), 반복(Iterations), 예외 처리(Exception Handling))로 그룹화합니다. 이를 통해 로컬(노드 수준) 및 글로벌(카테고리 수준) 해석 가능성을 모두 제공합니다.

주요 기여

• 통합 프레임워크: 본 논문은 프로토타입 기반 ICL 샘플링과 코드 해석 가능성을 위한 AST 기반 속성 부여를 결리한 최초의 연구라고 주장합니다.
• 효율적인 샘플링 전략: 제안된 공동 매니폴드 및 메트릭 학습 방식은 국소 데이터 구조에 기하학적으로 충실하면서도 의미론적으로 변별력 있는 ICL 데몬스트레이션을 샘플링하며, 반복적인 샘플링이나 전체 어휘 확률 저장에 따른 계산 비용을 피합니다.
• 메모리 효율적인 속성 부여: 프로토타입-그래디언트 속성을 활용함으로써, 토큰 확률을 사용하는 대신 출력 길이에 따라 성능이 저하되는 메모리 오버헤드 없이 구문 인식 신뢰도 매핑을 가능하게 합니다.
• 구문 기반 설명: 이 프레임워크는 생성된 코드를 해석 가능한 구문 영역으로 분할하여, 어떤 부분이 선택된 데몬스트레이션에 의해 가장 큰 영향을 받았는지 식별하는 사후(post-hoc) 설명 시스템을 제공합니다.

실험 결과

저자들은 Qwen3, Llama3.2, Falcon3와 같은 범용 모델과 Starcoder, Qwen2.5-Coder, CodeLlama와 같은 코드 특화 모델을 포함한 6개의 LLM을 대상으로 MBPP 및 MBPP+ 테스트 세트에서 방법을 평가했습니다.

• 성능: 제안된 프로토타입 기반 샘플링 방식은 대부분의 모델에서 pass@10 지표에 대해 베이스 모델(ICL 없음) 및 다른 샘플링 전략(다양성 기반, 유사도 기반, MBPP 퓨샷)보다 일관되게 우수한 성능을 보였습니다. 예를 들어, MBPP 데이터셋에서 제안된 방법은 CodeLlama-7B에 대해 0.122의 pass@10을 달ing했으며, 이는 베이스 모델의 0.116과 비교됩니다.
• 데몬스트레이션 민감도: 결과에 따르면 코드 생성 모델은 ICL 데몬스트레이션의 품질에 매우 민감합니다. 잘못 선택된 데몬스트레이션(예: 특정 맥락에서 유사도 기반 또는 다양성 기반 샘플링을 통한 것)은 성능을 저하시켜, 때로는 베이스 모델보다 낮은 결과를 초래하기도 했습니다.
• AST 분석 통찰: 분석 결과 구문 카테고리에 따라 뚜렷한 신뢰도 패턴이 나타났습니다. 모델은 일반적으로 Scope(범위), Data Structures(데이터 구조), **Functions(함수)**에서 높은 신뢰도를 보였습니다. 반면, **Exception Handling(예외 처리)**은 모든 모델에서 일관되게 가장 낮은 신뢰도 점수를 받았으며, 이는 견고한 에러 처리 로직을 생성하는 데 있어 체계적인 한계가 있음을 시사합니다.

의의 및 주장

본 논문은 효율적인 샘플링 전략이 단순히 선택 사항이 아니라 모델 성능에 필수적임을 강조합니다. 저자들은 자신들의 접근 방식이 고품질의 기하학적으로 충실한 프로토타입이 정확도를 높일 뿐만 아니라, 생성 과정의 해석 가능성도 향상시킨다는 것을 입증한다고 주장합니다. 생성된 코드의 특정 구문 영역을 샘플링된 프로토타입과 연결함으로써, 개발자는 모델 출력의 "이유"를 더 잘 이해할 수 있습니다.

저자들은 프로토타입이 Magicoder-OSS-Instruct 데이터셋에서만 독점적으로 샘플링되었다는 점을 언급하며 연구 범위에 대해 겸허한 태도를 유지했습니다. 이들은 자신들의 방법이 프로토타입 효능에 관한 데이터셋 품질을 순위 매기는 글로벌 메트릭 역할을 할 수 있음을 시사하며, 향-후 연구에서 이를 더 넓은 데이터셋으로 확장할 수 있다고 제안합니다. 또한, 이 프레임과이 설계 단계부터 "사전(pre-hoc) 해석 가능성"을 향해 확장될 수 있으며, 프로토타입 스티어링(steering)을 통해 모델의 행동에 영향을 미칠 수 있다고 제안하지만, 이는 현재 구현된 기능이 아닌 미래의 연구 방향입니다.