AUBER: Automated BERT Regularization

본 문서에서는 AUBER: Automated BERT Regularization 논문을 소개합니다. 논문에 대한 상세한 정보는 다음과 같습니다.

BERT 정규화 (BERT Regularization)

BERT는 2018년 제안된 모델로 많은 자연어 처리 문제에서 좋은 성능을 내고 있습니다. BERT의 학습은 수많은 데이터로 사전학습(pretrain) 후, 풀고자 하는 특정 문제에 맞추어 미세조정(fine-tuning)함으로써 이루어집니다.

BERT 모델은 트랜스포머(transformer)의 인코더(encoder)가 여러 층 쌓인 구조로 이루어져 있습니다. 그리고 각 인코더는 멀티 헤드 셀프 어텐션(multi-head self-attention) 구조를 갖습니다. 간단히 말해 어텐션 헤드(attention head)라는 블록 여러 개가 나란히 놓인 구조입니다. 가장 기본적인 모델인 BERT-base의경우, 12개의 어텐션 헤드로 구성된 인코더가 12층으로 쌓인 형태를 갖습니다. 즉, 어텐션 헤드가 12x12 형태로 배치되어 있습니다.

하지만 이러한 BERT는 파라미터가 굉장히 많은 큰 모델입니다. 이로 인해 과적합(overfitting)이 발생해 성능이 낮아지곤 합니다. 특히, 10,000개 이하의 훈련 데이터를 갖는 문제에 대해서는 과적합으로 인해 학습이 전혀 되지 않는 경우도 보고된 바 있습니다. 이는 BERT의 정규화의 필요성을 보여줍니다. 본 논문에서는 프루닝(pruning)을 통해 파라미터 수를 줄임으로써 문제를 해결하고자 합니다.

강화학습과 DQN (Reinforcement Learning and DQN)

강화학습은 현재 상태(state)에서 어떤 행동(action)을 취하는 것이 최적인가를 학습을 통해 파악하는 분야입니다. 이 때, 최적이란 보상(reward)을 최대화하는 상황입니다. 어떤 상태에 대해 최적의 행동이 무엇인지를 가리키는 함수를 정책(policy)이라 칭하며, 정책을 구한 뒤 그를 바탕으로 각 상태에서 최적의 행동을 취함으로써 보상을 최대로 만드는 방식으로 활용됩니다. 따라서, 풀고자 하는 문제 상황에 맞춰 적절한 상태, 행동, 보상을 설정해야 합니다.

그 중에서도 가장 대중적으로 사용되는 DQN은 딥 뉴럴 네트워크(deep neural network)를 이용해 정책을 파악합니다. DQN 네트워크는 상태를 입력으로 주면 가능한 모든 행동에 대해 얼만큼의 보상이 기대되는지를 예상하여 출력합니다. 그를 통해, 입력된 상태에 대한 최적의 행동, 즉 정책을 파악할 수 있습니다.

Proposed Method

Overview

본 논문에서는 강화학습을 활용해 프루닝을 진행함으로써 BERT 정규화를 자동화하는 기법, AUBER를 제안합니다. BERT의 어텐션 헤드가 다수 프루닝 되더라도 성능이 크게 떨어지지 않고, 심지어는 오르기도 한다는 현상을 바탕으로 AUBER는 어텐션 헤드 단위로 프루닝을 진행합니다. 이 때, 어떤 어텐션 헤드가 프루닝되는 것이 정규화에 유리할지를 파악하기 위해 강화학습을 사용합니다. 본 논문의 아이디어는 다음과 같습니다.
  • 강화학습을 사용해 어떤 어텐션 헤드가 프루닝되는 것이 정규화에 도움이 되는지 파악하는 과정을 휴리스틱(heuristics)에 의존하지 않고 자동화합니다. 이 때, 강화학습에 필요한 상태,행동, 보상 등을 적절하게 공식화하였습니다.
  • 고려해야 하는 프루닝 방법의 수가 지나치게 많아 모델 학습이 어렵다는 문제를 해결하기 위해 각 층별로 프루닝을 진행합니다.

어떤 문제 상황에 대해 미리 학습된 BERT 모델이 주어졌을 때, 첫 번째 층부터 마지막 층까지 순차적으로 어텐션 헤드 프루닝을 진행합니다. 프루닝은 DQN 모델의 지시에 따라 이루어집니다. DQN 모델은 현재 다루고 있는 층의 상태를 넣어주면, 해당 상태에서 어떠한 행동을 취하는 것이 좋을지를 알려줍니다. 이를 통해 해당 층의 프루닝이 완료되면 미세조정을 통해 프루닝 중 발생한 정보 손실을 보완합니다. 이 과정을 도식화하면 아래와 같습니다.

AUBER: Automated BERT Regularization

강화학습을 이용해 어떤 어텐션 헤드를 프루닝할지를 결정하기 위해서는 상태, 행동, 보상을 공식화하여야 합니다. 본 논문에서는 이들을 다음과 같이 공식화합니다.

  • 상태: 각 층의 상태는 한 층에 있는 어텐션 헤드 수 만큼의 원소를 갖는 벡터로 표현됩니다. 어떤 층을 처음으로 프루닝하는, 다시 말해 아직 프루닝하지 않은 상황에서는 해당 층의 어텐션 헤드들의 중요도를 반영한 벡터로써 상태를 나타냅니다. 각 어텐션 헤드는 밸류 행렬(value matrix)이라는 행렬을 포함합니다. 본 논문에서는 이 밸류 행렬에 각종 연산을 취함으로써 상태를 나타냅니다. 자세한 방법과 이론적인 증명은 논문[link]을 참고하시면 됩니다. 한편, 어텐션 헤드를 프루닝하는 상황에서는 대응되는 벡터의 원소를 0으로 만들어 상태를 업데이트 합니다.  
  • 행동: (한 층에 있는 어텐션 헤드 수+1)가지의 행동이 존재합니다. 예를 들어, BERT-base의 경우엔 총 13종류의 행동을 갖습니다. BERT-base를 기준으로 설명하자면, 첫번째부터 12번째 행동은 각각 대응되는 어텐션 헤드를 프루닝하는 것을 의미합니다. 한편, 마지막 행동은 프루닝을 그만두고 다음 층으로 넘어가는 것을 행동을 뜻합니다.
  • 보상: 성능을 개선하는 것이 목표이기 때문에, 행동을 취했을 때의 성능에서 행동을 취하지 않았을 때의 성능을 뺀 값이 보상이 됩니다. 다시 말해, 해당 행동을 취함으로써 얻는 성능 개선을 보상으로 둡니다.
첫 번째 층부터 순차적으로 위와 같이 공식화된 강화학습으로 어텐션 헤드 프루닝을 진행합니다. 각 층에 대한 프루닝이 완료되면 미세조정을 함으로써 정보 손실이 없도록 합니다. 이러한 과정을 마지막 층까지 반복함으로써 BERT 모델을 정규화합니다.

Experiments

본 논문에서는 훈련 데이터가 10,000개 이하인 네 가지의 자연어 처리 문제에 대해 다른 어텐션 헤드 프루닝 방식과 비교한 실험 결과를 제시합니다. 그 결과를 도식화한 것은 아래 그림과 같습니다. AUBER는 모든 문제에 대해 가장 좋은 성능을 내며 정규화에 성공하는 반면에 다른 어텐션 헤드 프루닝 방법들은 종종 성능을 떨어뜨리기도 하며 정규화에 실패하는 것을 확인할 수 있습니다.

그와 더불어 어텐션 헤드 프루닝이 어떤 방식으로 성능을 개선하는지를 확인하고자 MRPC 데이터에 대해 그 과정을 더욱 깊이 알아보았습니다. 각 층을 프루닝 후 미세조정할 시 성능이 어떻게 달라지는지를 도식화한 것은 아래와 같습니다. 다른 방법들은 모두 성능이 프루닝 전에 비해 감소하였으나, AUBER로 프루닝을 진행하자 일관되게 성능이 높아지며 좋은 성능에 도달하였습니다.

Conclusion

본 문서에서는 AUBER: Automated BERT Regularization 논문을 소개하였습니다. 해당 논문은 어텐션 헤드를 프루닝하여 BERT 모델을 정규화하였고, 이 과정을 강화학습을 활용해 자동화하였습니다. 또한, 실험을 통해 기존의 어텐션 헤드 프루닝 방법들보다 뛰어난 정규화를 이루었음을 확인하였습니다. 해당 논문에 대한 자세한 정보는 링크를 통해 확인할 수 있습니다.