Perceiver IO: Text Classification

지난 포스팅에 이어 Perceiver IO의 적용에 대해 코드 단위로 알아보도록 하겠다. Huggingface에서 구현한 라이브러리를 중심으로 설명을 한다. Perceiver IO의 기본개념에 대해서는 아래의 이전포스팅을 확인해주길 바란다. 

 

Perceiver IO 기본개념

https://oculus.tistory.com/59?category=500795 

Perceiver IO: a scalable, fully-attentional model that works on any modality

Perceiver IO: Image Classification

https://oculus.tistory.com/61?category=500795 

Perceiver IO: Image Classification

Perceiver IO: Optical Flow

https://oculus.tistory.com/62?category=500795

Perceiver IO: Optical Flow

 

---

Perceiver for text

이번 섹션에서는 Perceiver를 text classification에 적용해보기로 한다. Perceiver의 self-attention은 input size로 time/memory가 결정되지 않기 때문에, raw UTF-8를 직접적으로 모델에 넣어줄 수 있다. 기존의 BERT나 RoBERTa 모델의 경우 WordPiece, BPE, SentencePiece와 같이 input을 tokenization을 통해 변형을 해주었는데, 이게 필요가 없어진 것이니까 굉장한 장점을 가진다. BERT(512 sequence의 subword token을 사용)와의 공정한 비교를 위해 논문에서 저자는 2048 byte의 input sequence를 이용하였다. Batch dimension까지 있다고 생각하면, input의 shape은 (batch_size, 2048)이 된다. Input에는 텍스트 시퀀스에 대한 byte ID(BERT에서의 input_id 같은거)가 들어있다. 아래의 PerceiverTokenizer를 통해 text를 byte ID 시퀀스로 바꿀 수 있고, 2048 length까지 pad 된다. 

 

from transformers import PerceiverTokenizer

tokenizer = PerceiverTokenizer.from_pretrained("deepmind/language-perceiver")

text = "hello world"

inputs = tokenizer(text, padding="max_length", return_tensors="pt").input_ids

 

또한 PerceiverTextProcessor를 통해 input 뿐만아니라 absolute position embedding을 추가하여 embedding을 만들 수 있다. Decoder로는 PerceiverClassificationDecoder를 사용한다. PerceiverClassificationDecoder는 latent의 last hidden state를 classification logit 형태로 반환해준다. 최종적으로 classification model인 PerceiverForSequenceClassification은 다음과 같이 구현될 수 있고, 추가적인 postprocessor는 필요없다. 

 

from torch import nn
from transformers import PerceiverModel
from transformers.models.perceiver.modeling_perceiver import PerceiverTextPreprocessor, PerceiverClassificationDecoder

class PerceiverForSequenceClassification(nn.Module):
    def __init__(self, config):
        super().__init__(config)

        self.perceiver = PerceiverModel(
            config,
            input_preprocessor=PerceiverTextPreprocessor(config),
            decoder=PerceiverClassificationDecoder(
                config,
                num_channels=config.d_latents,
                trainable_position_encoding_kwargs=dict(num_channels=config.d_latents, index_dims=1),
                use_query_residual=True,
            ),
        )

위의 코드를 보면 trainable position encoding이 설정되어있는 것을 볼 수 있는데, 좀 더 자세히 살펴보면서 이해해보자. Initialization 단계에서 PerceiverModel은 아래와 같은 latent variable을 정의한다. 

from torch import nn

self.latents = nn.Parameter(torch.randn(config.num_latents, config.d_latents))

Perceiver IO 논문에서는 256 latents를 사용하고, latent의 dimension을 1280으로 설정하였다. Batch dimension을 추가한다면, (batch_size, 256, 1280) 사이즈를 가지는 latents를 가지게 된다. 먼저, preprocessor는 UTF-8 byte ID를 embedding vector로 추출하는 역할을 한다. 즉, PerceiverTextPreprocessor는 (batch_size, 2048) 사이즈의 input을 (batch_size, 2048, 768) 로 바꾸어준다. 여기서 각각의 byte ID가 768 사이즈의 벡터로 mapping 된다 (PerceiverConfig를 통해 사이즈를 변경할 수 있다). 

 

그 후, Perceiver IO는 latent (batch_size, 256, 1280) 과 preprocess된 input(batch_size, 2048, 768) 을 이용하여 각각 query와 (key, value)를 추출한다. 해당 cross-attention의 output은 query의 shape (여기서는 latent)과 같게되고, (batch_size, 256, 1280)이 된다. 

 

다음으로, 연속된 self-attention layer가 적용된다. 앞서 많이 언급했듯, 더이상 input length에 구애받지 않으며, input은 방금전 cross-attention을 할 때에만 사용이 된다. Perceiver IO 논문에서는 26개의 self-attention layer (head = 8)을 가진 single block이 사용되었다. 물론 transformer와 동일하게 self-attention의 output은 (batch_size, 256, 1280)으로 동일하다. 이 output은 "last hidden state"으로도 불린다. 

 

여기까지 하면 우리는 (batch_size, 256, 1280) 사이즈의 final hidden state를 구할 수 있다. 하지만 우리가 원하는 것은 (batch_size, num_labels) shape의 classification logit이다. 어떻게 만들수 있을까?

 

이는 PerceiverClassificationDecoder를 통해 이루어지는데, 앞서 input을 latent space로 mapping 했던 것과 유사한 방식을 채택하였다. Cross-Attention을 사용하는 것이다. 이번에는 latent variable이 key와 value를 만들고, decoder query를 원하는 output size로 준비해주고 여기에서 query를 뽑아낸다. 여기서는 (batch_size, 1, num_labels) 의 decoder query를 설정해주었다고 하자. 해당 tensor는 학습 시작시에 random initialized 되고, end-to-end로 학습이 된다. 이렇게 하면 cross-attention의 output으로 같은 사이즈인 (batch_size, 1, num_labels)이 나오게 되고, 이것을 단순히 (batch_size, num_labels)로 squeeze 해서 classification logit을 만들면 완성이다!

 

다음 포스팅에서는 Perceiver를 image에 적용하는 예제를 살펴보도록 하겠다!