[논문 리뷰] Do I Know This Entity? Knowledge Awareness and Hallucinations in Language Models
https://arxiv.org/pdf/2411.14257
ICLR 2025 Orals
2025 Feb. 8
0. Bobb’s 요약
SAE로 Known entity와 Unknown entity를 찾고 이를 통해 LLM의 Output Steering하기!
Unknown Entity의 의 activation을 높이면, 아는 것도 모른다고 대답하는 걸 볼 수 있음.
1. Introduction
However, an important subset of hallucinations occurs when models are prompted to generate information they don’t possess.
Hallucination이 발생하는 중요한 원인은 모델이 가지고 있지 않은 정보를 생성할 때이다.
그래서 모델이 이미 가지고 있는 지식인 지 아닌 지를 알 수 있는 ‘linear drections in representation sapace’을 찾는다.
이 directions는 모델이 답변을 거부하는 것과 rerlevant 하다.
Sparse Autoencoders(SAEs)를 이용해서 이 directions를 찾는다
2. Sparse AutoEncoder
Sparse AutoEncoder란?
⁍: activation function
→ 점프렐루 아키텍쳐 사용
- ⁍: Weight
- ⁍: bias vector
Loss
reconstruction: L2(유클리드 거리)
sparsity: L0(벡터에서 0이 아닌 요소의 개수) → 활성화된 뉴런이 많을 수록 Loss가 올라감
Steering with SAE Latents
결국 Activation function을 증가/감소 시키면, Model을 Steering 할 수 있음
이게 무슨 말이냐?? (그림을 통해 더 자세히 살펴보자)
SAE의 latents들 중에, Known entity latent와 Unknown entity latent가 있다.
이 latent를 찾아 activation 정도를 조작(increase/decrease)하면,
LLM의 환각 정도를 조절할 수 있음.
Wilson Brown(가상의 운동선수)의 마지막 token인 Brown의 SAE의 hidden state(latent space)에서 Known latent와 Unknown latent를 찾은 뒤,
Known latent의 activation을 increase 했더니,
LLM이 원래는 모른다고 하던 걸 환각을 보임.
What is Direction?
Activation된 latent에서의 Decoder Vector 값을 Directions라고 한다
3. Methodology
3.1. Data를 모음
농구선수, 영화, 도시, 노래 로 Entity를 나누고, 위키피디아에서 정보 긁어 모음
3.2. Data의 entities를 ‘Known’ 과 ‘Unknown’으로 division
템플릿을 Gemma 2 2B와 9B 모델에 템플릿을 넣어 마지막 “Attribute” 부분을 맞추는 지 확인
- 템플릿 예시
2개 이상 맞추면: known, 0개는: Unknown, 그 사이는 entity 버림
(fuzzy string matching을 이용해 채점)
이렇게 labeling을 통해 division 하면, 모델이 모르는 데 guess하거나, 알아도 recall을 못할 수 있음.
하지만, 우리는 완벽한 classification을 한다기 보다는, 알 때와 모를 때의 명확한 “차이”를 찾으려고 하는 거기 때문에 ㄱㅊ음
하나의 Entity 당 여러 attributes를 이용해서 구분
3.3 Activation layer 찾기
모델에 known과 unknown entities를 합쳐서 넣어줌.
we use the “residual stream” of the final token of the entity, ⁍ and ⁍.
- Residual Stream? 모든 layer마다, 하나의 layer에 모든 latent마다 activations를 계산한다.
각 latent마다 아래의 계산식을 통해
⁍:
Known entity의 l번째 layer의 j번째 latent에 대해
- activation 값이 0보다 크면 → 1
- activation 값이 0이면 → 0 → 예를 들어, known에 10개의 prompts가 있고, 1번 layer 3번 latent가 7번 켜졌다면, ⁍ 이 되는 것이다.
이렇게 latent마다의 score와 아래의 “latent separation score”를 구할 수 있다.
- ⁍
- ⁍
4. SAE uncover Entity Recognition Directions
9번 레이어까지는 쭉 상승함
→ 중간 layer에 known과 unknown를 구분하는 latent가 많이 몰려있음
→ 라마 3.1 8B에서도 같은 경향을 보임
- minimum sepraation score가 높을 수록 strong generalization capabilities.
- 모든 layer, 모든 latent에 separation score를 구함
- 희소성을 유지하기 위해 activate frequently가 2%가 넘는 latent는 제외함
- highest minimum separation score를 기록한 latent는 most generalized behavior를 보이며 이후 분석의 “핵”이 될 거임
5. Entity Recognition Directions Causally Affect Knoledge Refusal
그럼 이 entity recognition은 모델에 어떤 영향을 미칠까?
Base 모델에서 뿐만 아니라
Knowledge Refusal 정량화하기
다음과 같은 질문을 넣고
Unfortunately, I don’t have access to real-time information... 와 같이 대답하는 지 확인하기
entity recognition direction을 이용해, knowledge refusal을 이끌어내기 도전!
100 queries about unkown entities에 대해 답변을 거부한 횟수.
- Random: 10개의 random latent만 Steering
orthogonalized model Unknown latent: Unknown latent 쪽의 방향을 직교화함으로써 영향력을 줄임
가중치 직교화: 가중치 행렬의 각 행을 수정하여 지정된 방향 벡터 d에 수직이 되도록 함
모델의 모든 출력 행렬에 대해 이 작업을 수행하여 어떤 컴포넌트도 해당 방향으로 쓸 수 없도록 함. 그 결과 미지의 엔티티 방향이 가장 높은 직교화된 모델은 거부 응답이 크게 감소하여 이 방향이 모델의 지식 거부 동작에서 중요한 역할을 한다는 것을 알 수 있음.
- 거부 예시
6. Mechanistic Analysis
Entity Recognition Directions Regulate Attention to Entity
- 기존 연구: entity token에 대한 attention score와 Accuracy는 비례한다
early attention heads merge the entity’s name into the last token of the entity and downstream attention heads extract relevant attributes from the entity and move them to the final token position
초기 어텐션 헤드는 엔티티의 이름을 마지막 토큰으로 병합하고, 다운 스트림 어텐션 헤드는 엔티티에서 관련 속성을 추출하여 최종 토큰 위치로 이동시킨다
- Early Attention Head: Layer 1~13
- Downstream Attention Head: Layer 15 ~ 25
(a): “entity_tok”(엔티티 토큰) 위치들에서 중간 레이어(3-15층)에서 강한 활성화(파란색)가 나타남
이는 초기 어텐션 헤드가 엔티티 토큰들의 정보를 처리하고 있음
(b): 상위 레이어 (19~25)의 특정 헤드들에서 강한 활성화가 나타남
다운스트림 어텐션 헤드가 엔티티에서 추출한 속성 정보를 최종 토큰 위치로 이동시키는거지?
→ (b)는 최종 토큰 위치에 영향을 미치는 어텐션 헤드를 찾고 있었기 때문 (불충분 ..)
결론: 아래와 같은 패턴을 발견함! (기존 연구들에서 어느정도 밝혀진 거긴 한데 더 증거를 더함)
초기 레이어에서 엔티티 토큰들의 정보를 수집하고 병합함.
마지막 레이어에서 엔티티에서 추출한 속성 정보를 최종 토큰 위치로 이동시킴
아하 ~ 그래서 ~ 마지막 토큰의 layer들을 보는 거구나!
모델이 알고 있는 엔티티와 모르는 엔티티를 처리할 때 큰 차이가 있습니다. 알려진 엔티티에 대해서는 어텐션 점수가 더 높게 나타납니다. 이는 모델이 친숙한 정보에 더 집중한다는 의미.
L18H5와 L20H3 같은 ‘Attribute Extraction Head(속성 추출 헤드)‘에서 이러한 차이가 두드러진다
확실히 latent를 Steering 하면 Original Attention Score에 비해 변화가 있는 걸 볼 수 있다.
latent로 steering 했을 때 attention Score의 평균 값 변화
상위 레이어 헤드에서 증가/감소가 많음
Early Entity Recognition Directions Regulate Expressing Knowledge Uncertainty
Entity Recognition Direction이 Knowledge Uncertainty를 표현하는 데 영향을 미침.
모델한테 이 Entity 아냐고 물어봄
(왼쪽) 알고 있는 걸 넣어도 unknown latent steering 하면 내려가고
(오른쪽) 모르는 걸 넣어도 known latent steering 하면 올라감
7. Uncertainty Directions
Uncertainty가 높을 때 역시 Unknown latent가 Activation 된다
무슨 말이냐? → 진짜 문장에서 Uncertainty가 높을 때 를 말한다
→ Unkown latent Activation 정도. 틀렸을 떄 더 활성화 되는 걸 볼 수 있음