Image Foundation Model & Transfer methods

Intro

요즘 사회적으로 가장 핫한 주제가 ChatGPT와 같은 Foundation Model이다. 불과 몇 년 전만 해도 AI를 사용해서 실질적인 이익을 얻는 것은 아직 시간이 걸릴 것이라는 의견이 많았는데 OpenAI의 ChatGPT 열풍으로 변곡점을 맞이한 듯하다. 뿐만 아니라 최근 나오는 이미지 생성 모델들 또한 엄청난 성능을 보이며 많은 이들이 사용하고 있다. 덕분에 많은 기업들이 빅 모델(big model)을 만드는 것에 뛰어들고 있다(Fig 1). 이번 글에서는 자연어 생성 모델은 생략하고 이미지 생성(stable diffusion)과 이미지 분류(CLIP)에 관한 모델과 big 모델을 튜닝하기 위한 몇 가지 알고리즘을 소개한다.

Fig 1. Foundation models [1]

 

Zero-shot Image Classification

GPT 모델과 같은 자연어 생성 모델은 알고리즘 특성상 온라인에 있는 많은 글을 크롤링하여 학습시킬 수 있다. 그러다 보니 데이터의 양이 엄청나게 많아 학습도 잘 되고 zero-shot에서도 잘 작동하는 것이 특징이다.

Fig 2. GPT-3 training data [2]

반면 Image classification 모델 같은 경우 많은 논문에서 ImageNet 정도의 학습 데이터를 사용하는 것만으로 많은 데이터를 사용한 것으로 여긴다. 데이터의 차원이 높기 때문에 학습이 오래 걸리는 것도 한몫하지만, 보통의 classification 알고리즘은 모델의 head에서 각 class별 확률을 뽑아내서 정답과 비교하여 학습시키는 구조였기 때문에 자연어 모델과 같이 온라인에 있는 일반적인 데이터를 사용하지 못하고 labeling이 잘 되어 있는 데이터를 사용해야 했기 때문이다. ImageNet은 1400만 개 정도의 데이터밖에 없지만(Fig 3) GPT-3의 훈련 데이터는 500 billion에 육박한다. 압도적 차이다.

Fig 3. ImageNet Dataset [3]

첫 번째로 이런 한계를 극복한 OpenAI의 Learning Transferable Visual Models From Natural Language Supervision([4]) 논문을 소개한다. CLIP이라는 알고리즘으로 많이 알려져 있다. 이 알고리즘은 contrastive learning을 사용하여 온라인상의 많은 데이터를 사용할 수 있게 만들었고 zero-shot에서 매우 좋은 성능을 보인다. 논문에서 제시한 알고리즘은 Fig 4와 같다.

Fig 4. CLIP 알고리즘 [4]

먼저 온라인에 있는 사진과 사진에 대한 caption에 대한 정보를 함께 가져온다. 그리고 Image의 feature를 뽑는 모델과(Image Encoder) caption에 대한 feature를 뽑는 모델(Text Encoder)을 만든다. 각 모델에서 나온 feature는 Fig 4의 (1)의 그림처럼 각각 inner product 해준다. Contrastive learning처럼 같은 정보(image에 대응하는 caption, $I_{i} T_{i}$, 파란색 칸) 끼리는 maximize 하되 다른 정보(하얀색 칸)는 minimize 하도록 학습시키는 방식이다. 이러한 알고리즘은 이미 잘 동작하는 것이 많이 알려져 있기 때문에 좋은 아이디어인 듯하다. 

Test를 할 때는 Fig 4의 (2)와 (3)과 같이 사용하는데, 예를 들어 CIFAR10 데이터라고 하면 plane, car, dog 등 10개의 class가 존재한다. 각 class에 해당하는 caption (예를 들어 'a photo of plane')을 Text Encoder에 통과시켜 class 별 feature를 뽑고 classification 하고 싶은 Image를 Image Encoder에 통과시켜 feature를 뽑은 후 내적하여 가장 높은 값을 보이는 것이 예측 값이라고 생각하는 것이다. 실제로 이 알고리즘은 zero-shot에서 좋은 성능을 보였을 뿐 아니라 distribution shift 문제에서도 robust 한 성능을 보여주는 등의 좋은 결과로 주목받았다. 뿐만 아니라 이후에는 classification이 아닌 image generation 등의 모델에서도 쓰이며 성능을 높이는데 일조했다.

Stable Diffusion

Diffusion 모델은 original image $X_{0}$에 noise를 섞은 것을 모델을 통해 denoising 하는 과정을 배우는 알고리즘이며 VAE, GAN 등의 생성 모델보다 좋은 성능을 보여주며 많은 주목을 받았다. 하지만 초기 모델은 Fig 5와 같이 이미지를 통째로 denoising 하기 위해 많은 파라미터가 사용되어 computational cost가 큰 단점이 있었다.

Fig 5. Diffusion model architecture [6]

이를 극복한 것이 stable diffusion이다. Stable diffusion 모델의 알고리즘은 간단하다. 미리 auto encoder를 학습시키고 encoder를 사용하여 image를 압축시킨 후에 diffusion과 denoising을 하겠다는 것이 핵심이다. 이렇게 하면 computational cost가 훨씬 줄어드는 장점이 있을 뿐 아니라, Fig 6에서처럼 condition을 쉽게 latent에 추가하여 넣어줄 수 있다. (Condition으로 text를 넣어줄 때 위에서 봤던 CLIP과 같은 알고리즘 모델의 text encoder를 사용하면 더 좋은 성능을 보인다고 한다.)

Fig 6. Stable diffusion architecture [5]

Foudation Model Transfer Method

Foundation model은 보통 많은 데이터와 큰 모델을 사용하여 잘 작동한다. 하지만 파라미터가 너무 많아 down stream task에 적용하기 위해 tuning을 할 때는 문제가 된다. 모델을 통째로 tuning 하기에는 개인이나 작은 기업이 하기에 부담일 수 있고 효율적이지도 않다. 그래서 효율성을 개선한 방법들이 많이 나왔는데 몇 가지만 소개한다.

Adapter

가장 간단하게 사용할 수 있는 방법은 모델 사이에 layer를 끼워 넣는 방법이다. 이를 adapter라고 한다 [9]. 

Fig 7. Adapter [9]

보통 Fig 7과 같은 구조로 feed-forward layer 다음에 넣으며, 원래 모델은 freeze 시킨 후 adapter만 학습시키는 방식이다. 훈련할 파라미터가 훨씬 적어져서 효율적이라는 장점이 있지만 inference latency가 증가하게 되는 단점이 존재한다.

Prefix-tuning

자연어 생성 모델에서 사용되는 구조로 많은 논문에서 baseline으로 사용하고 있다. 그만큼 잘 동작하기도 하고 효율적인 알고리즘이다.

Fig 8. Prefix-tuning [7]

Prefix라는 말과 같이 모델의 input 앞에 추가 token에 해당하는 random 벡터를 넣고 이 벡터를 학습시키는 방식이다. 논문의 intuition을 보면 질문에 해당하는 맥락에 관한 정보를 같이 넣어주면 모델이 더 좋은 정보를 뱉어줄 것이라고 기대할 수 있기 때문에 이런 방법을 사용했다고 한다. 예를 들어 ChatGPT에게 '서울의 최고 명소는 어디야?' 같은 질문을 할 때, 앞에 '서울은 한국의 수도인데'라는 식으로 추가 정보를 넣어주면 더 좋은 대답을 얻을 수 있을 텐데, 이를 사람이 일일이 실험해 보긴 어려우니 학습을 통해 얻겠다는 방법이라고 볼 수 있다.

LoRA

LoRA는 Low-Rank Adaptation of Large Language Models의 줄임말이다. Low-Rank만 adaptaiton 하여 tuning 한다는 것인데 이를 이해하려면 INTRINSIC DIMENSIONALITY EXPLAINS THE EFFECTIVENESS OF LANGUAGE MODEL FINE-TUNING([10])에 나온 내용을 알아야 한다. 이 논문에서는 꽤 흥미로운 내용과 실험에 대해 소개해준다.

Auto Encoder 같은 차원 압축 알고리즘에 이미지 같은 고차원 데이터를 넣어 작은 차원의 벡터로 만들고 다시 원래 이미지로 복원해도 거의 같은 이미지가 나온다. 그 이유는 데이터가 고차원 공간에 존재하지만 대부분의 공간은 noise이고 실제 의미를 갖는 공간은 굉장히 작기 때문에, 낮은 차원에 mapping 해도 정보를 거의 잃지 않기 때문이다. 마찬가지로 모델도 굉장히 큰 차원에 존재하지만 낮은 차원에 mapping 할 수 있다고 생각할 수 있다. 그래서 원래 모델의 파라미터가 D개였다면 tuning 시에는 d(d<<D) 개의 파라미터만 사용하고 d의 파라미터를 D로 바꿔주는 random matrix ($P: R^{d} -> R^{D}$)를 사용했을 때에도 성능이 잘 나오는지 논문에서 실험을 해본다.

Fig 10. 파라미터 d를 사용하여 tuning 했을 때 결과 그래프 [10]

Fig 10은 그 결과를 보여주는데, x축이 차원 d를 나타내고 y축은 accuracy다. MRPC와 QQP 데이터 모두에서 d가 충분히 작아도 그냥 tuning을 하는 것에서 성능 저하가 크지 않음을 볼 수 있다.

이에 영감을 받아 LoRA가 만들어진다. 작은 차원의 파라미터만 사용해서 tuning 해도 충분히 좋은 finetuning 성능이 나온다는 것을 이용하여, Fig 11과 같이 matrix A와 B를 사용하여 훈련하겠다는 것이다.

Fig 11. LoRA architecture [8]

A와 B는 모두 learnable parameter이고 원래 dxd의 matrix를 가지는 모델과 다르게 r차원으로 한번 차원 축소하는 과정이 들어간다. 위의 알고리즘([10])과 다른 것은 random projection 방식을 쓰지 않는다는 것이다. 그러니 성능은 더 잘 나오게 된다. 이 알고리즘의 가장 큰 장점은 inference latency가 생기지 않는다는 것이다. 왜냐하면 아래 수식과 같이 tuning시 W(parameter)가 변하는 $\Delta W$를 $BA$로 mapping 한 것이므로 자연스레 더해서 계산하면 되기 때문이다.

$$ h = W_{0}x + \Delta Wx $$

$$ = W_{0}x + BAx $$

$$ = (W_{0}+BA)x $$

코드로 보면 Fig 12와 같이 작성되어 evaluation 시 원래 matrix와 BA를 미리 더하고 계산하는 식으로 만들 수 있다.

Fig 12. code of LoRA [11]

LoRA는 Language model에 적용하기 위한 모델이었지만, 작은 모델이라 쉽게 공유할 수 있다는 장점도 있고, 대부분의 network에서 사용가능하며, 원하는 down stream task이 변할 때마다 갈아 끼우며 사용가능하여 stable diffusion 커뮤니티에서 활발히 공유되며 사용 중이다.

Reference

[1] https://lastweekin.ai/p/gpt-3-foundation-models-and-ai-nationalism

GPT-3, Foundation Models, and AI Nationalism

[2] https://en.wikipedia.org/wiki/GPT-3#GPT-3.5

GPT-3 - Wikipedia

[3] https://paperswithcode.com/dataset/imagenet

Papers with Code - ImageNet Dataset

[4] https://arxiv.org/abs/2103.00020

Learning Transferable Visual Models From Natural Language Supervision

[5] https://arxiv.org/abs/2112.10752

High-Resolution Image Synthesis with Latent Diffusion Models

[6] https://arxiv.org/abs/2006.11239

Denoising Diffusion Probabilistic Models

[7] https://arxiv.org/abs/2101.00190

Prefix-Tuning: Optimizing Continuous Prompts for Generation

[8] https://arxiv.org/abs/2106.09685

LoRA: Low-Rank Adaptation of Large Language Models

[9] https://arxiv.org/abs/1902.00751

Parameter-Efficient Transfer Learning for NLP

[10] https://arxiv.org/abs/2012.13255

Intrinsic Dimensionality Explains the Effectiveness of Language Model Fine-Tuning

[11] https://github.com/microsoft/LoRA

GitHub - microsoft/LoRA: Code for loralib, an implementation of "LoRA: Low-Rank Adaptation of Large Language Models"