Pretrained Transformers as Universal Computation Engines

이번 포스팅에서는 꽤 최근에 UC Berkeley, FAIR, Google Brain에서 공동으로 발표한 Pretrained Transformers as Universal Computation Engines라는 논문에 대해서 리뷰하도록 하겠습니다.

 

논문의 내용을 한줄로 요약하자면 다음과 같습니다.

 

Language pre-trained Transformer (e.g. GPT)를 잘 fine-tuning하면 NLP가 아닌 다른 task (image classification 등)에서도 좋은 성능을 보여줄 수 있다!

 

누군가 한번쯤 해봤을 생각이지만, 단순히 생각했을 때는 잘 될 것 같지 않기 때문에 실제로 해본 사람이 없었던, 그런 느낌입니다. 이 논문은 실제로 가능하다는 것을 보여준 논문이라고 할 수 있고요. 방법을 조금 더 자세히 설명하자면 다음과 같습니다.

 

일단, 여러가지 pre-trained transformer중에서 GPT-2를 이용합니다. 그 다음, fine-tuning을 할 때 다음의 4가지 부분의 parameter에 대해서만 fine-tuning을 합니다.

 

• Output layer (Prediction layer): 당연히, language modeling으로부터 image classification으로 옮겨간다면 마지막 output layer가 바뀔 수 밖에 없습니다. Language modeling의 경우 마지막 output이 전체 vocabulary갯수가 될 것이고, 이를 image classification에 활용할 때는 전체 image의 label의 수가 되어야 합니다. 그렇기 때문에 이 부분은 아얘 새로운 layer로 갈아끼워야 합니다. 예를 들어서, pre-trained GPT-2를 CIFAR-10 classification으로 fine-tuning을 한다면 새로운 output layer는 768 * 10 = 7680개의 parameter를 가지게 됩니다. (bias가 포함된다면 10개가 더 추가됩니다.)
• Input layer (Embedding layer): 이 역시 마찬가지로, input이 문장에서 image로 바뀌면 달라질 수 밖에 없습니다. "어떻게" embedding하는지는 task마다 달라지는데, 이는 뒤에서 조금 더 자세히 설명하겠습니다. 가장 naive하게 CIFAR-10의 28*28 image를 하나의 벡터로 만든다면, flatten해서 32*32=1024차원의 벡터로 만들 수 있고 이를 d_model=768차원으로 바꾸기 위해서는 1024*768의 embedding matrix가 필요하게 됩니다. (실제로는 이렇게 하지 않습니다.)
• Layernorm parameters: Transformer에서 사용되는 layer normalization은 각 vector를 한번 normalize한 뒤 (평균을 빼고 표준편차로 나눠줍니다) affine transform(linear transform + bias)을 적용합니다. 이 affine transform에 해당하는 parameter역시 학습이 되는데, 이부분을 fine-tuning할 때 포함을 시켜줍니다. 실제 성능에서 가장 중요한 부분이 이부분을 fine-tuning하는 것이라고 하는데, 직관적으로 downstream task data의 statistic을 맞춰주는 역할을 한다고 보면 됩니다.
• Positional embedding: 실제로 학습된 positional embedding을 그대로 가져다가 써도 놀랍게도(!) 성능이 잘나왔다고 하는데, 이 역시 fine-tuning을 하면 성능이 조금 더 올라간다고 합니다. Learned positional embedding은 sequence_length * d_model만큼의 parameter를 요구하는데, 이는 전체 parameter 대비 매우 적은편이라서 fine-tuning하는 cost가 크기 않습니다.

그외의 transformer의 parameter로는 query, key, value embedding, feed-forward network등이 있는데 이부분이 transformer의 parameter 대부분을 차지합니다. 실제로, CIFAR-10에 GPT-2를 위에서 언급한 parameter만 fine-tune하게되면 전체의 0.086%에 해당하는, 매우 적은 숫자입니다.

 

그리고 나서는... 끝입니다! 위에서 언급한 부분 외에 다른 parameter는 모두 freeze하고 fine-tuning을 하면 됩니다. 모델 이름 역시 Frozen Pretrained Transformer(FPT)입니다(엄청 대충 지었네요). 논문에서 실험을 진행하는 fine-tuning tasks들은 다음과 같습니다.

 

• Bit Memory: 길이 1000짜리 bit string 5개가 입력으로 주어지고, 그 중 하나를 랜덤하게 골라서 0.5의 확률로 각 token을 masking했을 때, 원래의 token을 복구하는 task입니다.
• Bit XOR: 길이 5짜리 bitstring 2개가 입력으로 주어지고, 이 둘의 bit-wise XOR을 예측해야 합니다.
• ListOps: 수식을 계산하는 task입니다. sequence 길이가 512입니다. 매 우 어 렵 습 니 다
• MNIST: 매우 잘 알려진 숫자 이미지 분류 benchmark입니다. 각 이미지가 32 * 32의 크기인데, 이를 4 * 4의 patch 64개로 나눠서 입력으로 넣어줍니다. (input dimension 16)
• CIFAR-10: 역시 매우 잘 알려진 이미지 분류 benchmark입니다. MNIST와 똑같이 해줍니다.
• CIFAR-10 LRA: 긴 sequence를 위한 task로써 제시된 CIFAR-10의 변형 benchmark인데, 이미지를 flatten한 상태(길이 1024)를 입력으로 받아서 똑같이 이미지 분류를 합니다. Spatial inductive bias가 줄어들고 길이가 길어지기 때문에 CIFAR-10보다 훨씬 어렵습니다.
• Remote homology detection: 단백질 folding과 관련된 benchmark인데, 이에 대해서는 제가 잘 몰라서 이정도로 언급하고 넘어가겠습니다.

이제 실험 결과를 봅시다. 먼저 FPT와 downstream에 scratch부터 훈련시킨 transformer & LSTM을 비교하면 아래 표와 같습니다. 근데 이 표만 봐서는 FPT가 좋은것도 있고 나쁜것도 있어서 (심지어 MNIST와 CIFAR-10에서는 LSTM이 제일 좋습니다) 어찌된건지 알기 힘듭니다. 부연설명을 하자면, Full-trained transformer의 경우에는 모델의 크기가 커지면 (layer 12개정도) 작은 데이터셋에 대해서는 훈련 자체가 unstable하고 잘 안되었다고 합니다. 그래서 layer 3개짜리 모델을 사용하였는데, 이러한 과정에서 '적당한' 모델 크기를 찾는 것에 대한 cost가 발생하게 됩니다. 하지만, FPT의 경우 layer 36개짜리 커다란 모델 (GPT2-Large)도 훈련이 잘 되고, 모델 크기를 더욱 늘려 성능을 높이는것이 어렵지 않다고 합니다. 

FPT vs Fully-trained Transformer vs Fully-trained LSTM, Lu et al. 2021

그 다음은 pre-training task에 따른 성능입니다. FPT에서 사용한 (causal) language modeling 이외에 다음의 3가지 baseline을 제시합니다.

• Random: pre-training없이 random하게 initialize된 weight을 고정시킨 상태로 fine-tuning을 합니다. Fully-trained transformer와 다른점은 전체 parameter가 아니라 위에서 언급한 FPT의 parameter (input/output layer, layer normalization, positional embedding)만 fine-tuning을 합니다.
• Bit memory pretraining: Bit Memory task에 대해서 pre-training을 합니다.
• Image pre-training (ViT): 요즘 매우 매우 핫한 Vision Transformer(ViT)를 이용합니다. 참고로 ViT는 ImageNet-21K classification으로 pre-train되었고, autoregressive mask를 쓰지 않습니다.

아래의 결과를 보면, random의 경우 성능이 가장 안좋고, 대부분의 task에서는 FPT가 가장 좋은 성능을 보여줍니다. 다만 CIFAR-10의 경우에는 ViT가 훨씬 앞서있습니다.

 

Effect of pre-training modality, Lu et al. 2021

또한, language pre-training은 Random에 비해 속도도 빠르다고 합니다. Converge하는데 걸리는 gradient step의 수를 비교했을 때, 아래 표에서 볼 수 있듯이 많으면 40배까지 차이가 난다고 합니다.

Convergence speed, Lu et al. 2021

Overfitting과 Underfitting관점에서 봤을 때, 그냥 Transformer와 FPT, 그리고 Linformer를 비교하면 아래와 같이 train accuacy가 FPT가 가장 낮지만 test accuracy에는 큰 차이가 없음을(?) 알 수 있습니다. 이는 FPT의 모델의 크기를 키우면 좀 더 성능이 올라갈 가능성이 있는, 즉 좀 더 generalizable한 모델이라는 것을 알려줍니다. 실제로, 바로 다음의 표에서 크기를 키울수록 CIFAR-10의 accuracy가 증가함을 보여줍니다. (왜 하나는 CIFAR-10 LRA이고 다른 하나는 CIFAR-10으로 했는지 모르겠네요...)

Underfitting & Overfitting phenomena, Lu et al. 2021

Larger FPT works better, Lu et al. 2021

다음의 표는 Pre-training된 weight의 통계(mean, std)를 바탕으로 새로 initialize를 했을 때의 결과와 FPT, Random을 비교한 결과를 보여줍니다. (GPT2는 std가 0.02인 normal distribution으로 initialize합니다). 이 경우 Random보다는 좋았지만, FPT보다는 역시 별로였고, 특히 CIFAR-10에서는 Random보다 못한 결과를 보여줍니다. 이는 Language pre-training으로부터 오는 이득이 단순히 "더 좋은" initialization으로부터는 얻을 수 없음을 알려줍니다. (그래도 statistic-based보다 좀 더 좋은 방법이 있지 않을까..하는 생각이 들긴 했습니다.)

FPT vs Statistic-based initialization vs Random, Lu et al. 2021

처음에 언급했듯이 FPT는 Input/output layer, Layer normalization, positional embedding의 총 4개 부분만 fine-tuning을 하는데, 여기에 다른 부분도 추가적으로 fine-tuning했을 때 어떻게 되는지를 바로 다음 표에서 보여줍니다. Feedforward layer를 추가로 fine-tuning하면 성능이 올라가지만(CIFAR-10 LRA 제외), Attention layer까지 추가하면 오히려 성능이 안좋아지는 것을 볼 수 있습니다. Appendix(Table 11, Table 12)에서 4개를 하나씩 incremental하게 추가했을 때의 성능과 하나씩만 추가했을 때의 성능에 대한 ablation 결과를 보여주는에, 이를 보면 Layernorm을 fine-tuning하는 것이 가장 중요함을 알 수 있습니다. 실제로 layernorm은 데이터의 전체적인 통계를 control하는 부분이라고 볼 수 있기 때문에 어느정도 말이 되는 결과라고 볼 수 있습니다. 그리고 Output layer이외의 모든 부분을 고정하고 output layer만 fine-tuning할 경우(즉 transformer를 feature extractor로써만 사용하는것이고, 이 경우 input layer는 random initialize된 상태로 그대로 사용합니다) 속도는 훨씬 빨라지지만, 성능이 많이 안좋아지는 것을 보여줍니다(Table 10). (근데 input layer가 random이라면 너무 불공정한 비교가..아닌가 싶네요)

Additional fine-tuning, Lu et al. 2021

Ablation on fine-tuning parameters, Lu et al. 2021

Training only output layer, Lu et al. 2021

그 외에, Model architecture의 경우에는 Transformer 대신 LSTM을 이용하면 성능이 많이 뒤쳐지고, Bit 관련 task에서는 attention map이 interpretable하게 나타나는 것을 확인했다고 하는데, 이는 따로 설명하지 않고 마치도록 하겠습니다.

 

개인적으로는 결과에 비해서 왜 잘되는지에 대한 설명이 많지 않아 아쉬웠습니다. 이를 보완하기 위해서(?) 다양한 ablation을 진행한 것 같지만, language pre-trained transformer가 뭘 학습했길래 image에 대한 fine-tuning에서도 좋은 성능을 보여주는지에 대한 EDA를 좀 더 보여줬다면 더욱 도움이 되지 않을까 생각이 듭니다.