[논문 리뷰] How Transformers Learn Causal Structure with Gradient Descent

Eshaan Nichani, Alex Damian외 1명|arXiv (Cornell University)|2024. 02. 22.

Bayesian Modeling and Causal Inference인용 수 5

한 줄 요약

이 논문은 간단화된 두-layer 트랜스포머에서의 경사하강법이 첫 번째 어텐션 층에 잠재 인과 구조를 인코딩함으로써 인과 그래프를 학습하고, 맥락 내 Markov 체인 설정에서 유도 헤드가 나타남을 보이며; 어텐션의 gradient 정보가 상호정보를 반영한다는 것을 보인다.

ABSTRACT

The incredible success of transformers on sequence modeling tasks can be largely attributed to the self-attention mechanism, which allows information to be transferred between different parts of a sequence. Self-attention allows transformers to encode causal structure which makes them particularly suitable for sequence modeling. However, the process by which transformers learn such causal structure via gradient-based training algorithms remains poorly understood. To better understand this process, we introduce an in-context learning task that requires learning latent causal structure. We prove that gradient descent on a simplified two-layer transformer learns to solve this task by encoding the latent causal graph in the first attention layer. The key insight of our proof is that the gradient of the attention matrix encodes the mutual information between tokens. As a consequence of the data processing inequality, the largest entries of this gradient correspond to edges in the latent causal graph. As a special case, when the sequences are generated from in-context Markov chains, we prove that transformers learn an induction head (Olsson et al., 2022). We confirm our theoretical findings by showing that transformers trained on our in-context learning task are able to recover a wide variety of causal structures.

연구 동기 및 목표

• gradient 기반 학습이 트랜스포머에서 인과 구조를 어떻게 유도하는지 이해를 촉진한다.
• 잠재 그래프를 고정된 인과 구조로 정의한 시퀀스 무작위 과제를 소개한다.
• 그라디언트 하강법 하에서 두-layer 어텐션 전용 트랜스포머의 학습 다이내믹스를 분석한다.
• 어텐션 행렬의 기울기가 상호 정보(mutual information)를 포착하고 그래프 엣지를 드러낸다.
• 이 접근법이 다중 헤드 구조를 통해 비-트리 그래프에 일반화되는지 분석하고 분포 외(out-of-distribution) 성능을 평가한다.

제안 방법

• 간단화된 두-layer 분리된 트랜스포머를 정의하고 A^(1) 및 A^(2)에 초점을 맞춘 축소 모델을 구성한다.
• 토큰 위치에 대한 잠재 DAG로 정의된 인과 구조를 갖는 무작위 시퀀스 과제를 구성한다.
• 그라디언트 하강법이 첫 어텐션 층(A^(1))에 이를 인코딩함으로써 잠재 그래프를 회복한다는 것을 보인다.
• 첫 어텐션 층의 기울기가 토큰 간의 카테고리된 상호정보를 나타내고 데이터 처리 부등식으로 엣지 회복을 유도한다.
• 특수 사례 분석: 맥락 내 Markov 체인은 인-context 추정(transitions)을 수행하기 위한 유도 헤드를 개발한다.
• 다중 헤드 확장을 제공하여 비-트리 그래프를 헤드 간에 분산시키고 경험적으로 검증한다.

실험 결과

연구 질문

• RQ1 트랜스포머에 대한 경사하강법이 고정된 인과 그래프로 생성된 데이터에서 잠재 인과 구조를 회복할 수 있는가?
• RQ2 학습 중 트랜스포머의 어텐션 층 안에 인과 구조가 어떻게 인코딩되는가?
• RQ3 맥락 학습 시나리오에서 어떤 기본 원리(예: 유도 헤드)가 나타나는가?
• RQ4 잠재 그래프가 트리이 아닐 때 모델의 성능은 어떤가, 다중 헤드 설계가 이를 해결할 수 있는가?
• RQ5 학습된 모델은 분포 외(transitions)에 일반화되는가?

주요 결과

• 두-layer 분리된 트랜스포머에서 그라디언트 하강법은 잠재 인과 그래프를 첫 어텐션 층의 인접 행렬로 인코딩하는 것을 학습한다.
• 첫 어텐션 층의 기울기는 토큰 간의 카테고리된 상호정보에 해당하며 데이터 처리 부등식으로 엣지 학습을 그래프에 집중시킨다.
• 맥락 내 Markov 체인의 특수 경우에서 모델은 맥락 내 전이 추정을 수행하기 위한 유도 헤드를 개발한다.
• 그래프가 트리가 아닐 때, 다중 헤드 트랜스포머가 잠재 그래프를 헤드 간에 분산시켜 문제 해결 동작을 달성할 수 있다.
• 실험적으로 학습된 트랜스포머는 제안된 과제에서 다양한 인과 구조를 회복하고 전이(transitions)에 대한 분포 외 일반화를 보인다.
• 이론적 보장(Theorem 1 및 Theorem 2)은 명시된 가정하에서 모집단 손실 수렴과 OOD 일반화를 확립한다.