[논문 리뷰] Unveiling Transformers with LEGO: a synthetic reasoning task

Yi Zhang, Artūrs Bačkurs외 4명|arXiv (Cornell University)|2022. 06. 09.

Explainable Artificial Intelligence (XAI)인용 수 4

한 줄 요약

이 논문은 구조적 어텐션 메커니즘을 통해 연쇄적 논리적 추론을 따르는 능력을 평가하는 합성 추론 작업인 LEGO를 소개한다. 저자들은 동일한 토큰을 대상으로 하는 새로운 '연관성' 어텐션 패턴을 규명하고, 구조적 어텐션을 통해 사전학습이 성능 향상에 기여하는 방식을 입증하며, 대규모 사전학습에서 FLOPs를 줄이면서도 정확도를 유지하거나 향상시키는 LEGO 어텐션 모듈을 제안한다.

ABSTRACT

We propose a synthetic reasoning task, LEGO (Learning Equality and Group Operations), that encapsulates the problem of following a chain of reasoning, and we study how the Transformer architectures learn this task. We pay special attention to data effects such as pretraining (on seemingly unrelated NLP tasks) and dataset composition (e.g., differing chain length at training and test time), as well as architectural variants such as weight-tied layers or adding convolutional components. We study how the trained models eventually succeed at the task, and in particular, we manage to understand some of the attention heads as well as how the information flows in the network. In particular, we have identified a novel \emph{association} pattern that globally attends only to identical tokens. Based on these observations we propose a hypothesis that here pretraining helps for LEGO tasks due to certain structured attention patterns, and we experimentally verify this hypothesis. We also observe that in some data regime the trained transformer finds ``shortcut" solutions to follow the chain of reasoning, which impedes the model's robustness, and moreover we propose ways to prevent it. Motivated by our findings on structured attention patterns, we propose the LEGO attention module, a drop-in replacement for vanilla attention heads. This architectural change significantly reduces Flops and maintains or even \emph{improves} the model's performance at large-scale pretraining.

연구 동기 및 목표

• 연쇄적 추론 능력을 격리하고 평가할 수 있는 합성 추론 작업을 설계하는 것.
• 비관련 NLP 작업에서의 사전학습과 데이터셋 구성이 추론 작업에서의 일반화에 미치는 영향을 조사하는 것.
• 어텐션 헤드 행동을 식별하고 분석하는 것, 특히 새로운 구조적 어텐션 패턴인 '연관성'과 같은 패턴을 중심으로 한다.
• 사전학습이 추론 작업 성능을 향상시키는 이유를 이해하고, 이를 실험적으로 검증하는 것.
• 표준 다중헤드 어텐션의 즉각적인 교체로 사용할 수 있는 더 효율적인 어텐션 메커니즘인 LEGO 어텐션을 제안하는 것.

제안 방법

• LEGO 작업은 동일한 토큰의 연쇄를 추적하고 추론할 수 있도록 설계된 합성 추론 벤치마크이다.
• 저자들은 다양한 연쇄 길이를 가진 LEGO 데이터로 트랜스포머를 훈련하고, 활성화 및 어텐션 패턴 분석을 통해 어텐션 헤드 행동을 분석한다.
• 특정 헤드가 유일하게 동일한 토큰에 전역적으로 어텐션을 기울이는 새로운 '연관성' 어텐션 패턴을 도입하여 반복되는 요소의 정확한 추적을 가능하게 한다.
• 사전학습이 구조적 어텐션 패턴을 유도함으로써 성능 향상에 기여한다는 가설을 수립하고, 통제된 추론 실험을 통해 이를 검증한다.
• 표준 다중헤드 어텐션의 교체용으로 사용할 수 있는 LEGO 어텐션 모듈을 제안하며, 동일한 토큰에만 어텐션을 제한하여 FLOPs를 감소시킨다.
• 다양한 연쇄 길이를 가진 훈련 및 테스트 시나리오를 포함한 여러 데이터 제약 조건에서 실험을 수행하여 탄력성과 단순화된 학습 전략의 영향을 평가한다.

실험 결과

연구 질문

• RQ1 비관련 NLP 작업에서의 사전학습과 데이터셋 구성은 추론 작업에서 미지의 연쇄 길이로의 일반화 능력에 어떤 영향을 미치는가?
• RQ2 LEGO 작업 훈련 중에 어떤 특정 어텐션 패턴이 나타나며, 이러한 패턴은 토큰 연쇄에 대한 추론을 어떻게 지원하는가?
• RQ3 사전학습이 합성 추론 작업 성능을 향상시키는 이유는 무엇이며, 이는 구조적 어텐션 패턴의 출현 때문인가?
• RQ4 어떤 데이터 제약 조건에서 모델은 진정한 추론 대신 표면적인 통계적 신호에 의존하는 '단순화된 해결책'을 학습하게 되며, 이로 인해 분포 이질성에 대한 탄력성이 떨어지는가?
• RQ5 수정된 어텐션 메커니즘인 LEGO 어텐션은 대규모 사전학습에서 표준 어텐션 대비 더 높은 효율성과 성능을 달성할 수 있는가?

주요 결과

• 특정 헤드가 유일하게 동일한 토큰에 전역적으로 어텐션을 기울이는 새로운 '연관성' 어텐션 패턴이 규명되었으며, 이는 시퀀스 내 반복 요소의 정확한 추적을 가능하게 한다.
• 비관련 NLP 작업에서의 사전학습은 LEGO 작업 성능 향상에 기여하며, 이 성능 향상은 구조적 어텐션 패턴의 출현과 강하게 상관관계가 있다.
• 일부 데이터 제약 조건에서는 모델이 진정한 추론 대신 표면적인 통계적 신호에 의존하는 단순화된 해결책을 학습하여 분포 이질성에 대한 탄력성이 떨어진다.
• 제안된 LEGO 어텐션 모듈은 FLOPs를 크게 감소시키면서도 대규모 사전학습 벤치마크에서 성능을 유지하거나 향상시킨다.
• 연구 결과는 어텐션 패턴이 무작위가 아니며, 가중치 묶음이나 컨볼루션 구성 요소와 같은 아키텍처 선택이 추론 행동에 영향을 준다는 것을 확인한다.
• 결과적으로 사전학습의 이점이 단지 데이터 다양성 때문이 아니라, 특정한, 작업에 관계없는 어텐션 인덕티브 바이어스를 유도하기 때문일 수 있음을 시사한다.