[논문 리뷰] Multi-task Learning based Pre-trained Language Model for Code Completion

Fang Liu, Ge Li외 2명|arXiv (Cornell University)|2020. 12. 29.

Software Engineering Research참고 문헌 46인용 수 8

한 줄 요약

이 논문은 트랜스포머 아키텍처를 기반으로 한 다중 작업 학습 기반 사전 훈련된 언어 모델을 제안하여, 토큰과 그 타입을 동시에 예측함으로써 정적 임베딩의 한계와 열악한 식별자 완성 문제를 해결함으로써 코드 완성 성능을 향상시킨다. 이 모델은 두 개의 실세계 데이터셋에서 최신 기술 수준의 성능을 달성하며, 특히 정확한 식별자 식별 및 제안에 뛰어난 성능을 보인다.

ABSTRACT

Code completion is one of the most useful features in the Integrated Development Environments (IDEs), which can accelerate software development by suggesting the next probable token based on the contextual code in real-time. Recent studies have shown that statistical language modeling techniques can improve the performance of code completion tools through learning from large-scale software repositories. However, these models suffer from two major drawbacks: a) Existing research uses static embeddings, which map a word to the same vector regardless of its context. The differences in the meaning of a token in varying contexts are lost when each token is associated with a single representation; b) Existing language model based code completion models perform poor on completing identifiers, and the type information of the identifiers is ignored in most of these models. To address these challenges, in this paper, we develop a multi-task learning based pre-trained language model for code understanding and code generation with a Transformer-based neural architecture. We pre-train it with hybrid objective functions that incorporate both code understanding and code generation tasks. Then we fine-tune the pre-trained model on code completion. During the completion, our model does not directly predict the next token. Instead, we adopt multi-task learning to predict the token and its type jointly and utilize the predicted type to assist the token prediction. Experiments results on two real-world datasets demonstrate the effectiveness of our model when compared with state-of-the-art methods.

연구 동기 및 목표

• 소스 코드에서 맥락에 따라 달라지는 의미를 포착하지 못하는 정적 워드 임베딩의 한계를 해결하기 위해.
• 예측 시 타입 정보를 통합하여, 특히 식별자에 대한 코드 완성 성능을 향상시키기 위해.
• 코드 이해 및 코드 생성 작업을 모두 지원하는 통합 사전 훈련 모델을 개발하기 위해.
• 사전 훈련 중 하이브리드 목표 함수를 활용한 다중 작업 학습을 통해 맥락 기반 표현 학습을 향상시키기 위해.
• IDE에서 실시간 코드 완성에 대해 공동 토큰 및 타입 예측의 효과성을 입증하기 위해.

제안 방법

• 모델은 코드 토큰의 맥락 기반 표현을 학습하기 위해 트랜스포머 기반 신경 아키텍처를 사용한다.
• 마스크된 언어 모델링, 다음 토큰 예측, 타입 예측의 세 가지 목적 함수를 사용하여 대규모 코드 리포지터리에서 사전 훈련한다.
• 다중 작업 학습을 통해 다음 토큰과 그 데이터 타입을 동시에 예측하며, 타입 예측이 토큰 생성을 안내한다.
• 피니튜닝 단계에서 모델은 토큰 및 타입 예측 정확도 최적화를 통해 코드 완성에 적합하게 조정된다.
• 모델은 양방향 맥락과 어텐션 메커니즘을 활용하여 코드 시퀀스 내 장거리 종속성을 포착한다.
• 코드 이해 및 코드 생성 목표를 결합한 하이브리드 사전 훈련 전략을 통해 일반화 능력을 향상시킨다.

실험 결과

연구 질문

• RQ1 단일 작업 모델에 비해 공동 토큰 및 타입 예측이 코드 완성 정확도를 향상시키는가?
• RQ2 예측 시 타입 정보를 통합하면 식별자 완성 성능이 향상되는가?
• RQ3 다중 작업 사전 훈련된 언어 모델이 기존의 정적 임베딩 기반 모델보다 코드 완성에서 뛰어난 성능을 보일 수 있는가?
• RQ4 제안된 모델은 실세계 코드 완성 벤치마크에서 얼마나 효과적인가?
• RQ5 모델은 다양한 프로그래밍 언어와 코드 패턴에 대해 일반화 가능한가?

주요 결과

• 제안된 모델은 두 개의 실세계 코드 완성 데이터셋에서 최신 기술 수준의 성능을 달성하며, 이전의 최신 기술 수준 모델들을 능가한다.
• 모델은 식별자 완성 정확도를 크게 향상시켜 이전의 언어 모델 기반 접근 방식의 핵심 한계를 해결한다.
• 토큰과 타입의 공동 예측은 코드 완성에서 더 정확하고 맥락에 부합하는 제안을 가능하게 한다.
• 다중 작업 사전 훈련 전략은 희귀하거나 모호한 식별자에 대해 특히 표현 학습 능력을 향상시킨다.
• 모델은 다양한 코드 패턴과 프로그래밍 구조에 걸쳐 뛰어난 강건성과 일반화 능력을 보여준다.
• 제거 실험 결과, 타입 예측 및 다중 작업 학습의 포함 여부가 성능 향상에 가장 큰 기여를 한다는 것이 확인되었다.