[CoIn] 논문 리뷰 | Looking Beyond The TOP-1:Transformers Determine Top Tokens In Order (Lioubashevski et al., 2024)

https://arxiv.org/abs/2410.20210

Looking Beyond The Top-1: Transformers Determine Top Tokens In Order

 

Motivation.

 

해당 논문은 Transformer의 내부 동작에 대해 살펴보는 논문입니다. 그리고 이를 통해 알게된 사실을 통해 early exit이나, 언어 모델링을 향상시킬 수 있는 방향을 제안합니다. Attention mechanism을 기반으로하는 Transformer는 텍스트 생성이나 이미지 분류, 음성 인식 등을 포함한 다양한 모달리티에 대해서 뛰어난 성능을 보여왔고, 이는 내부 과정을 설명하려는 시도로 이어집니다. Geva et al., 2022 에 따르면 언어 영어에서 모델의 예측이 구성될 때 top-1에 대한 saturation events가 발생하여서 예측이 결정되면 이후 계층들에서도 고정된 채로 남는다는 사실이 밝혀졌습니다. 그래서 해당 논문에서는 top-1 예측에 대한 saturation event가 발생한 후에는 transformer에서 어떤 연산이 일어나는지에 대해 살펴봅니다.

 

이를 밝혀내기 위해 transformer의 각 hidden layer에 decoder의 vocab space로 사영하는 layer(hidden feature를 자연어 토큰으로 변환하는 layer)를 달아서, top-1을 넘어서 top-2, 3, 4 등의 상위 토큰들에 대해서도 saturation event가 발생하는지를 보입니다. 그리고 top-2, 3, ... 등에 대해서도 순서대로 saturation event가 발생한다는 사실을 발견했습니다.

Figure 1: An illustration of the proposed task-transition mechanism wherein the layers of the Transformer perform a changing number of tasks in order, so that task i is determining the i-th token in the final ranking, and the transition between task i and task i + 1 occurs at the corresponding i-th saturation layer. The transition is akin to a switch being flipped “on” and staying “on” for the remaining layers representing the i-th token being fixed from this point onward.

 

그림 1. 은 saturation event가 top-1, 2, 3, ... 의 순서대로 일어남을 나타내는 그림입니다. 위 그림에 따르면 "toy"라는 토큰이 31-th layer에서 saturation된 후 일명 "switch가 꺼지"면서 더 이상 토큰의 순위가 바뀌지 않습니다. 그리고 top-2인 "ball"이라는 토큰은 38-th layer에서 saturation되고, 다음 토큰, ... 이런 현상이 발생합니다.

 

Figure 2: Schematic of our framework and visualization of the ordered saturation of the top-k token on GPT2-XL. The hidden states from each layer are projected onto the vocabulary space using the unembedding matrix E, then sorted in descending order and treated as rankings. The saturation effect is marked separately for each token in the top-4 of the final ranking, emphasizing the fact that the 2nd token saturates after the 1st token, the 3d token saturates after the 2nd token and so on. The dashed line represents the previously established saturation event of the top-1 token.

 

그림 2. 는 실험의 과정을 그림으로 나타낸 것입니다. GPT2-XL과 같은 decoder-only 아키텍처에 대해서 각 layer마다 vocab space로 사영해주는 Embedding layer를 달아서 각 layer마다의 token 순위를 출력합니다. 그리고 이를 통해 top-1, 2, ... 토큰들이 언제 saturation이 되는지를 알 수 있습니다.

 

그리고 이를 언어 모델 뿐만 아니라 vision이나 audio에 대해서도 검증하며, GPT와 같은 decoder-only 구조 뿐만 아니라 ViT와 같은 encoder-only, Whisper와 같은 encoder-decoder 구조에 대해서도 이 사실을 검증합니다. 그리고 이 사실이 학습과는 별개로 transformer 구조 자체에 내재된 특성임을 추가적인 실험을 통해 증명합니다.

 

Experiments.

 

실험에 앞서 해당 논문의 저자는 saturation layer에 대해 정의를 먼저 합니다:

 

(1) 1st Saturation layer(;Geva et al., 2022): 만약 모델의 top-1 예측이 계층 l-th layer 이후의 모든 후속 계층들에서 일정하게 유지된다면, 입력의 index i에 대해 saturation layer는 l-th layer에서 발생합니다. 다시 말해, N개의 layer를 갖는 모델에서, 모든 계층 l'에 대해서 l < l' ≤ N 을 만족할 때, l' 에서 토큰의 top-1 예측이 변하지 않는 경우 saturation event는 l-th layer에서 발생한다고 얘기합니다.

 

(2) k-th Saturation layer: 1st Saturation layer의 정의를 확장합니다. 예를 들어, 그림 2. 에서 1st saturation layer는 44로 최상위 토큰("toy")이 그곳에서 결정되기 때문입니다. 이와 유사하게 두 번째 토큰("ball")이 45-th layer에서 결정되기 때문에 2nd Saturation layer는 45입니다. 

 

이런 정의 하에, 해당 논문의 저자들은 top-k의 k들에 대한 포화 계층을 계싼하고, 각 k에 대한 k-th saturation layer의 순위(rank)를 계산합니다. 또한 해당 논문의 저자는 상위 k개 토큰이 순서대로 saturation되는 밑바탕에 작업 전이(task transition)가 있다고 가정합니다. 예를 들어, 모델은 top-1을 예측하는 task를 수행한 후, top-2를 예측하는 task로 전이(transition)가 발생하는 것입니다.

 

이를 증명하기 위해 각 layer에 simple classifier를 달아서 해당 계층이 "몇"번째 saturation layer에 해당하는지를 예측하도록 한 후, 그 결과를 임베딩의 평균과 분산을 갖는 정규분포(임베딩의 통계적인 정보는 유지하는 방향으로의 랜덤 임베딩)로부터 똑같은 예측을 수행한 결과를 비교합니다. 이를 통해 모델의 임베딩 자체가 task에 대한 정보를 담고 있다고 증명합니다.

 

Results.

 

 

Figure 3: Average rank of the k-th saturation layer among the saturation layers for k=1,..,5 with standard error bars. Asterisks indicate statistically significant differences between consecutive token ranks (*** represents p < 0.001), based on an independent samples t-test.

 

그림 3. 은 Text generation(GPT2-XL, decoder-only)와 Vision(ViT-L/16, encoder-only), audio(Whisper-large, encoder-decoder)에 대해 top-k token과 k-th saturation layer에 대한 관계를 설명하는 그림입니다. 이에 따르면 top-k의 rank(순위)와 saturation layer의 순위가 단조적으로 증가하는 모습을 보여줍니다. 이를 통해 top-k 번째 token이 saturation되는 순서가 고정되어있음을, 순서대로 task가 수행됨을 보여줍니다.

 

audio(c)에 대해서는 그 관계가 뚜렷해보이지는 않는데, 해당 논문의 저자들은 encoder-decoder 구조로 인해 decoder 수준에서 encoder의 정보가 계속되서 사용되기 때문에 transformer mechanism의 구조적인 특징이 희석되었다고 분석했습니다. 즉, decoder는 encoder의 정보를 이용해서 높은 순위(top-1, 2)에 대한 예측을 강화하는 방향으로 학습이 되기 때문에 상위 순위에 대해서는 안정적인 반면, 조금 낮은 순위(top-3, 4, 5)에 대한 예측은 decoder 레이어가 진행되면서 외부(encoder) 정보가 지속 주입되어 작업 경계가 흐려지는 효과가 발생한다는 뜻입니다.

 

Table 1: Accuracy of task number logistic regression classifier showing that in all modalities the layer embeddings contain information about the task number.

 

표 1. 은 각 layer의 임베딩을 이용해서 해당 layer가 몇 번째 saturation layer인지를 예측하는 task의 결과입니다. 이때 random embedding은 평균과 표준편차를 유지하며 가우시안 분포를 따라 랜덤하게 임베딩을 변경한 뒤 예측을 수행하는 결과로, layer embedding에 비해 성능이 현저히 떨어짐을 알 수 있습니다. 이를 통해 모델의 임베딩 정보가 saturation에 대한 정보를 담고 있다고 생각할 수 있습니다.

 

Analysis.

 

해당 논문의 저자들은 transformer의 이런 특징을 이용하여 모델이 다음 작업(다음 top-k 예측을 수행하도록 하는 task)으로 전환할 수 있도록 모델에 개입(intervention) 절차를 사용할 수 있다고 얘기합니다. 저자들은 각 saturation layer가 다음 task로 전이하라는 신호를 가지고 있고, 모든 후속 계층이 이전 작업이 완료되었고 해당 토큰이 고정되었다는 정보를 담고 있다고 가정합니다. 이를 증명하기 위해, 샘플 s1의 1st saturation layer의 출력을 샘플 s2에 주입하여, s2의 task가 전환되었는지를 확인합니다.

 

만약 s1의 1st saturation layer가 10층이고, s2의 1st saturation layer가 15층이라면, s1의 1st saturation layer의 출력을 s2에 대해서 같은 layer(10층)에 주입하여 s2에 대해서 1st saturation layer가 기존(15층) 대비 변했는지를 확인합니다.

 

Figure 5: Flipping the Top-1 Switch. The percentage of examples where the top-1 saturation occurred at the injected layer after the intervention, shown as a function of the layer from which the injected activations were taken, relative to the original saturation layer (e.g., −2 means activations were taken from two layers before the original saturation layer).

 

그림 5. 는 각 모달리티(text와 vision, audio)들에 대해서 injection을 수행한 후 task가 전환되었는지를 확인한 것입니다. 이에 따르면 주입된 출력이 원래 saturation layer보다 이른 계층에 주입되면 task switch가 발생하지 않지만 원래 layer이후에 주입되면 task switch가 발생합니다.

 

Practical Applications.

 

이런 사실은 두 가지 실용적인 응용으로 이어질 수 있습니다:

 

(1) early-exit 전략: 해당 논문의 저자는 task-trainsition classifier에 기반한 새로운 토큰 단위 동적 추론 방법을 제안하며, 여기서 각 토큰의 조기 종료 계층은 classifier가 task 2에 속한다고 예측하는 가장 이른 계층으로 정의됩니다. 즉, top-1 prediction만 알면 되는 다음 토큰 예측 task의 경우는 top-2를 예측하는 task 2로의 전환이 일어나면 추론을 종료하여(어짜피 이후에 top-1 token은 변하지 않으니까) 효율적인 추론을 달성합니다.

 

실효성을 입증하기 위해, Schuster et al., 2022 이 소개한 두 가지의 local confidence 측정치와 비교합니다. 우선 softmax 방법입니다. 이는 각 layer의 출력에 softmax를 취해 인접하는 두 layer들의 출력 차이가 갑자기 커지면 task가 변경되었다고 판단하여 exit을 하는 방법입니다. 또 다른 방법은 hidden-state saturation이라는 방법입니다. 이는 각 layer의 출력에 cosine similarity를 취해 그 차이가 갑자기 커지면 task가 변경되었다고 판단하여 exit을 하는 방법입니다.

 

Figure 6: Performance-efficiency trade-off comparison of different confidence measures against a static baseline (where all layers are used for each token) and a local oracle measure (where the early exit is at the top-1 saturation layer). The graph shows softmax and state confidence measure results at different thresholds. Our method achieves the highest next-word prediction accuracy out of all early-exist methods while providing significant speedup compared to the baseline

 

그림 6. 은 성능과 속도 증가에 대해서 각 방법론들을 비교한 결과입니다. 제안한 방법론이 Schuster의 두 방법론에 비해 적절한 speed up으로 좋은 성능을 보임을 알 수 있습니다.

 

(2) 언어 모델링: 만약 top-1 token을 예측하는 task를 수행하는 layer가 너무 많고, top-2 token을 예측하는 task를 수행하는 layer의 개수가 적다면 top-2 token의 신뢰도가 떨어진다고 예측할 수 있습니다. 즉, saturation layer가 어디에 위치해있는지를 통해 top-k token의 신뢰도를 측정할 수 있습니다. 

 

이를 검증하기 위해 논문의 저자는 top-2 토큰의 saturation layer가 출력 layer보다 최소 7개 계층 이전에 위치한 샘플과 top-2 토큰이 saturation되지 못하고 마지막 계층을 통과하는 샘플을 비교합니다. 첫 번째 샘플의 경우 top-2 token의 정확도가 31.99%인 반면, 두 번째 샘플의 경우는 17.14%로 saturation layer를 통한 토큰의 신뢰도 예측이 합리적임을 증명합니다(top-1이 오답인 경우, top-2가 얼마나 믿을 만 한지를 보려는 목적입니다).