Explaining Attention with Program Synthesis

문제 정의
알고리즘적 관점에서 딥 뉴럴 네트워크가 수행하는 계산을 이해하는 것은 머신러닝 분야의 중대한 과제로 남아 있다. 해석 가능성에 대한 기존의 접근 방식은 주로 뉴런이나 피처(feature)에 의미를 부여하는 방식에 의존해 왔으며, 이는 인간이 정의한 개념을 위한 프로브(probe)를 사용하는 하향식(top-down) 방식이나, 입출력 요약을 활용하는 상향식(bottom-up) 방식으로 나뉜다. 이러한 방법들은 강력하지만, 신경 계산에 대한 완전하고 형식적인 설명을 제공하는 데는 한계가 있다. 또한, 중간 구성 요소에 자연어 설명을 라벨링하는 것은 모호하거나 정형화하기 어려운 해석을 초래할 수 있다. 저자들은 수십억 개의 파라미터를 가진 모델의 복잡성과 자연어 설명 사이의 간극을 메울 수 있는 매개체가 필요하다고 주장한다. 즉, 인간이 읽을 수 있으면서도 동시에 형식적으로 검증 가능한 매개체가 필요하다는 것이다.

방법론
본 논문은 트랜스포머 언어 모델의 어텐션 헤드(attention head) 계산을 설명하기 위해 **프로그램 합성(program synthesis)**에 기반한 프레임워크를 제안한다. 핵심 목표는 입력 텍스트 시퀀스($X$)를 특정 어텐션 헤드의 로직과 유사하게 매핑할 수 있는 실행 가능한 파이썬 프로그램($\pi$)을 찾는 것이다.

이 프레임워크는 네 가지 단계로 진행된다:

1. 어텐션 맵 추출 (Attention Map Extraction): 저자들은 타겟 모델로부터 학습 시퀀스(구조적 단순성을 위해 TinyStories 데이터셋 사용) 전반에 걸친 실제 어텐션 활성화 행렬을 추출한다. 이들은 주목할 만한 토큰 쌍 간의 상호작용에 집중하기 위해 상위 2.5%의 어텐션 가중치만을 필터링하여 추출한다.
2. 프로그램 합성 (Program Synthesis): 보조적인 대규모 언어 모델(LLM)이 합성 에이전트로서 추출된 어텐션 패턴과 입력 시퀀스를 바탕으로 프롬프트된다. 이 에이전트는 추출된 패턴을 재현하는 실행 가능한 파이썬 함수를 생성하는 임무를 맡는다. 합성 에이전트는 언어적 및 수치적 처리를 위한 NumPy, spaCy, NLTK와 같은 라이브러리를 사용할 수 있다.
3. 정제 및 순위 지정 (Refinement and Ranking): 후보 프로그램들은 구문론적 유효성을 검증받으며, 젠슨-샤논 거리(Jensen-Shannon Distance, JSD)를 사용하여 실제 어텐션 패턴과 비교되어 점수가 매겨진다. 상위 후보들은 실제 패턴과 예측된 패턴을 대조하여 프로그램을 수정하도록 하는 피드백 조건부 정제 루프를 거친다.
4. 선택 (Selection): 각 헤드에 대한 최적의 프록시(proxy) 프로그램은 홀드아웃(held-out) 데이터 시퀀스에 대한 유사성(Intersection over Union, IoU로 측정)을 최대화함으로써 선택된다.

주요 기여

• 프로그래밍적 해석 가능성 (Programmatic Interpretability): 저자들은 뉴럴 어텐션 헤드를 상징적이고 실행 가능한 코드로 대체하는 방법을 도입하여, 자연어 설명에 대한 형식적으로 검증 가능한 대안을 제시한다.
• 교체를 통한 인과적 검증 (Causal Validation via Replacement): 상관관계만을 관찰하는 기존 연구와 달리, 이 접근 방식은 학습된 신경 구성 요소를 합성된 프로그램으로 직접 교체함으로써 설명의 유효성을 검증한다.
• 확장성 (Scalability): 이 방법은 BERT-Base, GPT-2-Small, TinyLlama-1.1B, Llama-3B의 네 가지 서로 다른 트랜스포머 아키텍처에 적용되어, 다양한 모델 크기와 아키텍처(양방향 및 인과적 모델)에 대한 적용 가능성을 입증하였다.

결과

• 높은 정렬도 (High Alignment): 테스트된 모든 모델의 상당 부분의 어텐션 헤드가 높은 정확도를 가진 실행 가능한 프로그램에 의해 근사될 수 있다. 가장 성능이 좋은 프로그램들은 GPT-2에서 69%, TinyLlama에서 74%, Llama-3B에서 79%의 평균 IoU 점수를 달성했다.
• 모델 규모 및 아키텍처: 저자들은 자기회귀(autoregressive, 디코더) 모델의 어텐션 헤드가 BERT와 같은 양방향(bidirectional, 인코더) 모델보다 적합시키기 더 쉽다는 것을 관찰했다. 또한, 모델의 규모가 커질수록 적합 품질이 향상되는데, 이는 대규모 모델일수록 헤드 간의 기능적 전문화가 더 크게 나타날 수 있음을 시사한다.
• 인과적 충실도 (Causal Fidelity): 어텐션 헤드를 합성된 프로그램으로 교체했을 때, 모델은 기능적 성능을 유지한다. 구체적으로, 헤드의 최대 25%를 교체하더라도 퍼플렉시티(perplexity) 증가율은 16%에 불과했다.
• 다운스트림 성능 (Downstream Performance): 어텐션 헤드의 30~40%를 최적의 적합 프로그램으로 교체하더라도 HellaSwag, PIQA, SciQ, ARC-Easy, Social IQA, COPA를 포함한 다운스트림 질의응답 벤치마크 성능이 크게 저하되지 않는다.
• 상관관계 (Correlation): 프로그램의 IoU 유사도와 헤드 교체 시 발생하는 퍼플렉시티 증가 사이에는 강한 음의 상관관계(Spearmann $r > 0.9$)가 존재하며, 이는 IoU가 기능적으로 충실한 프로그램을 선택하는 신뢰할 수 있는 프록시임을 나타낸다.

의의 및 주장
본 논문은 최첨단 언어 모델에서도 상당한 비율의 어텐션 패턴이 상징적인 용어로 이해될 수 있다고 주장한다. 주요 의의는 학습된 신경 구성 요소가 실행 가능한 코드를 통한 상징적 대리물(symbolic surrogates)로 교체될 수 있으며, 이를 통해 모델의 동작을 크게 변화시키지 않을 수 있다는 점을 입증했다는 데 있다. 이는 현대의 LM 기반 프로그램 합성 방법이 딥 네트워크가 어떻게 작동하는지에 대한 더 넓은 질문에 접근하는 유효한 프레임워크를 제공함을 시사한다.

저자들은 본 연구를 트랜스포머의 완전한 상징적 특성화를 향한 단계로 규정하며, 이를 통해 모델의 동작을 가중치 조작이 아닌 논리를 통해 직접 추론하고, 수정하고, 테스트할 수 있게 된다고 본다. 그러나 논문은 현재 모든 모델의 어텐션 헤드가 완전히 특성화된 것은 아니며(많은 헤드가 여전히 낮은 IoU 점수를 기록함), 현재의 프로그램 라이브러리가 효과적이기는 하나 아직 완전하지 않다는 점을 인정하며 한계점에 대해 겸허한 태도를 유지한다. 이 연구는 트랜스포머 아키텍처의 완전한 상징적 증류(symbolic distillation)를 목표로 하는 미래 연구의 토대를 마련한다.