트랜스포머 뉴런의 읽기-쓰기 관계로 분류한 약화 뉴런과 그 과대한 영향력

무엇인가

이 논문은 트랜스포머 기반 대규모 언어모델 내부의 MLP 뉴런이 실제로 어떤 입력-출력 함수를 구현하는지를 다룬다. 기존 뉴런 해석 연구는 주로 두 갈래였다. 하나는 뉴런이 어떤 문맥에서 활성화되는지를 관찰하는 것이고, 다른 하나는 출력 가중치가 어떤 방향을 쓰는지를 분석하는 것이다. 저자들은 이 둘만으로는 뉴런의 메커니즘을 온전히 설명하지 못한다고 본다. 뉴런은 언제 켜지는지, 무엇을 쓰는지뿐 아니라 잔차 스트림에서 읽어들이는 방향과 다시 써 넣는 방향 사이의 관계로도 정의되기 때문이다. 특히 SwiGLU, GEGLU 같은 게이트 활성화 함수는 입력 쪽 가중치가 두 개라서 ReLU 모델과 사정이 다르며, 게이트 활성화 함수에 대한 입력-출력 분류 체계는 이전에 제대로 만들어지지 않았다.

어떻게 동작하나

제안 방법은 단순하다. 각 뉴런에서 세 개의 가중치 벡터, 즉 게이트 가중치 w_gate, 입력 가중치 w_in, 출력 가중치 w_out를 꺼내고 이들 사이의 코사인 유사도 세 쌍을 계산한다. 뉴런이 잔차 스트림에 더하는 양은 Swish(<w_gate, x>) 곱하기 <w_in, x> 곱하기 w_out이므로, w_in과 w_gate는 뉴런을 발화시키는 방향, w_out은 뉴런이 되돌려 쓰는 방향으로 읽을 수 있다. cos(w_in, w_out)이 강한 음수면 뉴런이 탐지한 방향을 약화시키는 것이므로 약화 뉴런(weakening neuron)이라 부른다. 이를 바탕으로 강화, 약화, 조건부 강화, 조건부 약화, 비례 변화(proportional change), 직교 출력이라는 분류를 제시한다. 조건부 강화와 조건부 약화는 w_in과 w_out이 같은 방향이지만 w_gate가 그와 직교해, 특정 조건이 잔차 스트림에 있을 때만 해당 방향을 강화하거나 약화한다. 실제 코사인 값은 0이나 ±1에 딱 맞지 않으므로 임계값 τ=±0.5로 가장 가까운 원형 클래스에 배정하고, cos(w_in, w_gate)가 나머지 두 코사인과 어긋나는 경우에는 atypical을 붙인다. 무작위 기준선으로는 무작위 초기화 모델에 해당하는 i.i.d. 가우시안 벡터와, 레이어별로 코사인을 뒤섞은 베이스라인을 쓴다.

무엇과 다른가

가중치 기반 분석은 24개의 비 GLU 트랜스포머 모델(인코더-디코더, 인코더 전용, 디코더 전용, ReLU와 GELU, 언어 및 비언어 학습 데이터)과 12개의 GLU 기반 LLM(SwiGLU와 GEGLU)에 적용됐다. 그중 9개 큰 모델의 레이어별 cos(w_in, w_out) 중앙값을 보면 패턴이 일관된다. 초·중반 레이어에서는 중앙값이 0보다 크게 높아 강화 뉴런이 다수를 차지하고, 후반 레이어에서는 중앙값이 0 아래로 내려가 약화 뉴런이 상대적 다수가 된다. Llama-3.2-3B를 자세히 보면 입력 조작 뉴런, 즉 출력이 입력과 직교하지 않는 뉴런이 전체의 25%이고 초·중반 레이어(0부터 세어 7~11층)에서는 50%까지 올라간다. Llama에서는 입력 조작 뉴런의 80% 이상이 조건부 강화 하나에 속한다. 반면 후반 레이어의 입력 조작 뉴런 중에서는 약화 뉴런이 큰 비중을 차지하고 비례 변화 뉴런도 중요한 부류로 등장한다. 직교 출력으로 분류된 뉴런들도 실제로는 코사인이 유의성 임계값을 넘는 경우가 많아, 어느 정도 입력 조작을 수행한다고 저자들은 해석한다.

어떻게 쓰나

활성화 빈도 분석은 OLMo-7B-0424, Gemma-2-2B, Llama-3.1-8B, Llama-3.2-3B 네 모델과 Dolma에서 뽑은 2천만 토큰으로 수행됐다. 결과는 cos(w_in, w_out)과 게이트 값이 양수일 빈도 사이의 거의 선형적인 음의 관계다. OLMo에서 상관계수는 -0.91로, Gurnee 등이 GELU 모델에서 보고한 최대 -0.69보다 강하다. 즉 조건부 강화 뉴런은 거의 활성화되지 않고 약화 뉴런은 자주 활성화되며, 이 효과는 주로 중간 레이어에서 나온다. 절제 실험은 OLMo-7B와 Llama-3.2-3B에서 수행됐고, 각 실행마다 약화 뉴런 수만큼의 뉴런을 다른 클래스에서 지우거나 같은 레이어의 무작위 뉴런을 지우는 베이스라인과 비교했다. OLMo-7B의 약화 뉴런은 526개다. 손실과 출력 분포 엔트로피 같은 지표에서 가장 큰 영향을 준 것은 약화 뉴런이었는데, 이는 수백 개에 불과한 작은 클래스라는 점에서 예상 밖이다. 게다가 이 효과는 활성화 빈도만으로 설명되지 않는다. 약화 뉴런의 드물고 작은 음수 게이트 값조차 모델 출력에 영향을 준다.

전제와 한계

사례 연구는 OLMo-7B의 28.4737(강화)과 31.9634(약화, 32층 중 마지막 층) 두 뉴런을 다룬다. 약화 뉴런의 효과를 가장 극적으로 보여주는 예시는 "Yesterday (21 December) the Government announced a package of support for hospitality and leisure businesses that are losing trade because of the O" 다음에 "mic"(Omicron)을 예측하는 경우다. 깨끗한 실행에서 mic와 유사 토큰은 절제 실행보다 최대 12점까지 점수가 올라가고, 반대로 점수가 크게 깎이는 토큰은 없다. 즉 약화 뉴런의 사분면 (iii) 활성화가 정답 토큰을 밀어 올려 출력 분포를 날카롭게 만든다. 이 효과는 단일 뉴런으로는 재현되지 않고 8개의 약화 뉴런 부분집합으로 재현되며, 뉴런의 직접 효과가 관측 효과와 일치하지 않아 간접적이고 누적적인 효과로 보인다. 저자들은 게이트 값과 입력 값의 부호에 따라 네 사분면을 나눠 해석한다. 사분면 (iii)(x_gate<0, x_in<0)은 이론적으로 예상되지 않던 영역인데도 강화에 가까운 행동을 보이며, 사분면 (i)(둘 다 양수)은 중첩으로 생긴 원치 않는 방향을 약화시키는 오류 수정, 사분면 (ii)(x_gate>0, x_in<0)는 잔차 스트림의 상충 정보를 해소하는 충돌 해결로 해석한다. 사분면 (iv)의 활성화는 작고 드물며 절제해도 영향이 거의 없다.

실무 관점에서 이 논문이 주는 것은 뉴런 단위 해석과 개입을 위한 값싼 사전 지표다. 활성화 통계를 대량으로 수집하지 않고도 가중치 벡터 몇 개의 코사인만 계산하면, 어떤 뉴런이 모델 출력에 큰 영향을 줄지 후보를 추릴 수 있다. 특히 후반 레이어의 약화 뉴런은 수가 적지만 절제 시 손실과 엔트로피에 큰 변화를 만들기 때문에 가지치기나 모델 편집에서 무시하면 안 되는 대상이다. 또한 게이트 값이 음수인 구간이 기능을 담지 않는다는 통념을 재고해야 한다는 점도 실무적으로 중요하다. 조건부 절제를 여러 층에 동시에 적용할 때는 상류 뉴런 절제가 하류 뉴런의 활성화 부호를 바꿔 절제 대상 여부 자체를 흔들 수 있다는 경고도 함께 새겨둘 만하다.

저자들이 밝힌 한계는 분명하다. 이 방법은 단일 뉴런의 파라미터만 보는 해석이라 간단하고 효율적이지만 범위가 좁고, 다른 뉴런 분석 기법을 대체하는 것이 아니라 보완하는 것이다. 가중치의 수학적 유사도가 의미적 유사도와 일대일로 대응하지는 않으며, 부록 H에서는 거의 직교하는 벡터가 매우 비슷한 개념을 나타내는 사례도 찾았다. 또한 잔차 연결이 있어야 레이어 간 표현을 비교할 수 있다는 전제가 깔려 있어, 잔차 연결 없이 트랜스포머를 학습시키는 최근 방법이나 하이퍼 커넥션을 쓰는 모델에는 직접 적용할 수 없다. 절제 실험은 다음 토큰 예측 수준의 효과만 보여줄 뿐 뉴런 활성화에서 출력까지 이어지는 중간 단계를 밝히지 못하고, 조건부 절제는 여러 층에 걸쳐 적용할 때 취약하다. 가중치 분석은 다양한 모델에 적용했지만 전문가 혼합(MoE) 모델과 9B보다 큰 모델은 테스트하지 않았고, 계산량이 큰 코퍼스 실험은 더 적은 수의 모델에 집중했다. 무엇보다 저자들은 자신들의 해석 일부가 추측이며, 특히 약화 뉴런에 대한 이해가 아직 제한적이라고 인정한다. 중간 레이어의 패턴을 신뢰도 관리로 설명하는 가설도 확증된 증거가 없고 후반 레이어에는 적용되지 않는다.