CLS 토큰이란? / 파이썬 BERT CLS 임베딩 벡터 추출 예제

[CLS] 토큰이란?

BERT, RoBERTa 등의 언어 모델에서 문장 토큰들이 인코딩된 결과를

활용하는 것이 유용한 경우가 많은데 가장 첫 위치에 문장 공통 토큰인 [CLS]를 두어

해당 위치의 임베딩 결과를 대표 임베딩으로 주로 사용하게 되었습니다.

 

BERT-base 크기의 모델의 경우 토큰 길이는 512이고 각 토큰 위치에서의 임베딩 결과는

768 차원 벡터로, 문장 input 1개에 대한 총 output 텐서의 shape는 512 * 768이 됩니다.

 

이 중 가장 앞 위치의 768 차원 벡터를 가져와 해당 문장의 대표 임베딩 결과로 활용하는 것이고,

흔히 NLP 관련 논문에서 CLS 토큰 임베딩으로 일컫는 것이 바로 이것입니다.

 

반응형

 

CLS 임베딩 벡터 추출 코드 예제

먼저, BERT-base 모델을 선언해보겠습니다.

from transformers import BertModel

model = BertModel.from_pretrained("bert-base-uncased")

 

batch size는 8로 가정해보고, 각 문장의 토큰화가 진행된 결과 벡터인

input_ids와 attention_mask는 아래와 같이 주어졌다고 가정해 보겠습니다.

import torch

# batch size = 8 가정, token length = 512
input_ids = torch.randint(30000, size = (8, 512))
attention_mask = torch.ones(8, 512)

 

우선 BERT 인코더를 통과한 output tensor는 아래 코드와 같이 last_hidden_state를

조회하여 가져올 수 있습니다.

 

여기서 output 텐서의 shape는 batch size * 512 * 768인 것을 확인할 수 있습니다.

embedding = model(input_ids = input_ids, attention_mask = attention_mask)

print(embedding.last_hidden_state.shape) # torch.Size([8, 512, 768])

 

이 중, CLS 토큰 임베딩 결과는 가장 앞의 토큰 부분의 벡터를 슬라이싱해서

가져와주시면 됩니다.

 

batch 내 각 문장에 대한 768 차원의 대표 임베딩 결과가 추출된 것이므로

최종 shape는 batch size * 768이 됩니다.

# batch 내 각 문장의 CLS 토큰 임베딩 추출
cls_embedding = embedding.last_hidden_state[:, 0, :]

print(cls_embedding.shape) # torch.Size([8, 768])