[24’ ICLR] DoLa: Decoding by Contrasting Layers Improves Factuality in Large Language Models
카테고리: Language
🔍 Abstract
ICLR 2024에서 발표된 새로운 Contrastive Decoding 관련 논문이다. 해당 논문에서는 특히 Factuality를 개선하기 위해 Contrastive Decoding을 사용한다. 해당 논문의 핵심 관찰과 그 방법론은 다음과 같다.
- Factual Knowledge는 특정한 Transformer Layer에 집중되어 있는 것으로 보이며, Early Layer에 비해 Late Layer에서 더 Factual Token을 얻을 가능성이 높다. 즉, 중간쯤의 어느 Layer에서 Factual Knowledge를 가져오게 된다.
- 따라서, Early Layer와 Late Layer 간의 Contrastive Decoding을 통해 Factual Token을 증폭하여 Hallucination을 줄일 수 있다.
해당 논문은 CD(Contrastive Decoding)에 이은 논문으로, 가능하다면 먼저 해당 논문 리뷰를 읽는 것을 추천한다.
1. Key Observations
저자들은 Decoding에 따라 Token의 Distribution이 어떻게 변하는지 관찰하였다. 결국 Last Layer와 Intermediate Layer 사이의 JSD(Jensen-Shannon Divergence)를 계산하였다. 그 결과, 두 가지 패턴이 관찰되었다.
- Pattern 1: 중요한 날짜나 엔티티를 예측할 때, 즉 Factual Knowledge가 필요할 때에는 중간 Layer에서 갑자기 JSD가 증가하고, 많은 경우 JSD가 late layer까지 높다. 이는 late layer에서도 모델이 prediction을 바꾸고 있다는 뜻으로, 여기서 Factual Knowledge를 참조하고 있다는 것을 의미한다.
- Pattern 2: 반면, function word 즉
was, the, to, in등을 예측할 때나 Input Question의 token을 인용하는 경우에는 JSD가 금방 낮아진다. 이는 해당 Token이 일종의 ‘easy’ token이기 때문이다.
여기서 Pattern 1의 예시를 조금 더 자세히 살펴보면 다음과 같다.
위 문제의 정답은 Olympia이지만, Early Layer에서는 Seattle의 Probability가 높다. 이것은 Seattle이 syntactic하게 적절한 답이기 때문이다. 그러나 Late Layer에서는 Factual Knowledge를 참조하면서 Olympia의 확률이 높아지게 된다.
2. DoLa: Decoding by Contrasting Layers
2.1. Main Methods
따라서 저자들은 plausible하지만 factual하지 않은 premature layer를 선택하여 contrastive decoding을 수행하고자 한다. 그러나 이러한 layer는 매번 달라지기 때문에, 여기서는 dynamic premature layer selection을 제안한다.
먼저 Probability Distribution 사이의 거리를 JSD로 정의한다. 여기서 $N$은 last layer를 말한다.
여기서 JSD가 가장 큰 Layer를 premature layer $M$으로 선택한다. 이때 모든 Layer에 대해 JSD를 계산하는 것이 비효율적일 수 있어, Bucket $\mathcal{J}$ 내에서 선택하도록 한다. 예를 들어 LLaMA-7B의 경우 [0, 32) 전체 대신 [0, 16)이나 [16, 32)로 나누어 계산하는 식이다.
다른 방법은 Dynamic Layer Selection이 아닌 가장 Optimal한 Layer를 찾아 Layer를 고정하는 것이다. 저자들은 이를 DoLa-Static이라고 부르며, 이 경우에는 더 많은 hyperparameter search가 필요하고 dataset-specific하기 때문에 Dynamic Layer Selection보다 덜 robust하다고 한다.
2.2. Adaptive Plausibility Constraint
추가로 저자들은 Contrastive Decoding에서 사용한 APC를 사용하였다. 구체적인 내용은 CD 논문을 참고하길 바란다. 동일하게 $\alpha = 0.1$을 사용하였다.
추가로, 저자들은 repetition 문제를 교정하기 위해 simple한 repetition penalty를 적용하였다고 한다.
3. Experiments
3.1. Main Results
저자들은 Multiple-choices, open-ended generation tasks 모두에서 성능을 평가하였다. 전반적으로 성능이 향상되는 모습이다. 한편, CD의 경우 모델의 크기가 커짐에 따라 Instruction Following 능력이 올라가 오히려 “I have no comment.”의 비율이 높아져 성능이 하락하는 현상이 있었다. 반면, DoLa는 이러한 현상이 없었다.
3.2. Premature Layer Selection Strategies
DoLa-Static의 경우 동일한 Benchmark의 다른 Validation Set에서조차 Optimal Layer가 달라 hyperparameter Setting이 Robust하지 않다는 것을 알 수 있다. 반면, Dynamic Layer Selection은 이러한 문제를 해결할 수 있다.
3.3. Latency & Throughput
추가로, DoLa Decoding 방법은 거의 동일한 Latency와 Throughput을 보여주었다. 이는 DoLa가 추가적인 계산량이 적다는 것을 의미한다.
💡 Summary
DoLa(Decoding by Contrasting Layers)는 Contrastive Decoding을 LLM 내부적으로 수행하여 Factuality를 개선시키는 굉장히 간단하지만 효과적인 방법이다. 해당 논문의 내용을 요약하면 다음과 같다.
- Factual Knowledge는 중간의 특정 Layer에서 얻어지며, 따라서 Early Layer와 Late Layer 간의 Contrastive Decoding을 통해 Factual Knowledge를 증폭시킬 수 있음
- Early Layer는 Last Layer와의 JSD가 가장 먼 Layer를 선택하여 DoLa를 수행함
댓글 남기기