Review

TL; DR

Summary

• 연구진

  

• github: https://github.com/submarat/removing-layer-norm

• 인용수: 5

• openreview forum : https://openreview.net/forum?id=VPtHqcafIY
  • 리뷰어들에게는 8,8,8,6점을 받았지만, AC는 meta reiew 에서 “rejection-like” meta review를 주었음! 한번 읽어보시는 걸 추천 ~~

Background & Motivation

• Layer Normalization Layer

  

  • ⊙ : Hadamard (element-wise) product
  • 학습 단계에서의 stabilize를 위한 Layer 임
    • 선행연구에서는 LN이 confidence regulation에 기여한다고 밝히기도 했음

      https://proceedings.neurips.cc/paper_files/paper/2024/hash/e21955c93dede886af1d0d362c756757-Abstract-Conference.html (NIPS’24)

    • 미리 학습해둔 값을 사실상 linear transformation처럼 쓸 수 있는 Batch normalization 과 달리, LN은 inference때 매번 수행해야 하는 연산임. (=non-linear function)
    • LN의 non-linearity로 인해, model의 mechanistic interpretability가 방해됨
      • mechanistic interpretability 란? 모델을 더 작은 component의로 분해하고, 각 component의의 개별적인 효과와 상호 작용을 이해하는 것
      • why? LM이 residual stream activation에 따라 각 component(=sublayer, head, …)가 영향을 받기 때문
        • Eq (1) 에서, 현재 시점에서의 평균을 빼고 std를 나눠서 정규화하는데, 각 단계에서의 평균, std가 다 다르기 때문에 뒤 단계까지 영향을 끼침
      • 그래서 기존 interpretability 연구에서는 NL을 linear transformation으로 근사화해서 수행했음
        • 참고 논문:
          • https://arxiv.org/abs/2406.17759
          • [ICLR’24] Copy Suppression: Comprehensively Understanding an Attention Head https://openreview.net/forum?id=g8oaZRhDcf
        • but, 정확하지 않고, 이렇게 학습된 모델은 실제 LLM과 다름
      • 또는 아예 LN 을 없애고, element-wise tanh function을 사용함
        • 참고 논문: https://arxiv.org/abs/2503.10622
        • but, 여전히 non-linear function이기 때문에 interpretability 연구에 부적합

      ⇒ 이미 학습된 실제 transformer에서 LN을 제거한 버전을 분석해보자 !!

Contributions (What they’ve revealed)

• transformer로부터 LN을 제거하였음 & LN layer 없이도 작동하며, original model과 유사한 cross-entropy loss를 달성할 수 있음을 보임

  • FakeLN block 정의

    

    • standard deviation 대신에, fixed scalar $\bar{\sigma}_{avg}$ 사용
    • BS: batch size

  • 구체적인 removal process
    1. 0번째 MLP layer의 LN ( $LN^0_{MLP}$) 를 제거하고 fine-tuning
    2. 1번째 MLP layer의 LN ( $LN^1_{MLP}$) 를 제거하고 fine-tuning
    3. …
    4. 0번째 query/key LN ( $NL^0_{qk}$) 을 제거 하고 fine-tuning
    5. 1번째 query/key LN ( $NL^1_{qk}$) 을 제거 하고 fine-tuning
    6. …
    7. 0번째 value LN ( $NL^0_{v}$) 을 제거 하고 fine-tuning
    8. 1번째 value LN ( $NL^1_{v}$) 을 제거 하고 fine-tuning
    9. …
    10. Final NL $NL^f$제거
    • 한번에 없애지 않는 이유?

      : 모든 LN을 한 번에 없애면 모델 성능이 회복 불가능하게 붕괴하기 때문에, LN block을 하나씩 제거하고 loss spike가 가라앉을때까지 (일정 step만) fine-tuning함

    • LN layer 제거를 보다 안정적으로 수행하기 위한 Auxiliary loss 활용

      

      • λ: hyperparameter
      • Limitations에서 언급하길, GPT2-large/XL model을 안정화하기 위해 추가적으로 도입한 loss라고 함 !!

  • 학습 결과
    • training loss

      

    • 학습 이후에도 비슷한 cross entropy loss를 달성함

      

• LN layer가 없을 때 model의 interpretability가 향상되었음을 검증

  : 기존 interpretability 분석 연구에서 많이 활용되는 direct logit attribution (DLA), attribution patching 사용

  1. LN layer가 없을 때 DLA error가 50% → 0%로 감소함
     • direct logit attribution(DLA)란?
       • 어떤 component의 direct effect(=특정 component가 중간 component를 거치지 않고 output에 주는 효과)를 선형 근사로 추정하는 방법
       • 기존 LN은 nonlinearity를 가지고 있기 때문에, DLA가 근사화됨

         ⇒ 이예 따라 DLA와 DE간의 차이(=error)가 있을 수 밖에 없음!

     • DLA가 DE와 평균적으로 얼마나 어긋나는가를 평가하기 위해 Normalized Mean Absolute Error (NMAE) 활용하여 측정

       

       • w/LN 49.07% → w/FakeLN 0.00%로 감소함

  2. attribution patching은 LN layer 유무와 무관함
     • activation patching이란?

       : 어떤 내부 activation이 정말로 원인인지 확인하기 위해, 정답을 잘 맞히는 clean prompt에서 나온 activation을 corrupted prompt에서 나온 activation에 같은 위치로 대체하여, 각 component의 결과를 관찰하는 방법

       

       ⇒ 정확하지만, 실제로 activation을 하나하나 바꿔 넣으면서 확인해야 하기에 연산량 커짐

     • attribution patching이란?

       : activation patching의 first-order Taylor approximation

       

       ⇒ 근사하는 과정에서 attribution patching errors 발생

     • 결과

       

       • LN free model이라고 해도, attribution patching errors 는 여전함

• 추가 분석 수행

  1. LN layer가 residual stream geometry를 어떻게 바꾸는가
     • 선행연구에서, “first position token의 hidden representation L2 norm이 유난히 크며, 이는 attention sink의 핵심 mechanism임”을 밝혔음
       • 참고: 🥈Interpreting the Repeated Token Phenomenon in Large Language Models ( https://openreview.net/forum?id=WVth3Webet&noteId=hWnUO0wR7D)
     • NL-free model에서도 같은 현상이 나타나는지를 관찰함

       

       • LN-free model은 first token을 다른 token과 거의 동일하게 취급함
       • attention sink rate도 55.3% → 45.3%로 감소함

  2. LN-free model은 조금 더 overconfident함
     • LN-free model을 개발하는 와중에, GPT-2 Medium의 경우, 출력 분포의 평균 엔트로피는 2.86 → 2.53 으로 감소함을 발견
     • 이에 따라, entropy neuron(=confidence neuron)을 분석해봄
       • 기존 연구와 동일한 방법으로 발견 및 분석 수행함

         : weight norm이 크고, 모든 output logit에 거의 같은 영향을 주는 (token ranking 자체를 바꾸지 않는) neuron

         • Universal Neurons in GPT2 Language Models
         • Confidence regulation neurons in language models
       • 결과 분석

         

         • entropy neuron인 1083, 1108, 3144의 cross entropy loss 값이 NL-free에서 크게 감소함

           ⇒ 즉, confidence neuron이 기능을 하지 않음