실리콘 속의 친구와 할머니: 언어 모델 내 엔티티 셀의 위치 특정

언어 모델이 방대한 지식을 어떻게 저장하고 꺼내 쓰는지에 대한 근본적인 메커니즘을 규명했다. 특정 엔티티에만 반응하는 '엔티티 셀'을 찾아내어 이를 억제하거나 활성화함으로써 모델의 지식을 정밀하게 제어할 수 있는 새로운 경로를 제시한다.

Transformer 모델이 수많은 정보를 처리할 때, 특정 개념이 모델 전체에 넓게 퍼져 있는지 아니면 특정 지점에 집중되어 있는지는 오랜 의문이었다. 본 논문은 뇌 과학의 '할머니 세포' 가설을 인공지능에 적용하여, 특정 엔티티(예: 버락 오바마)를 처리할 때만 유독 안정적으로 활성화되는 개별 뉴런이 존재함을 찾아냈다.

연구팀은 엔티티에 대한 다양한 질문 프롬프트를 던졌을 때 레이어별 뉴런의 활성화 패턴을 분석했다. 이때 Attention Mechanism이 입력 토큰들을 결합하여 의미를 형성하는 초기 단계에서, 특정 MLP 뉴런이 해당 엔티티의 '정체성'을 대표하는 신호를 생성한다는 사실을 발견했다. 이는 모델이 텍스트의 겉모습(Surface Form)에 상관없이 추상적인 엔티티 개념으로 정보를 압축하여 처리함을 의미한다.

결과적으로 이 뉴런은 지식 회로의 '입구' 역할을 한다. 이 뉴런 하나만 끄면 모델은 해당 엔티티에 대해 아무것도 기억하지 못하는 상태가 되며, 반대로 이 뉴런의 신호를 강제로 주입하면 모델이 전혀 다른 맥락에서도 해당 엔티티와 관련된 정보를 출력하도록 유도할 수 있다. 이는 복잡한 LLM 내부 지식을 단일 뉴런 단위에서 정밀하게 조작할 수 있는 가능성을 열어준다.

엔티티 셀 위치 특정(Localization)을 위해 특정 엔티티에 대한 100여 개의 속성 질문 프롬프트를 생성한다. 각 프롬프트 실행 시 엔티티 토큰 위치에서의 MLP 뉴런 활성도를 기록하고, 프롬프트 간 활성도가 가장 안정적이고 높은 뉴런을 후보로 선정한다.

활성도 표준화 및 안정성 점수 계산을 수행한다. [입력 프롬프트 x에 대해 레이어 l, 뉴런 j의 활성도 a_lj(x)를 측정] → [일반적인 프롬프트 집합에서의 평균과 표준편차를 이용해 z-score로 변환] → [K개의 프롬프트에 대한 평균의 제곱을 표준편차로 나눈 안정성 점수 S_lj를 산출] → [이 점수가 가장 높은 뉴런이 엔티티 정체성을 대표하는 지표가 된다].

인과적 검증을 위해 부정적 절제(Negative Ablation)와 제어된 주입(Controlled Injection) 기법을 적용한다. 부정적 절제는 [특정 뉴런의 활성값에 음수 계수 alpha를 곱하여] → [해당 엔티티 정보의 흐름을 차단하고] → [모델의 답변 정확도 하락을 측정하여] → [해당 뉴런의 필요성을 확인한다]. 제어된 주입은 [빈칸 토큰 위치에 특정 엔티티 셀의 활성값을 덮어씌워] → [모델이 해당 엔티티를 떠올리게 유도함으로써] → [단일 뉴런의 정보 복구 능력을 평가한다].

Qwen2.5-7B 모델 분석 결과, 식별된 엔티티 셀의 99%가 0~5번 레이어에 집중되어 있었다. 이는 모델이 추론 초기 단계에서 입력된 텍스트를 추상적인 엔티티 정체성으로 빠르게 변환함을 시사한다.

부정적 절제 실험에서 200개 엔티티 중 131개가 특정 뉴런 하나만 억제해도 해당 엔티티에 대한 지식만 선택적으로 망각하는 '엔티티 특화 건망증' 현상을 보였다. 대조군 엔티티에 대한 지식은 1.0에서 0.996으로 거의 유지된 반면, 타겟 엔티티의 지식 보유량은 0.123까지 급감했다.

제어된 주입 실험에서는 단일 뉴런 활성화만으로도 63.3%의 확률(pass@5)로 올바른 엔티티 관련 답변을 복구해냈다. 이는 엔티티 정보가 모델 내부에 매우 압축적이고 국소적인 형태로 존재함을 나타낸다.

표면 형태 강건성 테스트에서 'Barack Obama', 'Obama', 'FBI' 등 다양한 변형과 다국어 표기(라틴어, 히브리어, 중국어 등)에 대해서도 동일한 뉴런이 최상위 활성도를 유지하며 정형화된 표현을 제공함이 확인됐다.

본 연구는 Transformer의 MLP 블록이 키-값 메모리 역할을 한다는 가설을 확장하여, 특정 사실이 아닌 엔티티 정체성에 고정된 '엔티티 셀'의 존재를 규명했다. 이는 특정 관계(Relation)가 아닌 엔티티 그 자체에 반응하는 뉴런을 타겟팅한다는 점에서 기존 지식 편집 연구와 차별화된다.

아키텍처 측면에서 엔티티 셀은 MLP의 down-projection 직전의 채널 값으로 정의된다. 이 뉴런은 잔차 연결(Residual Stream)에서 특정 패턴을 감지하는 감지기(Detector)와 해당 스트림에 엔티티 일관적 표현을 쓰는 기록기(Writer)의 쌍으로 작동하며, 초기 레이어에서의 활성화는 하위 토큰들을 하나의 의미 단위로 통합하는 역할을 수행한다.

위치 특정 알고리즘은 중요도 가중 안정성 기준(Importance-scaled stability criterion)을 사용한다. 높은 평균 활성도는 보상하고 프롬프트 간 높은 상대적 변동성은 페널티를 주어, 다양한 문맥에서도 변하지 않는 엔티티의 핵심 신호를 추출한다. 이는 단순 활성도 기반 탐색보다 훨씬 정밀한 타겟팅을 가능하게 한다.

실험 결과는 Qwen 계열 모델에서 가장 뚜렷하게 나타났으며, OLMo나 Llama 등 타 모델에서는 국소화 패턴이 더 넓게 퍼져 있거나 깊은 레이어에서 발견되는 경향을 보였다. 이는 사전 학습 데이터의 구성이나 토큰화 방식이 엔티티 표현의 국소화 정도에 영향을 미칠 수 있음을 시사한다.

본 연구는 PopQA 데이터셋에 한정되어 있으며, Qwen2.5-7B 모델에서 가장 강력한 증거가 발견되어 모든 모델로의 일반화는 추가 검증이 필요하다. 또한 첫 번째 토큰 확률 기반의 메트릭을 주로 사용하여 긴 답변이 필요한 복잡한 사실 관계를 완벽히 포착하지 못할 수 있다.