[ICLR 2020] Depth-Adaptive Transformer 논문 정리

ICLR 2020 NLP 논문 읽기 (5)

[paper]

Abstract만 읽고 정리해보기

• 최신 large-scale task에 대한 sequence-to-sequence은 각 input sequence가 처리하기 쉬운지 어려운지와 상관없이 고정된 수의 연산(computation)을 수행
• 따라서, 우리는 Transformer 모델의 다른 단계에서 output prediction을 할 수 있게 학습하고 특정 sequence에서 얼마나 많은 연산이 필요한지 예측하는 다양한 방법을 연구
• 동일한 집합의 layer를 반복적으로 적용하는 Universal Transformers의 dynamic computation과 달리, 우리는 모든 단계에서 서로 다른 레이어를 적용하여 연산 양과 모델 크기 둘 다 조정
• IWSLT German-English translation에서 이 접근법은 decoder 층의 1/4 미만을 사용하면서 기존 잘 학습된 Transformer 만큼의 정확도를 보임

Introduction

• 요즘 neural sequence model의 size: billions of parameters
• WMT’19 news machine translation task winning entry: two billion parameters의 앙상블 모델
• 큰 모델은 어려운 examples에 대해서 잘 수행하고, 작은 모델은 쉬운 examples에 대해서 잘 수행함 → 하지만, 현재 모델은 입력이 쉽든 어렵든 상관없이 같은 양의 연산을 수행

따라서, 추론에 적절한 speed-accuracy trade off을 달성하기 위해 각 입력에 layer 수를 맞추는 Transformer를 제안 ← Graves (2016; ACT) 의 extend version (ACT는 RNN에 dynamic computation을 도입한 모델)

• 각 단계 별로 서로 다른 레이어를 적용하고, halting module의 다양한 디자인과 학습 타겟을 연구
• 대규모 task에서 좋은 성능을 얻기 위해 간단한 oracles을 통해 명시적으로 지도 (oracles: a definition for an entity who knows the correct answer to all questions)

Related work 소개 + 우리는 다르다!

• Universal Transformers (UT)또한 ACT의 dynamic computation에 의존
  하지만, UT는 같은 layer에 dynamic computation을 반복적으로 적용
  → 네트워크 깊이를 추정하기 위한 다양한 방법을 고려하여 각 단계에서 다른 layer를 적용
• Dehhgani et al. (2018) 은 단계 수 (the number of steps)를 고정
  → 정확도를 잃지 않고 상당한 속도 향상을 보여주기 위해 단계 수를 다양하게 함
• UT는 standard Transformer만큼 많은 weight를 포함하는 layer를 사용
  이는 여러 번 적용되기 때문에 속도에 영향을 미침
  → 개별 layer의 크기를 증가하지 않음
• Huang et al. (2017) , Bolukbasi et al. (2017) 의 자원 효율적인 객체 분류(resource efficient object classification)를 dynamic computation 결정이 미래 computation에 영향을 미치는 구조화된 예측으로 확장
• Computer vision 관련 work는 Teerapittayanon et al. (2016), Figurnov et al. (2017), Wang et al. (2018)이 있음
  일찍 exit하거나 layer를 skip하는 동적 routing의 아이디어를 탐구

Challenge

• encoder는 standard Transformer encoder를 사용하고, decoder network에서 다양한 연산 양으로 수행
• Dynamic computation의 경우, self-attention에 어려움이 있음 ← 향후 이전 time-steps에서의 생략된 layer가 필요할 수 있기 때문
• 이를 해결하기 위해 두 가지 접근법을 갖고 실험 (2장)
• Decoder network에서 전체 시퀀스에 대해 또는 토큰 단위로 계산량을 조절하기 위한 다른 방법을 연구 (3장) — model likelihood에 의해 supervise 되는 multinomial, binomial classifiers 또는 argmax가 단순히 model score를 thresholding 하는 것 뿐만 아니라 이미 올바른지 여부
• IWSLT 14 German-English translation와 WMT’14 English-French translation에서의 실험 결과 (4장) — well tuned baseline models의 성능이면서 76% less computation

Anytime Structured Prediction

• 서로 다른 layer에서 prediction을 하는 모델
• computer vision에서는 anytime prediction (Huang et al., 2017)으로 잘 알려져 있음, 이를 structured prediction으로 확장

Transformer with multiple output classifiers

Training multiple output classifiers

• Dynamic computation은 모델이 단지 1개의 final classifier 대신, N개의 exit classifier를 사용할 수 있도록 함.
• 그 중 몇 모델은 각 time-step마다 서로 다른 classifier를 선택할 수 있음
  이는 sequence length 내 가능한 output classifier의 조합이 exponential number가 됌
• decoder network을 학습하기 위한 두 가지 가능한 방법을 고안 (아래 Fig)

Training regimes for decoder networks able to emit outputs at any layer

• Aligned training은 모든 classifier를 동시에 최적화하고 self-attention에서 요구되는 이전 hidden states를 이용할 수 있다고 가정
  하지만, test time에서 모든 token에 대해 다른 exit을 선택하여 잘못 정렬된 상태가 되는 경우가 종종 발생(?)
• 대신, mixed training은 주어진 문장에 대해 몇 개의 sequences of exits을 샘플링하고 다른 layer의 hidden states에 모델을 노출시킴

Adaptive Depth Estimation

• 모델이 어디서 멈추고 다음 token을 출력할 decoder block을 예측하기 위한 매커니즘
• 또한, 언제 모델이 멈춰야 좋은 speed-accuracy trade-off을 달성하는 지
• Sequence-specific depth: 동일한 block을 사용하여 모든 출력 token을 decode
• Token-specific depth: 각 개별 token에 대하여 별도의 exit을 결정

Variants of the adaptive depth prediction classifiers

Sequence-specific depth

Multinomial classifier가 어느 depth에서 멈출 지 결정 (그림에서는 depth 2에서 멈춤)

• Likelihood-based: the likelihood of the entire sequence after each block
  Problem: 모델이 이미 가장 높은 점수를 올바른 target에 할당했는지 여부를 무시
• Correctness-based: 정확한 예측에 가장 큰 점수를 할당하는 가장 낮은 블록을 선택
  각 block에 대해 sequence에 대해 올바르게 예측된 token의 수를 세고, 가장 많은 수의 올바른 token을 가진 block을 선택

Token-specific depth

Token-specific — Multinomial: 각 token 별로 첫번째 block 다음에 멈출 depth를 예측

Token-specific — Geometric-like: 각 token 별로 매 block마다 멈출지 말지를 binary decision으로 결정 (C: continue, S: stop)

• Likelihood-based: exit classifier가 가장 높은 likelihood + regularization term 을 갖는 block을 선택
  Problem: 현재 결정이 미래 time-steps에 미치는 영향을 무시
  → RBF 커널로 likelihood를 완화
• Correctness-based: time-step t와 주변 위치에서의 prediction의 correctness

Confidence thresholding

• model predictions에 thresholding
• 학습 없이 valid set에서 BLEU를 maximize하는 threshold 선택

My Discussion

이전에 정리했던 Reformer와 같이 Transformer의 변형 중 하나이다. Depth-Adaptive Transformer는 제목 그대로 Depth를 Adaptive하게 조정할 수 있는 Transformer를 제안하여 decoder layer의 약 3/4를 reduce하는 효과를 보였다. Adaptive Depth Estimation 파트에서 얼마나 많은 고민을 하고 고생, 실험을 했는지 논문에서 느껴졌다. 확실히 Introduction을 주의 깊게 읽은 다음에 전체를 읽으니 이해가 빨라졌다.