🤖 AI

Inference Time Causal Probing in LLMs

본 논문은 기존 분류기 의존적 접근법보다 더 신뢰할 수 있는 인과적 개입을 달성하기 위해 모델의 고유 출력을 활용하여 LLM 은닉 상태를 직접 조정하는 프로브 없는 경계 기반 방법인 숨겨진 상태 기반 마진 개입 (HDMI) 을 제안하며, 동시에 텍스트 편집을 위한 룩어바 변형을 도입한다.

원저자: Sadegh Khorasani, Saber Salehkaleybar, Negar Kiyavash, Matthias Grossglauser

게시일 2026-05-11

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Sadegh Khorasani, Saber Salehkaleybar, Negar Kiyavash, Matthias Grossglauser

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

이 논문은 쉬운 언어와 일상적인 비유를 사용하여 설명합니다.

큰 그림: "블랙박스" 수정하기

대형 언어 모델 (LLM) 을 주방에 있는 초지능이지만 약간 불투명한 셰프로 상상해 보세요. 이 셰프는 놀라운 이야기를 쓸 수 있지만, "왜 이 수프에 소금을 넣었나요?"라고 물으면 명확한 답을 하지 못할 수도 있습니다. 그들은 그저 "맛이 좋다는 것을 안다"고만 할 뿐입니다.

과학자들은 이 셰프가 어떻게 결정을 내리는지 이해하고자 합니다. 구체적으로, 그들은 다음과 같은 질문을 하고 싶어 합니다: 셰프가 실제로 동사 'run'을 결정할 때 '복수형'(예: 'cats') 개념을 사용하는 것일까, 아니면 그것이 단순한 우연일까?

이를 테스트하기 위해 연구자들은 **인과성 탐지 (Causal Probing)**라는 기법을 사용합니다. 이는 "만약에" 실험과 같습니다. 그들은 "좋아, 주어를 단수형 대신 복수형이라고 가정해 보자. 셰프가 동사를 'runs'에서 'run'으로 바꾸는가?"라고 말하고 싶어 합니다.

구식 방법의 문제점: "번역가" 문제

과거에 이 실험을 수행하려면 연구자들은 번역가(프로브라고 함) 를 고용해야 했습니다.

그들은 이 번역가를 훈련시켜 셰프의 비밀 노트 (은닉 상태) 를 읽고 주어가 단수형인지 복수형인지 추측하게 했습니다.
그런 다음 번역가의 피드백을 사용하여 셰프의 노트를 살짝 밀어 의미를 변경했습니다.

결함: 이 번역가는 셰프의 언어를 완벽하게 구사하지 못할 수 있습니다. 번역가는 '복수형'이 어떻게 생겼는지에 대한 자신만의 규칙을 가지고 있을 수 있으며, 이는 셰프가 실제로 생각하는 방식과 일치하지 않을 수 있습니다. 이는 다른 브랜드의 자동차를 위한 매뉴얼을 사용하여 자동차 엔진을 수리하려는 것과 같습니다. 올바른 버튼을 누를 수도 있지만, 엔진의 실제 배치를 이해하지 못해 다른 부분을 고장 낼 수 있습니다.

새로운 해결책: HDMI(직접적인 "밀기")

저자들은 HDMI(Hidden-state Driven Margin Intervention, 은닉 상태 기반 마진 개입) 라는 새로운 방법을 제안합니다.

비유: 번역가를 고용하는 대신, HDMI 는 셰프의 머릿속에 직접 이야기합니다.

목표: 셰프가 "runs"(단수형) 라고 말하려 합니다. 당신은 그들이 "run"(복수형) 이라고 말하기를 원합니다.
구식 방법: 번역가에게 "복수형" 버튼을 찾아 누르라고 요청합니다.
HDMI 방법: HDMI 는 셰프의 최종 결정 과정 (출력) 을 살펴봅니다. 그것은 "run"이라는 단어가 조금 더 확률이 높고 "runs"는 조금 더 확률이 낮아지도록 필요한 정확한 수학적 "밀기"를 계산합니다. 이는 두 단어 사이의 차이인 **마진 (margin)**을 살펴봄으로써 이를 수행합니다.

더 나은 이유:

번역가 불필요: 개념을 이해하기 위해 별도의 AI 를 훈련시킬 필요가 없습니다. 모델의 자체 두뇌를 사용하여 출력을 변경하는 방법을 파악합니다.
정밀성: 모델의 자체 "기하학"(단어를 어떻게 자연스럽게 배열하는지) 을 사용하기 때문에, 모델이 실제로 생각하는 방식과 완벽하게 일치합니다.
효율성: 긴 우회로가 아닌 핸들을 빠르게 튕기는 것과 같은 간단하고 빠른 계산입니다.

"선제적" 업그레이드: LA-HDMI

이 논문은 텍스트 편집을 위한 LA-HDMI(Lookahead HDMI, 선제적 HDMI) 라는 세련된 버전도 소개합니다.

상황: 당신이 이야기를 쓰고 있다고 상상해 보세요: "The cat was sleeping."(고양이가 잠을 자고 있었다.) 당신은 이를 "The cats were sleeping."(고양이들이 잠을 자고 있었다.)로 바꾸고 싶습니다.

문제: 단순히 "was"를 "were"로 변경하면 문장이 깨질 수 있습니다. 왜냐하면 그 앞의 단어 ("The") 나 그 뒤의 단어도 문법적 흐름을 유지하기 위해 변경될 필요가 있을 수 있기 때문입니다.
LA-HDMI 해결책: 이 방법은 현재 단어만 보는 것이 아니라 앞으로 내다봅니다. 변경을 가하는 동안 문장의 다음 몇 단계를 시뮬레이션합니다.
- 그것은 "was"를 "were"로 변경하면 "The"를 "The"로 변경할 필요는 없지만 (변경 없음), 명사가 변경된다면 명사 앞의 관사를 조정해야 할 수도 있음을 파악합니다.
- 단어조차 쓰이기 전에 모델의 숨겨진 생각들을 부드럽게 조종하여, 변경된 단어뿐만 아니라 전체 문장이 매끄럽고 유창하게 유지되도록 합니다.

이는 지금 당장 왼쪽으로 돌아라고 알려주는 GPS 가 아니라, 나중에 막다른 길에 갇히지 않도록 전체 경로를 계산하는 GPS 와 같습니다.

효과가 있었는가? (결과)

저자들은 문법과 논리를 테스트하도록 설계된 두 가지 주요 "체육관"(데이터셋) 에서 이를 테스트했습니다.

LGD Agreement: 주어와 동사의 일치 여부 확인 (예: "he is" 대 "they are").
CausalGym: 복잡한 문법 퍼즐 세트.

점수판:
그들은 두 가지를 측정했습니다.

완전성: 우리는 성공적으로 의미를 변경했는가? (예: "runs"가 "run"이 되었는가?)
선택성: 우리는 실수로 다른 것들을 망가뜨렸는가? (예: 실수로 시제나 전체 문장의 의미를 변경했는가?)

판결:
HDMI 는 일관되게 구식 방법보다 높은 점수를 받았습니다. 문장의 나머지 부분을 망가뜨리지 않고 정확히 변경해야 할 것을 변경하는 데 더 뛰어났습니다. 이는 Llama 와 Pythia 와 같은 다양한 유형의 AI 모델에서 잘 작동하여 강력한 도구임을 입증했습니다.

요약

구식 방법: AI 의 마음을 어떻게 바꿀지 추측하기 위해 별도의 도구를 사용함 (종종 부정확함).
HDMI: AI 의 마음을 바꾸기 위한 완벽한 밀기를 계산하기 위해 AI 의 자체 출력을 사용함 (정확하고 효율적).
LA-HDMI: 편집된 단어뿐만 아니라 전체 문장이 자연스럽게 흐르도록 보장하기 위해 "수정구"를 사용하여 텍스트를 매끄럽게 편집하는 HDMI.

이 논문은 이 "직접적인 밀기" 접근 방식이 AI 모델이 어떻게 생각하고 쓰는지 이해하고 통제하는 더 신뢰할 수 있는 방법이라고 결론지었습니다.

기술 요약: LLM 의 추론 시간 인과적 프로빙 (HDMI)

1. 문제 제기

본 논문은 대규모 언어 모델 (LLM) 이 예측을 위해 내부 표현을 어떻게 활용하는지 이해하는 데 사용되는 기존 인과적 프로빙 (causal probing) 방법의 한계를 다룹니다.

상관관계 프로빙의 한계: 전통적인 방법은 숨겨진 상태 (hidden states) 에 보조 분류기 ("프로브") 를 훈련시켜 언어적 속성 (예: 주어 단수/복수) 을 감지합니다. 이러한 방법들은 무엇이 인코딩되었는지를 보여주지만, 모델이 다음 토큰을 생성할 때 해당 정보를 어떻게 사용하는지는 증명하지 못합니다.
인과적 프로빙의 한계: 기존 인과적 접근법은 속성을 변경하기 위해 (예: 단수를 복수로 변경) 숨겨진 상태에 개입하고 출력의 변화를 관찰하려 시도합니다. 그러나 이러한 방법들은 일반적으로 개입 방향을 결정하기 위해 별도의 프로브를 훈련하는 데 의존합니다 (예: 프로브의 기울기를 사용하여 라벨을 반전).
- 단점: 이는 추가적인 훈련 비용을 발생시키고, 속성별 지도 학습을 요구하며, 불일치 (misalignment) 의 위험을 초래합니다. 프로브의 분류 경계가 모델의 내부 "예측 기하학 (predictive geometry)"과 일치하지 않을 경우, 모델의 생성을 효과적으로 유도하지 못하거나 관련 없는 속성을 부수적으로 변경하는 개입이 발생할 수 있습니다.

저자들은 모델을 재훈련하거나 보조 프로브를 훈련하지 않고 추론 시간에만 개입이 적용되는 추론 시간 인과적 프로빙 (Inference-Time Causal Probing) 프레임워크를 제안합니다.

2. 방법론: 숨겨진 상태 기반 마진 개입 (HDMI)

핵심 기여는 생성 모델의 네이티브 출력 헤드를 직접 활용하여 숨겨진 상태를 유도하는 프로브 없는 기울기 기반 기법인 HDMI입니다.

2.1 핵심 메커니즘

HDMI 는 모델의 원래 선호하는 이어짐인 "소스 (source)" 토큰과 원하는 반사실적 이어짐인 "타겟 (target)" 토큰 사이에 마진 목적 함수 (margin objective) 를 정의하여 작동합니다.

목적 함수: 로그 오즈 (logit) 마진 $L(x) = \phi(x)_\tau - \phi(x)_\sigma$ 를 최대화합니다. 여기서 $\phi(x)$ 는 로그 오즈 벡터, $\tau$ 는 타겟 토큰 인덱스, $\sigma$ 는 소스 토큰 인덱스입니다.
개입: 이 방법은 특정 층 $\ell$ 의 숨겨진 상태 $h_\ell(x)$ 에 대해 기울기 상승 (gradient ascent) 을 수행하여 타겟 토큰의 확률을 높이고 소스 토큰의 확률을 낮춥니다.
계산 효율성: 최종 층의 경우, 기울기는 폐쇄형 행렬 - 벡터 곱 $\nabla_{h_L} L(x) = W_U^\top (e_\tau - e_\sigma)$ 로 표현됩니다. 이는 반복적인 프로브 훈련의 필요성을 제거합니다.
다중 토큰 확장: 이 방법은 타겟 집합 ( $T^+$ ) 과 소스 집합 ( $T^-$ ) 에 대한 로그 오즈를 합산하여 $L_{set}(x) = \sum \phi(x)_{\tau_i} - \sum \phi(x)_{\sigma_j}$ 로 정의함으로써 토큰 집합 (예: 여러 개의 유효한 복수 동사) 을 지원합니다.

2.2 텍스트 편집을 위한 Lookahead HDMI (LA-HDMI)

사용자가 수정된 시퀀스를 제공하는 텍스트 편집 문제를 해결하기 위해 저자들은 LA-HDMI를 도입합니다.

과제: 표준 자기회귀 (autoregressive) 생성은 $t$ 단계의 숨겨진 상태가 $t' > t$ 인 미래의 편집을 인식할 수 없도록 합니다.
해결책: LA-HDMI 는 기대 토큰 임베딩을 통해 기울기를 역전파하는 "lookahead" 메커니즘을 도입합니다.
- 이는 유창성을 유지하기 위해 낮은 온도 $\beta_f$ 를 사용하는 순방향 통과 (forward pass) 와 기대 임베딩을 계산하기 위해 높은 온도 $\beta_g$ 를 사용하는 기울기 경로 (gradient path) 를 분리합니다.
- 목적 함수 $J_t(h_t)$ 는 미래 편집 위치 ( $S_{max}$ 단계 ahead) 에 대한 마진 목적 함수를 누적하여, 현재 숨겨진 상태 $h_t$ 를 조정함으로써 사용자의 편집 방향으로 미래 생성을 유도하면서도 편집되지 않은 영역의 유창성을 유지합니다.

3. 평가 프레임워크

본 논문은 Canby 등 [2024] 의 프레임워크에 따라 두 가지 주요 지표를 사용하여 개입을 평가합니다:

완전성 (Completeness): 숨겨진 상태에서 타겟 속성 (예: 주어 단수/복수 변경) 이 성공적으로 전환되었는지를 측정하며, 별도의 검증 프로브로 확인합니다.
선택성 (Selectivity): 개입 후 관련 없는 속성 (방해 변수, 예: 비주어 명사의 단수/복수) 이 안정적인지 측정합니다.
신뢰도 (Reliability): 완전성과 선택성의 조화 평균으로, 개입 품질의 전체적인 척도로 작용합니다.

4. 주요 결과

저자들은 LGD agreement 코퍼스와 CausalGym 벤치마크를 사용하여 Meta-Llama-3-8B-Instruct와 Pythia-70M에서 HDMI 를 평가했습니다.

성능: HDMI 는 대부분의 작업에서 프로브 기반 베이스라인 (AlterRep, FGSM, PGD) 보다 일관되게 높은 신뢰도를 달성합니다.
- 일치 suites (예: 주어 - 동사 일치) 에서 HDMI 는 종종 완벽한 완전성과 높은 선택성을 달성합니다.
- 베이스라인과 달리 HDMI 의 성능은 개입용 프로브의 품질에 독립적입니다. HDMI 는 프로브를 사용하지 않기 때문입니다. AlterRep 및 PGD 와 같은 베이스라인은 프로브 정확도가 낮을 때 더 작은 모델 (Pythia-70M) 에서 성능이 크게 저하되는 것을 보였습니다.
강건성: HDMI 는 다양한 모델 크기와 작업에서 일관된 신뢰도를 유지하는 반면, 프로브 기반 방법은 특정 속성과 모델에 민감한 혼합된 결과를 보였습니다.
텍스트 편집: LA-HDMI 는 문법적 유창성을 유지하고 필요한 문맥을 조정 (예: "a"를 "an"으로 변경) 하면서 사용자가 지정한 편집 (예: "girl"을 "owl"로, "was"를 "were"로 변경) 방향으로 생성을 유도하는 능력을 성공적으로 입증했습니다.

5. 의의 및 주장

본 논문은 HDMI 가 프로브 의존적인 접근에서 모델 네이티브 인과적 프로빙으로의 전환을 나타낸다고 주장합니다.

정렬 (Alignment): 모델의 자체 출력 헤드를 사용하여 개입 방향을 정의함으로써, HDMI 는 교란이 모델의 실제 예측 기하학과 정렬되도록 보장하여 외부 프로브의 불일치 위험을 피합니다.
효율성: 이 방법은 보조 분류기 훈련의 오버헤드 없이 숨겨진 상태에 대한 기울기 상승만 요구하는 계산적으로 경량화된 방법입니다.
일반성: 이 접근법은 자연스럽게 다중 토큰 속성으로 확장되며, lookahead 기울기를 통한 유창한 텍스트 편집 메커니즘을 제공합니다.

저자들은 하이퍼파라미터 (단계 크기, 온도, 지평선) 에 대한 민감성과 장거리 텍스트 편집 작업에서 기울기가 약화될 수 있는 가능성 등 한계를 인정합니다. 그들은 완전성을 희생하지 않고 선택성을 더욱 향상시키기 위해 적응형 층 선택 및 단계 크기 스케줄링에 대한 향후 연구를 제안합니다.