[논문 리뷰] Thinking Like Transformers

Gail Garfinkel Weiss, Yoav Goldberg외 1명|arXiv (Cornell University)|2021. 06. 13.

Advanced Neural Network Applications인용 수 11

한 줄 요약

이 논문은 주어진 시퀀스 처리 작업의 최소 레이어 수와 헤드 수를 예측할 수 있는 기호 프로그래밍 언어인 RASP(Restricted Access Sequence Processing Language)를 소개한다. RASP는 어텐션과 피드포워드 연산을 일급 원자로 추상화하여 트랜스포머 인코더의 계산 행동을 모델링한다. 이를 통해 RASP 프로그램은 정렬, 히스토GRAM 계산, Dyck-언어 인식 등의 작업을 해결하기 위해 필요한 최소 레이어 수와 헤드 수를 예측할 수 있으며, 훈련된 트랜스포머는 RASP에서 유도된 어텐션 패턴을 따라 99% 이상의 정확도를 달성한다.

ABSTRACT

What is the computational model behind a Transformer? Where recurrent neural networks have direct parallels in finite state machines, allowing clear discussion and thought around architecture variants or trained models, Transformers have no such familiar parallel. In this paper we aim to change that, proposing a computational model for the transformer-encoder in the form of a programming language. We map the basic components of a transformer-encoder -- attention and feed-forward computation -- into simple primitives, around which we form a programming language: the Restricted Access Sequence Processing Language (RASP). We show how RASP can be used to program solutions to tasks that could conceivably be learned by a Transformer, and how a Transformer can be trained to mimic a RASP solution. In particular, we provide RASP programs for histograms, sorting, and Dyck-languages. We further use our model to relate their difficulty in terms of the number of required layers and attention heads: analyzing a RASP program implies a maximum number of heads and layers necessary to encode a task in a transformer. Finally, we see how insights gained from our abstraction might be used to explain phenomena seen in recent works.

연구 동기 및 목표

• 트랜스포머 인코더의 형식적 계산 모델을 제공함. 이는 RNN에 대한 유한 오토마타와 유사함.
• 도메인 특화 언어를 통해 작업을 기호 프로그램으로 표현함으로써 트랜스포머의 능력에 대한 고수준 추론을 가능하게 함.
• 특정 시퀀스 처리 작업을 구현하기 위해 트랜스포머에서 필요한 최소 레이어 수와 어텐션 헤드 수를 예측함.
• 훈련된 트랜스포머가 RASP에서 도출된 솔루션과 정확히 일치하는 어텐션 패턴을 학습할 수 있는지 검증함으로써 기호 추론과 신경망 행동을 연결함.
• 트랜스포머 변종의 경험적 현상과 효율적 아키텍처의 한계를 RASP 추상화를 통해 설명함.

제안 방법

• 시퀀스에 대한 선택, 필터링, 집계를 위한 원자 연산을 갖춘 최소한의 프로그래밍 언어로 RASP 설계. 어텐션과 피드포워드 제약 조건을 반영함.
• RASP 프로그램을 트랜스포머에서 요구하는 어텐션 패턴과 레이어 구조를 정의하는 등가의 어텐션 및 MLP 연산으로 컴파일함.
• 출력에 대한 표준 크로스 엔트로피 손실과 RASP 솔루션을 따라야 할 어텐션 패턴에 대한 MSE 손실을 함께 사용하여 트랜스포머 훈련.
• RASP 프로그램을 사용해 주어진 작업에 대해 필요한 최소 레이어 수와 헤드 수에 대한 이론적 하한을 도출함.
• 예측된 최소 크기에서 트랜스포머를 훈련하고, 헤드 또는 레이어 수를 줄였을 때의 정확도 저하를 측정함으로써 RASP 프로그램의 타당성 분석.
• Dyck-k 언어 인식을 위한 이전 작업을 재구성하고 개선함으로써 RASP를 적용하여 모든 k에 대해 Dyck-k를 완전히 인식할 수 있음을 증명함.

실험 결과

연구 질문

• RQ1 RNN에 대한 유한 오토마타와 유사하게 트랜스포머 인코더의 계산 행동을 모델링할 수 있는 기호 프로그래밍 언어를 구축할 수 있는가?
• RQ2 기호 프로그램에 의해 예측된, 주어진 시퀀스 처리 작업을 수행하기 위해 트랜스포머에서 필요한 최소 레이어 수와 어텐션 헤드 수는 얼마인가?
• RQ3 훈련된 트랜스포머가 주어진 작업에 대해 RASP 프로그램에서 도출된 어텐션 패턴을 정확히 학습할 수 있는가?
• RQ4 모델 크기를 줄였을 때 성능에 어떤 영향을 미치며, RASP는 이러한 실패 지점을 예측할 수 있는가?
• RQ5 RASP를 사용하여 트랜스포머 변종의 경험적 행동을 설명하거나 예측할 수 있는가? 예를 들어 효율적 아키텍처에서 성능 저하를 예측할 수 있는가?

주요 결과

• 이중 히스토GRAM, 정렬, 가장 빈번한 토큰 태깅 등의 작업에 대해 RASP 프로그램을 어텐션 패턴으로 컴파일한 결과, 훈련된 트랜스포머가 99% 이상의 테스트 정확도로 이를 재현함.
• 출력과 어텐션 모두에 대한 이중 감독을 받은 트랜스포머는 이중 히스토GRAM 작업에서 99.9%, 정렬 작업에서 99.8%의 정확도를 달성하여 RASP에서 지정한 어텐션 패턴과 밀도 있게 일치함.
• RASP가 예측한 최소 크기 이하로 모델 크기를 줄였을 경우 정확도가 급격히 떨어짐 — 예를 들어, 한 레이어를 더 줄였을 때 역순 작업에서 41.2%의 정확도를 기록함 — 이는 RASP가 아키텍처의 한계를 예측하는 데 유용함을 시사함.
• RASP 추상화 덕분에 이전 작업보다 더 간결하고 정확한 Dyck-k 언어 인식 구조를 설계할 수 있었으며, 모든 k에 대해 트랜스포머가 Dyck-k를 완전히 인식할 수 있음을 증명함.
• 정렬 작업에서 단일 레이어, 단일 헤드를 가진 트랜스포머가 근사적으로 완벽한 정확도를 달성한 것은 균일한 어텐션 패턴 덕분이었으며, 이는 버킷 정렬의 형태를 효과적으로 수행한 것으로 해석됨.
• RASP 분석에 따르면 '효율적'인 트랜스포머는 복잡한 정보 흐름이 필요한 작업에서 실패할 수 있으며, 이는 RASP 분석이 이를 정확히 예측함.