[논문 리뷰] Learning Transformer Programs

Dan Friedman, Alexander Wettig외 1명|arXiv (Cornell University)|2023. 06. 01.

Explainable Artificial Intelligence (XAI)인용 수 13

한 줄 요약

이 논문은 기계적으로 해석 가능하도록 설계된 채로 제약된 Transformer를 훈련시키고, 성능을 크게 저하시키지 않으면서 인간이 읽을 수 있는 프로그램(Python/RASP-영감)으로 변환합니다. 이는 맥락 내 학습, 알고리즘적 과제, NLP에서의 결과를 보여주고, 코드 수준의 디버깅을 통해 해석 가능한 회로를 제공합니다.

ABSTRACT

Recent research in mechanistic interpretability has attempted to reverse-engineer Transformer models by carefully inspecting network weights and activations. However, these approaches require considerable manual effort and still fall short of providing complete, faithful descriptions of the underlying algorithms. In this work, we introduce a procedure for training Transformers that are mechanistically interpretable by design. We build on RASP [Weiss et al., 2021], a programming language that can be compiled into Transformer weights. Instead of compiling human-written programs into Transformers, we design a modified Transformer that can be trained using gradient-based optimization and then automatically converted into a discrete, human-readable program. We refer to these models as Transformer Programs. To validate our approach, we learn Transformer Programs for a variety of problems, including an in-context learning task, a suite of algorithmic problems (e.g. sorting, recognizing Dyck languages), and NLP tasks including named entity recognition and text classification. The Transformer Programs can automatically find reasonable solutions, performing on par with standard Transformers of comparable size; and, more importantly, they are easy to interpret. To demonstrate these advantages, we convert Transformers into Python programs and use off-the-shelf code analysis tools to debug model errors and identify the "circuits" used to solve different sub-problems. We hope that Transformer Programs open a new path toward the goal of intrinsically interpretable machine learning.

연구 동기 및 목표

• 고위험 작업에서 감사 및 디버깅을 위해 본질적으로 해석 가능한 Transformer 모델의 필요성을 동기화합니다.
• Deterministic하게 인간이 읽을 수 있는 프로그램으로 매핑될 수 있도록 제약 하에서 Transformer를 훈련하는 프레임워크를 제안합니다.
• Transformer 프로그램이 다양한 알고리즘적 과제와 NLP 작업에서 경쟁력 있는 성능으로 문제를 해결할 수 있음을 보입니다.
• 훈련된 모델에서 실행 가능한 Python/RASP-형 프로그램을 자동으로 추출하여 회로 수준의 디버깅을 가능하게 합니다.

제안 방법

• 각 모듈이 고정된 변수 집합을 읽고 전용 직교 부분공간에 쓰도록 하는 해소된 잔류 흐름 제약을 도입합니다.
• Hard attention으로 구현된 이산적이고 해석 가능한 모듈(범주적 주의 헤드)을 정의하고, 최적화 과정에서 Gumbel-Softmax를 사용해 Relax합니다.
• 각 주의 헤드를 RASP 유사 predicate-aggregate 원시 명제로 매핑하고, 이산 가중치(πK, πQ, πV, Wpredicate)에 대한 분포를 학습하며 Gumbel 재매개변수를 사용해 샘플링합니다.
• 훈련 후에 이산 가중치를 최대화하고 주의 헤드를 select_closest 원시 함수로 변환해 Python 프로그램을 결정적으로 추출합니다.
• 단어 임베딩, 수치 주의, 피드포워드/조회유사 층을 포함하도록 프레임워크를 확장하여 프로그램 레퍼토리를 넓힙니다.
• 연습 코드와 디버깅 워크플로를 포함하여 해석 가능한 프로그램으로의 매핑에 대한 확장 및 세부 정보를 제공합니다.

실험 결과

연구 질문

• RQ1 제약 조건 하에서 해석 가능한 프로그램으로 결정적으로 매핑될 수 있는 Transformer 모델을 학습할 수 있는가?
• RQ2 이런 Transformer 프로그램이 해석 가능성을 유지하면서 맥락 내 학습, RASP 스타일 알고리즘 과제, NLP 벤치마크를 어느 정도까지 해결할 수 있는가?
• RQ3 읽기 가능한 Python/RASP 유사 코드로 변환될 때 학습된 프로그램과 회로의 질적 구조는 무엇인가?
• RQ4 표준 Transformer와 비교했을 때 다양한 난이도의 과제에서 정확도와 해석 가능성 측면에서 Transformer 프로그램은 어떤 차이를 보이는가?

주요 결과

데이터셋	설명	예시	k	L	H	M	정확도(Acc.)
Reverse	문자열 뒤집기.	reverse("abbc") = "cbba"	8	3	8	2	99.79
Histogram	각 토큰에 대해 시퀀스에서 해당 글자 등장 횟수.	hist("abbc") = "1221"	8	1	4	2	100.0
Double hist.	각 토큰에 대해 같은 히스토그램 값을 가진 고유 토큰의 수.	hist2("abbc") = "2112"	8	3	4	2	98.40
Sort	입력을 사전 순으로 정렬합니다.	sort("cbab") = "abbc"	8	3	8	4	99.83
Most-Freq	빈도 순으로 배열된 고유 입력 토큰, 동점을 위치로 구분.	most_freq("abbc") = "bac"	8	3	8	4	75.69
Dyck-1	각 위치 i에 대해 i까지의 입력이 Dyck-1(T)인가, 유효 접두사(P)인가, 아니면 유효하지(F)인가.	dyck1("()())") = "PTPTF"	16	3	8	2	99.30
Dyck-2	위와 같은 내용의 Dyck-2 버전.	dyck2("({})(}") = "PPPTPF"	16	3	4	4	99.09
(표는 본문에 따라 계속됩니다)

• Transformer 프로그래밍은 비슷한 규모의 표준 Transformer에 비해 여러 과제에서 합리적인 성능을 달성합니다.
• RASP 스타일 과제에서 여러 과제에서 99% 이상 정확도를 달성하지만 더 긴 입력에서 예외가 있습니다.
• 맥락 학습 토이 과제에서 모델은 헤드를 조합해 induction-head 동작을 재현하도록 학습하고 테스트 정확도에서 완벽에 가깝습니다.
• CoNLL-2003 NER에서 Transformer 프로그램은 표준 Transformer와 유사한 F1를 달성하고 unigram 베이스라인을 능가합니다.
• 추출된 Python/RASP 유사 프로그램은 해석 가능한 회로와 특징 가중치를 노출해 디버깅 및 회로 분석에 도움을 줍니다.
• 단조로운 증가나 더 큰 어휘를 다루는 경우 표준 Transformer가 일반적으로 Transformer 프로그램보다 우수하다는 trade-off가 존재하며, 확장에는 도전이 있습니다.