Learning to Remember Rare Events

Title:

Learning to Remember Rare Events

Overview:

Google Brain 팀에서 ICLR 2017에서 발표할 본 논문은 memory-augmented neural network (MANN)에 관한 논문입니다. 기존의 MANN들은 rare event에 대한 generalization performance 측면에서 한계점을 갖고 있는데요, 본 논문에서는 그러한 단점을 극복하기 위한 새로운 memory module을 제안합니다. 제안하는 memory module은 구조상 확장성이 좋아서, 다양한 종류의 supervised NN 모델에 쉽게 추가할 수 있고, 언급한 rare event (e.g., extreme한 경우에는 one-shot learning)에 대한 성능 향상을 기대할 수 있습니다.

지난 포스팅에서 언급했듯이, 이 논문의 main idea는 Neural Episodic Control (NEC) 논문과 거의 동일한데요, 차이점이라고 한다면 NEC는 reinforcement learning에 이 idea를, 본 논문은 supervised learning model에 이 idea를 적용해봤다는 것 정도입니다;; 참고로, 본 논문이 NEC보다 먼저 공개된 논문입니다. (NEC는 2017.03.06에 arXiv에 올라왔고 아마도 ICML 2017에 제출한 것으로 보입니다.)

Introduction:

Deep learning 성공 스토리의 대표적 application 중 하나인 neural machine translation (NMT)의 성능은 human translation performance에 근접해 있습니다. 하지만, 이는 평균적인 성능 측면에서이고, rare word (e.g., Dostoevsky) 에 대해서는 training 시 해당 단어를 포함한 혹은 해당 단어와 관련된 example들이 부족했을 것이기에 그 성능이 평균성능에 훨씬 못 미치게 되죠. 하지만, human은 rare event에 대해서도 문제없이 잘 배웁니다.

저자들은 이 문제를 극복하기위해, one-shot learning을 가능케하는 memory module을 제안합니다. 제안하는 memory module은 key-value pair 형태의 dictionary 구조이며, 여기서 value는 given example에 대한 ground-truth target (label) 입니다.

이러한 memory module을 포함하는 다양한 supervised learning model의 경우, 학습이 진행되면 진행될 수록 memory module의 활용도가 증가하며, memory module에 저장되어있는 과거 data 를 잘 활용하여 보다 정확한 prediction을 가능하게 해줍니다. 즉, new example이 주어지면, network은 해당 example에 대한 정보를 memory module에 기록해두기 때문에 설령 관련 example이 학습시 단 한번만 등장했다고 하더라도 필요한 시점에 꺼내서 사용할 수 있기때문에, one-shot learning performance 향상을 기대할 수 있습니다.

저자들은 제안하는 방식을 몇몇 supervised learning model에 적용하여 성능비교를 해보고, 제안하는 방식의 장점을 실험적으로 보여줍니다.

Memory module (basic):

제안하는 memory module은 아래와 같이 (key, value, age)로 구성됩니다.

M = (K_{memory-size \times key-size}, V_{memory-size}, A_{memory-size})

여기서 key-size는 memory query vector size라고 생각하면 됩니다. normalized query vector q (||q|| = 1)가 주어지면, query q에 대한 nearest neighbor를 memory module M으로 부터 찾을 수 있죠.

NN(q,M) = \arg\max_i q \cdot K[i]

여기서 memory module에 저장될 key와 query key 모두 normalized vector이기 때문에, nearest neighbor w.r.t cosine similarity가 되겠죠. 이렇게 해서 top-k nearest neighbors를 찾아 sorting한 결과를 (n_1, ... , n_k) = NN_k (q,M)이라고 하면, query q와 가장 similar한 key에 대한 value는 V[n_1]이 됩니다. 더불어, d_i = q \cdot K[n_i] 로부터 softmax(d_1 \cdot t, ... , d_k \cdot t)의 softmaxed weights (여기서 t는 softmax function의 temperature입니다)를 구해, sorted key에 대한 value들의 weighted sum을 통해 memory의 confidence에 대한 단일 signal을 생성할 수 있죠.

Memory module (how to train):

Query vector q가 주어지면, paired label인 (supervised) value v도 주어집니다. 우리는 memory module에서 이 ground-truth value v를 갖는 key가 given query q와 nearest neighbor가 되도록 memory module을 꾸준히 update해주면 됩니다!

스크린샷 2017-03-27 오후 8.57.38

먼저, 위에서 언급한대로 given query q에 대해 k nearest neighbors (n_1, ... , n_k)를 계산하고, 이 중 V[n_p]=v가 되는 smallest index p에 대한 key n_ppositive neighbor, V[n_b] \neq v가 되는 smallest index b에 대한 key n_b를 negative neighbor라고 했을 경우 (단, top-k 안에 positive neighbor 가 없는 경우, memory module안에 value v를 갖는 아무 key vector를 골라서 positive neighbor라고 하면 됩니다.), loss는 아래와 같이 정의됩니다.

loss(q,v,M) = [q \cdot K[n_b] - q \cdot K[n_p] + \alpha]_{+}

즉, given query q에 대한 positive key와의 similarity를 maximize함과 동시에 negative key와의 similarity를 minimize하기 위한 loss 함수이고, 두 similarity사이의 distance가 margin \alpha를 넘어가면 loss를 zero로 만듦으로써 학습시 gradient가 더이상 전파되지 않도록 합니다.

Memory module (how to update):

Memory module update는 이미 V[n_1]=v인 경우 (즉, query vector q의 nearest neighbor의 value v가 이미 query의 ground-truth label인 경우)와 그렇지 않은 경우로 나눌 수 있습니다.

먼저 V[n_1]=v인 경우는, K[n_1] \leftarrow \frac{q+K[n_1]}{||q+K[n_1]||} 이렇게 current key와 query q와의 averaging followed by normalization 으로 update를 하고, 해당 age array A[n_1] \leftarrow 0를 해주어 최근에 update됐다는 표시를 age array에 해둡니다 (나중에 LRU update rule에 걸리지 않도록).

V[n_1] \neq v인 경우는, memory에 비어있는 공간이 있으면 그 곳에 update하면 되고, 그렇지 않다면 age array의 값이 가장 큰 곳 (즉, least recently used memory element)을 찾아 그곳에 update하면 됩니다. K[n'] \leftarrow q, V[n'] \leftarrow v, A[n'] \leftarrow 0 이렇게요.

그리고, LRU update를 위해 every memory update 마다 non-updated memory index들의 age 값을 1씩 증가시켜 줍니다.

Experiments:

제안하는 memory module은 1) 어떤 layer의 output을 query로 사용할지와 2) memory module의 output을 어떻게 사용할 지만 결정하면, 기존의 supervised neural network model에 자유롭게 붙여서 사용할 수 있습니다. 본 논문에서는 이 관점에서 세 가지 실험을 통해 제안하는 idea를 검증합니다.

  • CNN with memory

convolution-ReLU, max pooling, fully connected layer를 포함하는 간단한 구조의 CNN architecture에 제안하는 memory module을 추가하여 실험을 했습니다. 1) memory module query는 base CNN의 출력을 사용했고, 2) memory module에서 출력되는 nearest neighbor를 final prediction에 직접 사용했습니다.

스크린샷 2017-03-28 오전 11.44.43

one-shot learning 검증에 가장 많이 사용되는 Omniglot dataset을 사용했는데요, Omniglot dataset은 50개의 alphabet에 대해 1623 character로 구성된 dataset이며, alphabet 하나 당 평균 character수가 매우 적은 dataset이기 때문에, 본 연구의 취지에 잘 부합하는 dataset이죠. 아래 결과를 보시죠. N-way M-shot은 N개의 class (alphabet)에 대해 M개의 example (character)을 학습데이터로 사용했을 경우를 말하는데요, 보시다시피 다양한 경우에 대해 기존의 방식들보다 superior한 성능을 보여주고 있습니다.

  • seq2seq with memory

Google Neural Machine Translation (GNMT) model에 제안하는 memory module을 추가하여 실험을 했습니다. 그림을 보면 쉽게 이해할 수 있는데요,

스크린샷 2017-03-28 오전 11.44.18

GNMT는 기본적으로 attention-based encoder-decoder network이구요, 이 중 encoder part는 base network을 그대로 사용합니다. decoder part에 제안하는 memory module을 활용하는데요, 1) memory module query는 attention vector를 사용했고, 2) memory module에서 출력되는 value를 각 decoder time-step에서의 LSTM 출력과 combine하여 사용했습니다.

스크린샷 2017-03-28 오전 11.45.00

제안하는 방식으로부터, NMT에서 얻을 수 있는 성능 향상이 어느정도 인지 검증하기 위해 실험 세팅을 customize했는데요, test set을 even lines와 odd lines로 나누고, 이 중 even lines를 context set으로 사용해서 odd lines에 대한 성능을 확인했습니다. 즉, 학습을 마친 GNMT+memory model에 대해 context set으로 memory module을 update하고, odd lines로 성능을 검증한거죠. 보시다시피 좋은 성능을 보여주고 있는 것을 확인 할 수 있습니다. (가장 밑에 결과는 모든 test set을 context로 본 경우이고, 말도안되는 BLEU score를 보여주고 있죠. 이건 just for debugging 정도로 생각하고 넘어가도 좋을 것 같습니다.)

  • extended neural GPU with memory

먼저 extended neural GPU (NIPS 2016 @ Google Brain)를 이해하기 위해서는 neural GPU (ICLR 2016 @ Google Brain)를 알아야 하는데요, 이 내용은 또 하나의 포스팅이 필요한 내용이기 때문에, 여기서는 스킵하겠습니다. 다음 포스팅에서 이 두 논문에 대한 내용을 다루도록 하겠습니다 (아래는 그림만 첨부하고, 나중에 수정하도록 하죠!)Neural GPU와 Extended Neural GPU 포스팅이 추가되었습니다.

Neural GPU와 extended neural GPU 모두 Convolutional Gated Recurrent Unit을 사용하는 neural network model입니다. sequential data에 대한 학습을 time-sequence로 학습하지 않고, time-step 0 에서 모든 입력 sequence를 한 번에 받아 처리할 수 있는 parallel streaming 구조라 GPU라는 이름을 갖고있죠. extended neural GPU는 기존의 neural GPU의 NMT에서의 낮은 성능 문제를 극복하고자 제안한 encoder-decoder구조를 갖고 있습니다. 즉, 이 전 time step의 output element prediction을 다음 time step의 output element prediction에 사용할 수 있도록 tape memory p_N을 base neural GPU + encoder-decoder architecture에 추가했습니다. 자세한 내용은 해당 논문 포스팅에서 확인하시면 됩니다. 본 논문에서는 extended neural GPU가 rare event에 robust한 model이 될 수 있도록, 제안하는 memory module을 추가한 architecture에 대해 성능을 검증 합니다. Extended neural GPU는 이전 step의 output element prediction결과를 tape memory p의 해당 position에 기록해두고, 현재 step의 output element prediction을 수행할 때 이전 state s와 tape memory p를 활용할 수 있도록 되어있는 반면, 본 논문에서의 memory-augmented extended neural GPU는 tape memory p의 각 decoding stage에서의 update rule이 좀 다릅니다. 아래 그림을 보시죠.

스크린샷 2017-03-28 오전 11.44.31

o_2 prediction을 예를 들면, o_2 prediction에 사용되는 입력은 이전 decoding stage의 output state d_1과 tape memory p_1입니다. Baseline extended neural GPU에서의 tape memory p_1은, 앞서 언급했듯이 이전 prediction o_1p_0의 해당 위치에 replacement하여 얻게 됩니다. 본 논문에서는 이 구조를 바탕으로 현재 prediction을 하려고 하는 o_2에 해당하는 d_1 위치의 정보로 augmented memory query를 하고, 여기서 나오는 value를 p_1의 해당 위치에 update합니다.

스크린샷 2017-03-27 오후 10.41.30

위와 같은 문제에 적용해 볼 수 있는데요, 입력을 받아 특정 값을 출력하는 문제 입니다. A 와 B 는 그대로 출력하면 되고, 7-digit sequence는 특정 function에 의해 converting된 값을 출력하면 됩니다. 여기서 특정 function은 학습 전 random하게 정해지고, 학습동안 fix됩니다. 예를들면, 132332 (1982의 4진수 표현)을 f(1982)=3726 이 되는 random function을 통해 322032 (3726의 4진수 표현)을 출력하면 되는 거죠. 이러한 synthetic data는 무한히 많이 만들어낼 수 있는데요, 이 예제에서는 16k 개의 학습데이터를 만들어 사용했습니다. memory가 단순히 이 값들을 look-up-table (LUT) 형식으로 기억하는 방식으로 학습된다면, test set에 대한 성능이 매우 안좋을 것입니다.

스크린샷 2017-03-28 오전 11.44.53

보시다시피 baseline extended neural GPU에서는 성능이 낮은데, extended neural GPU + memory 의 경우 (제일 밑에), 성능이 매우 높아지는 것을 확인할 수 있습니다.

Discussion:

본 논문에서는 기존의 memory-augmented neural network이 갖는 단점 (rare event에 대한 generalization performance 저하)을 극복하는 새로운 구조의 memory module을 제안했고, 다양한 supervised learning model에 적용하여 그 성능이득을 확인했습니다.

나름 중요한 문제지만 이러한 내용을 검증할 수 있는 dataset이 생각보다 많이 없다는 것에 대해 저자들이 아쉬워했는데요, 아마 이상한 실험셋팅들 (synthetic data, GNMT의 even lines as a context set)로 부터 이미 느끼셨을 수도 있습니다.

본문에 언급한 neural GPU, extended neural GPU 논문도 꽤 재밌는데요, 다음 포스팅에서 이 내용을 다뤄보도록 하겠습니다. 포스팅은 여기에!

2 thoughts on “Learning to Remember Rare Events

Leave a Reply

Fill in your details below or click an icon to log in:

WordPress.com Logo

You are commenting using your WordPress.com account. Log Out / Change )

Twitter picture

You are commenting using your Twitter account. Log Out / Change )

Facebook photo

You are commenting using your Facebook account. Log Out / Change )

Google+ photo

You are commenting using your Google+ account. Log Out / Change )

Connecting to %s