Last_hidden_state와 Logit

본 글은 [Open-Up] 오픈소스 소프트웨어 통합지원센터로부터 지원받아 작성하였습니다.

 

일반적으로, Last_hidden_state는 logit을 생성하기 위한 입력으로 사용된다.
일반적인 흐름: Input → ... → Last Hidden State → Linear Layer → Logits → Softmax → Probabilities

 

[Last_hidden_state]:
- 정의: 일반적으로 모델의 마지막 트랜스포머 층(layer)의 출력을 말한다.
- 특징: 
  - 고차원의 벡터 표현 (예: BERT-base에서는 768차원).
  - 각 입력 토큰에 대해 하나의 벡터가 있다.
  - 입력 시퀀스의 의미적, 문맥적 정보를 포함하고 있다.
- 용도: 
  - 다양한 downstream 작업을 위한 특성(feature) 추출에 사용된다.
  - 문장 임베딩, 토큰 분류 등의 작업에 직접 사용될 수 있다.


[Logit]:
- 정의: 최종 출력 층(일반적으로 선형 층)의 출력으로, softmax 함수를 적용하기 전의 raw 점수.
- 특징: 
  - 모델의 어휘 크기와 같은 차원을 갖는다. (예: GPT-2의 경우 50,257차원).
  - 각 가능한 다음 토큰에 대한 점수를 나타낸다. (CLM을 기준으로 설명)
  - 정규화되지 않은 확률 분포.
- 용도: 
  - 다음 토큰 예측을 위한 직접적인 입력으로 사용됨.
  - Softmax 함수를 통해 확률 분포로 변환됨.



[주요 차이점]:
- 차원: Last_hidden_state는 모델의 hidden size 차원을, logit은 어휘 크기 차원을 갖는다.
- 의미: Last_hidden_state는 입력의 의미적 표현을, logit은 다음 토큰에 대한 예측 점수를 나타냄.(CLM / next token prediction 기준)
- 활용: Last_hidden_state는 다양한 작업에 범용적으로 사용될 수 있지만, logit은 주로 다음 토큰 예측에 특화되어 있다.


[임베딩으로서의 사용]:
Last_hidden_state는 그 자체로 임베딩으로 사용될 수 있다.
Logit은 일반적으로 임베딩으로 직접 사용되지 않지만, 특정 상황에서는 logit의 패턴이 유용한 정보를 제공할 수 있다.

 

---

 

LLM을 기준으로 예시를 들어 설명해보자. 

embedding_size(representation 차원)가 4096인 LLM(ex. mistral7B, LLama2 8B 등)에 10개의 토큰을 입력하는 상황. 
이때, batch_size = 1로 가정.

[last_hidden_state의 크기]:
- 크기: (1, 10, 4096)
- 설명:
  - 1: 배치 크기 (일반적으로 배치 처리를 위해 첫 번째 차원이 있습니다)
  - 10: 입력 시퀀스의 길이 (토큰 수)
  - 4096: hidden state의 차원

[logits의 크기]:
- 크기: (1, vocab_size)
- 설명:
  - 1: 배치 크기
  - vocab_size: 어휘 크기 (가능한 모든 토큰에 대한 점수)
  - 참고로, last_linear_layer는 (4096, vocab_size)의 가중치를 갖고 있음.

 

---

 

이제 모델을 특정해보자,

Mistral 7B 모델을 기준으로 Last_hidden_state와 Logit의 차이를 확인해보자. 

from transformers import AutoModelForCausalLM
 
model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-Instruct-v0.1")
model.to("cuda")

AutoCasulLM으로 model을 호출하면, 아래 빨강색 부분까지 추가되어 모델을 불러온다.

아래 빨강색 Layer가 바로, logit을 생성하는 last_linear_layer이다.
차원을 보면, (4096, vocab_size = 32000)인걸 알수 있다.

그렇다면, Last_hidden_state가 출력되는 부분은 어디일까?
바로 파랑색 박스 부분이다. 

Last_hidden_state를 마지막 트랜스포머 층(layer)의 출력이라고 말했는데,
MixtralDecoderLayer를 모두 통과한후, 정규화를 위한 MistralRMSNorm()을 거친다. 
여기를 거치고 나면, 정규화된 Last_hidden_state 또는 last representation을 얻을수 있다. (물론, MistralRMSNorm() 이전에도 뽑아서 쓸수 있다)

 

 

 

참고로, huggingface의 transformer는 logit과 last_hidden_state를 편하게 뽑아 쓸수 있도록 지원한다.
출처 : https://huggingface.co/docs/transformers/main_classes/output

# 예시)
outputs = model(**inputs, labels=labels)

outputs .logit()
outputs .last_hidden_state()

---

만일, 모델의 header를 제거해서 부르고 싶다면?

 

from transformers import AutoTokenizer, AutoModel, AutoConfig

# mitral 7b instruct LLM 모델
model = AutoModel.from_pretrained(
    "mistralai/Mistral-7B-Instruct-v0.1",device_map="cuda"
)

### e5-7b-instruct 임베딩 모델
# model = AutoModel.from_pretrained( 
#     "intfloat/e5-mistral-7b-instruct",device_map="cuda"
# )

model

 

이런 방식으로 Last_hidden_state를 최종 출력으로 하는 모델로 변경할수 있다.

(LLM을 임베딩 모델로 바꾸고 싶을때라던지....에 활용될 수 있다. LLM2Vec 논문 리뷰)