🍭

(220210) Diary: Head Selection on Transformer & Variational Inference

Head Selection on Transformer using Variational Inference

(Meta, 2021) Pay Better Attention to Attention: Head Selection in Multilingual and Multi-Domain Sequence Modeling

Transformer Layer의 Key Component를 고르자면 역시 Multi-Head Attention이 아닐까 싶다. 데이터 간의 상관성을 다양한 관점에서 계산하여 복잡한 Data Patterns을 포착한다.

(별개의 이야기로) NLP에서 Multi-Domain (이하 Domain이라 작성하지만 Lingual의 의미도 내포함) Training이 빈번하게 사용되는데, Domain 간의 유사성이 낮은 경우, 개별 Domain에서의 성능이 저하되는 Negative Interference가 발생한다.

Negative Interference의 한 원인으로 Domain 간 Shared Parameters에서 발생하는 Gradients Conflict를 생각할 수 있는데, 본 논문에서는 Domain별로 서로 다른 Heads를 선택하는 Head Selection 방법론을 제안하며 해당 문제를 해결한다.

(Domain Experts와 같은 느낌으로 LM을 Modular하게 활용)

이 때, Domain 별로 어떤 Heads를 선택할 것인지를 학습하기 위해 Variational Inference를 사용한다.

(VAE에서의 그 친구가 맞음!)

Input Sequence로부터 해당 (n)차원이 n번째 Head를 선택하는지를 표현하는, Bernoulli 분포의 Discrete Latent Variable, z를 모델링할 때 각 Head가 선택될 확률이 같음을 가정하여 Inference Network를 학습한다.

또한, 모델이 End-To-End 미분 가능하도록 Gumbel-Softmax를 차용한다.

(Tabular Data에 관심 있을 적부터 리뷰하려 하였으나, 여태까지 하지 않은 Gumbel-Softmax.. 올해에는 공부할 예정!)

Selection 과정에서는 선택 확률이 높은 Top-H개의 Heads를 선택하게 된다.

Heads의 순서를 고려하지 않는 Subset Strategy, 고려하여 선택하는 Group Strategy가 존재한다. (위 그림 참조)

FAIR의 Multi-Lingual 혹은 Multi-Domain 실험은 개인적으로 너무 생소하여 분석이 어려움..

전반적으로 좋은 성능을 보이는 것으로 생각되어 자세한 내용은 생략!

이 논문을 리뷰한 이유는 Variational Inference (VAE)를 처음 접했을 때 매우 흥미롭게 공부했던 기억이 있고, 다루기 쉬운 분포의 Latent를 활용하여 Output을 조절하는 점이 현재 업무 (Controlled Generation)에 영감을 주지 않을까 해서이다.

예상대로 흥미로웠고, 당분간 비슷한 연구들을 찾아 공부할 예정이다.