[Paper] Diet Code Is Healthy: Simplifying Programs for Pre-Trained Models of Code

안녕하세요 왕새우 여러분, 오늘 살펴볼 논문은 2022년 ESEC/FSE에 발표될 Diet Code Is Healthy: Simplifying Programs for Pre-Trained Models of Code 입니다. Program Simplification을 통해 source code 관련 task 에서의 computation을 줄이는 것이 목표입니다.

 

요새는 Software Engineering task들을 해결하기 위해 언어 모델이 많이 사용되고 있습니다.  또한 모든 task에 대하여 전부 처음부터 학습시키는 경우보다 미리 어느 정도 학습을 시켜놓은 모델을 가져와서 학습을 시키는 pre-trained model을 많이 사용합니다. 

 

헌데 본 논문은, 이러한 pre trained 모델을 학습시키는데 너무 많은 자원이 필요하다고 말합니다. CodeBERT의 경우 125,000,000 개의 파라미터가 필요하고, 학습시키는데 28 시간이 걸린다고 합니다. 회사차원에서는 충분히 scalable 하겠지만, 개인 차원에서는 조금 버거운게 맞는 것 같습니다. 

 

1. CodeBERT

우선 CodeBERT에 대하여 간략하게 알아보고, 어떻게 computation을 줄일지 살펴보겠습니다.

1-1. CodeBERT는 code와 natural language 의 bi-modal 한 데이터, 코드만 있는 uni-modal 한 데이터 두 종류를 사용합니다. 

bi-model 데이터. 출처: 논문

1-2. MLM (Masked Language Model)

Pretrain task 로, MLM을 사용합니다. 랜덤 두 토큰을 가리고, 맞춰보도록 합니다. 

이런 식으로 <MASK>를 맞추게 시킵니다.. 출처: CodeBERT github

1-3. RTD (Replaced Token Detection)

출처: ELECTRA 논문

NL, Code Generator를 따로 두고, MASK 토큰에 대한 예측을 하게 한 후, pretrain 모델이 각 토큰이 generator에 의해서 만들어진 것인지, 오리지널 토큰인지 맞추도록 합니다. RTD task는 효율적인 학습을 도울 수 있는 것으로 알려져 있습니다. RTD의 장점에는,

1) 모든 토큰에서 loss가 나와서 학습이 효과적 (MLM은 마스크에서만 나옴)

2) inference시에는 MASK 가 없는데, 실제 use case와 비슷한 환경

이 있습니다.

 

2. token, stmt attention

논문에서는 pretrained model에서 높은 attention 값을 갖는 토큰이 중요하다고 가정합니다. 이를 활용하기 위해 우선 각 token, 그리고 source code의 statement 마다 attention 값을 만듭니다. Token attention 값은 head, layer 마다 평균이 취해지게 되고, stmt attention 값은 내부에 속한 token 값의 weighted average 가 취해집니다.

아래 그림에 모아진 attention 값에 따라 highlight 을 해두었습니다. (blue: token, red: stmt)

출처: 논문

3. Simplification Method

세 가지 simplification method를 제안합니다.

3-1. Word Dropout

출처: 논문

랜덤하게 input dropout

 

3-2. Frequency Filtering

top-k 만 남기고 drop

 

3-3. Attention based code pruning (메인)

우선, 각 statement를 가방, 가방의 value를 attention value, 가방의 weight을 length 로 두고 knapsack problem 을 풉니다. 

여기서 총 가방의 weight 을 원하는 길이보다 조금 더 길게 잡습니다.

그 이후에는, 원하는 length 까지 greedy 하게 attention값이 낮은 token을 drop 합니다.

(b): knapsack 푼 이후, (c): token drop 이후 출처: 논문

4. 결과

출처: 논문

결과적으로 CodeBERT, CodeT5에 대하여 DietCodeBERT, DietCodeT5를 보면, Finetuning Time (FT Time), 그리고 FLOPs가 크게 줄어듬에도 MRR (Code search metric)은 크게 감소하지 않는 것을 볼 수 있죠. 본 논문의 방식을 사용하면 들어온 코드를 확 줄여버림으로써 메모리나 시간을 아낄 수 있으면서도, 성능은 어느 정도 유지할 수 있다는 것을 보였네요.