[논문리뷰] R-Drop: Regularized Dropout for Neural Networks (with Python)

AI4Code 캐글 대회에서 R-Drop 방법을 사용해 막판에 성능을 끌어올릴 수 있었다. 따라서 오늘은 R-Drop 논문에 대해 리뷰해보고자 한다.

 

논문 링크

https://arxiv.org/pdf/2106.14448.pdf

 

 

---

다양한 정규화 기술들은 딥러닝 모델의 과적합을 방지한다. 그 중, Dropout은 네트워크 아키텍쳐와 작업 유형에 관계없이 가장 일반적으로 사용되는 기술 중 하나이다. Dropout은 훈련 중 신경망에 숨겨진 Unit의 일정 비율을 떨어뜨리는 것으로 구현된다.

 

논문 저자는 이런 Dropout의 무작위성을 활용하도록 설계된 간단하지만 강력한 정규화 기법인 R-Drop을 소개한다. 소개된 정규화 기법은 Vanila Transformer 모델을 사용하여 WMT14 번역 데이터 셋에서 SOTA를 달성했다고 한다.

 

R-Drop regularization

R-Drop regularization 방법은 간단하다. 훈련단계에서 데이터 셋이 모델을 두번 통과하게 된다. 두 출력 사이의 양방향 Kullback-Leibler(KL) 발산을 최소화하여 두 출력 분포가 일관되도록 훈련한다.

 

 

* 쿨백-라이블러 발산(Kullback–Leibler divergence, KLD)은 두 확률분포의 차이를 계산하는 데에 사용하는 함수로, 어떤 이상적인 분포에 대해, 그 분포를 근사하는 다른 분포를 사용해 샘플링을 한다면 발생할 수 있는 정보 엔트로피 차이를 계산한다. 상대 엔트로피(relative entropy), 정보 획득량(information gain), 인포메이션 다이버전스(information divergence)라고도 한다.

 

 

따라서 기존의 cross entropy를 쓰는것과 다른점은 다음과 같다.

• 데이터 셋이 모델을 두번 통과하게 되어 예측 결에 대한 logit 값을 2번 추출한다.
• 양방향 KL 발산 항이 손실함수에 2번 추가 된다.

 

최종 손실 항은 negative log-likelihood loss(cross-entropy) L_NLL과 양방향 KL 발산 L_KL을 결합한 아래 방정식과 같다. KL 발산은 KL(P_1, P_2), KL(P_2, P_1) 양쪽에서 측정되고 평균이 계산됩니다. α는 L_KL을 제어하기 위한 중량 계수이다.

 

 

 

코드는 다음과 같이 구현될 수 있다.

from bisect import bisect

def compute_kl_loss(p, q, pad_mask=None): 
    p_loss = F.kl_div(F.log_softmax(p, dim=-1), F.softmax(q, dim=-1), reduction='none')
    q_loss = F.kl_div(F.log_softmax(q, dim=-1), F.softmax(p, dim=-1), reduction='none')
    
    if pad_mask is not None:
        p_loss.masked_fill_(pad_mask, 0.)
        q_loss.masked_fill_(pad_mask, 0.)

    p_loss = p_loss.sum()
    q_loss = q_loss.sum()

    loss = (p_loss + q_loss) / 2
    return loss

 

 

실제 사용은 아래와 같이 사용하면 된다.

 

... 생략 ...
    criterion = torch.nn.L1Loss() #L1, L2, MSE, ...
    scaler = torch.cuda.amp.GradScaler()
    
    print_idx = 0
    score_tb = pd.DataFrame() # 성능 log 기록용 테이블
    
    for e in range(epochs):
        model.train()
        tbar = tqdm(train_loader, file=sys.stdout)
        loss_list = []
        labels = []

        for idx, data in enumerate(tbar):
            inputs, target = read_data(data)

            if r_dr == True: # R-Droup을 쓸 경우
                with torch.cuda.amp.autocast():
                    # 각 모델별 logit값 생성
                    logits = model(*inputs)
                    logits2 = model(*inputs)

                    ce_loss = 0.5 * (criterion(logits, target) + criterion(logits2, target))
                    kl_loss = compute_kl_loss(logits, logits2)

                    loss = ce_loss + TRAIN_CFG['kl_alpha'] * kl_loss

            else: # R-Droup을 안쓸 경우
                with torch.cuda.amp.autocast():
                    pred = model(*inputs)
                    loss = criterion(pred, target)
                    
 ... 생략 ...

 

 

저자들은 훈련 비용을 절약하기 위해 입력 미니 배치 x'는 ([x, x])의 연결이며 inference는 실제로 두 번 호출되지 않는다고 제안한다.

 

Theoretical Analysis

저자가 소개한 R-Drop의 정규화 효과를 이론적으로 분석해보고자 한다.

저자들은 L_KL이 네트쿼크에서 매개 변수의 자유를 제약한다는 이론적인 분석을 제공한다. 수식은 다음과 같이 정의된다.

 

 

h¹(x): 입력 x가 있는 NN의 l번째 레이어의 출력값

ξ¹: 베르누이 분포 B(p)에서 추출한 랜덤 벡터

Dropout은 다음과 같이 표현될 수 있다.

h¹_ξ¹ (x) = (1/p) ξ¹ ⊙ h¹(x)

 

신경망은 아래와 같이 표현될 수 있다.

 

 

드롭아웃으로 손실을 최소화하는 것(L_NLL)은 드롭아웃을 사용하지 않을 때에 비해 모델의 야코비안 행렬(Colin et.al, 2020)을 제어하여 모델 복잡성을 제한하는 것으로 이미 알려져 있다. 

 

저자들은 다층 퍼셉트론과 트랜스포머와 같은 fully-connected 신경망의 경우, 양방향 KL 발산 손실의 제약은, 네트워크의 모든 매개 변수가 동일하도록 제한하는 것과 동일하다는 것을 증명한다.

 

하나의 숨겨진 노드에서만 서로 다른 두 개의 하위 구조 P1과 P2를 생각해보자. P1에만 포함되는 숨겨진 노드는 h_i이고, P2에만 포함되는 노드는 h_j이다. KL(P1 || P2)=0라고 제약조건을 걸 경우, P1(y|x) = P2(y|x)와 같으며, h_i=h_j는 이전에 대해 참이어야 한다. 따라서, KL(P1 || P2)은 h_i=h_j일 때 최소화된다.

 

"결론적으로, 혼합 손실 L을 최적화하는 것은 매개 변수의 자유도가 가장 작은 손실 L_NLL을 최소화하여 과적합을 방지하고 일반화 능력을 향상시킬 수 있다."

 

 

Ablation Study

실험의 경우 Transformer기반 모델을 활용하여 실험을 진행을때, NLP분야와 Image 분류에서 실험했을 때 R-Drop을 적용한 결과가 더 좋게 나왔다.

 

소감

AI4Code 캐글 대회에 참가했을 땐 R-Drop을 쓴게 좋은 성과를 냈지만, Image와 Text데이터가 결합된 데이콘 대회에선 (리더보드 기준) R-Drop을 안쓴 결과가 더 좋게 나왔었다.

 

차이점을 보면 AI4Code에선 Transformer기반 언어 모델을 사용했고, 데이콘 대회에선 Text의 경우 Transformer기반인 모델을 사용하고, Image의 경우 CNN 기반 모델을 사용했다.

 

따라서 R-Drop의 경우 CNN 기반 모델에선 다른 Crossentropy 손실함수를 쓰는게 더 나을것 같기도 하다.