기술 요약: 어텐션(Attention)을 위한 하세 다이어그램(Hasse Diagrams)

문제 정의

트랜스포머 모델의 어텐션 마스크 설계는 전통적으로 경험적 휴리스틱이나 인과적 마스크(causal masks)에 대한 임의적인 수정 방식에 의존해 왔습니다. FlexAttention과 같은 연산자가 임의의 마스크를 지원하기는 하지만, 임의의 마스크에 의해 유도되는 정보 흐름 구조에 대한 체계적인 형식적 분석은 부족했습니다. 기존의 그래프 이론적 접근 방식은 마스크를 임의의 유향 그래프로 취급하는 경우가 많았으며, 이는 수학적 구조를 도출하기에는 너무 광범위했습니다. 또한, 투 스트림 어텐션(two-stream attention, XLNet)과 같은 설계가 설계 공간을 확장했음에도 불구하고, 일련의 학습 태스크로부터 새로운 어텐션 메커니즘을 직접 도출할 수 있는 체계적인 방법론은 존재하지 않았습니다. 핵심 과제는 멀티 레이어 트랜스포머의 정보 흐름을 특징짓고, 어텐션 마스크 및 병렬 학습 체계를 설계하기 위한 구성적 방법(constructive method)을 제공하는 엄밀한 수학적 프레임워크를 구축하는 것입니다.

방법론

본 논문은 트랜스포머 아키텍처, 순서론(order theory), 그리고 그래프 이론을 연결하는 이론적 프레임워크를 제안합니다. 방법론은 다음 두 가지 주요 단계로 진행됩니다.

1. 부분 순서(Partial Order)로서의 정보 흐름 특징화

저자들은 분석의 단위를 단일 어텐션 레이어에서 잔차 연결(residual connections)을 포함한 전체 멀티 레이어 트랜스포머 블록으로 확장합니다.

• 모델링: 단일 트랜스포머 레이어는 어텐션 마스크와 잔차 연결에 의해 결정되는 입력 위치에서 출력 위치로의 정보 흐름을 나타내는 이분 유향 그래프(bipartite directed graph)로 모델링됩니다.
• 수렴: $L$개의 레이어를 가진 멀티 레이어 트랜스포머의 경우, 도달 가능성 행렬(reachability matrix) $R^{(L)}$은 레이어 인접 행렬 $A$의 불리언 거듭제곱($R^{(L)} = A^L$)으로 정의됩니다. 본 논문은 $L \to \infty$일 때, $R^{(L)}$이 유일한 극한값 $R$으로 수렴함을 증명합니다.
• 전순서에서 부분 순서로: 잔차 연결으로 인해 극한 행렬 $R$은 반사적(reflexive)입니다. 도달 가능성의 추이성(transitivity)과 결합하여, $R$에 의해 정의된 관계는 **전순서(preorder)**를 형성합니다. 저자들은 상호 도달 가능한 위치들(동치 클래스)을 식별하고 그 몫 집합(quotient set)을 형성함으로써, 이 클래스들 위에 **부분 순서(partial order)**를 유도합니다.
• 하세 다이어그램(Hasse Diagram): 트랜스포머의 정보 흐름 구조는 이 부분 순서의 하세 다이어그램으로 시각화됩니다. 다이어그램의 노드는 위치들의 동치 클래스를 나타내며, 엣지는 직접적인 의존성(covering relations)을 나타냅니다.

2. 최소 공통 상위 그래프를 통한 구성적 설계

병렬 학습 태스크의 설계를 그래프 이론적 문제, 즉 일련의 하세 다이어그램들의 **최소 공통 상위 그래프(minimal common supergraph)**를 찾는 문제로 재정의합니다.

• 태스크 정의: 학습 태스크는 (입력 시퀀스, 레이블 매핑, 어텐션 마스크)의 삼중항으로 정의됩니다. 밀집 태스크(dense tasks)는 마스크의 인접 행렬이 극한 도달 가능성 행렬과 일치하는 태스크입니다.
• 노드 동치성: 서로 다른 하세 다이어그램에 있는 두 노드는 상류 의존 구조(upstream dependency structures)가 동형(isomorphic)이고 입력 멀티셋(multisets)이 동일할 경우 동등한 것으로 간주됩니다.
• 최소 병합 태스크: 주어진 태스크 군에 대해, 저자들은 모든 정점 집합의 합집합을 구하고, 이를 동치 클래스로 분할하며, 대표 원소를 선택하여 새로운 하세 다이어그램을 형성함으로써 "최소 병합 태스크"를 구성합니다. 이 다이어그램은 최소 공통 상위 그래프를 나타냅니다.
• 희소화(Sparsification): 이론적 명확성을 위해 유도는 밀집 마스크를 가정하지만, 본 프레임워크는 동일한 하세 다이어그램 구조를 공유하면서 여러 레이어를 거쳐 수렴해야 하는 희소 최소 병합 태스크를 구축하는 것을 허용합니다.

주요 기여

1. 이론적 프레임워크: 본 논문은 충분히 깊은 트랜스포머의 정보 흐름이 하세 다이어그램으로 시각화 가능한 부분 순서로 수렴함을 입증합니다. 이는 임의의 어텐션 마스크에 대한 최초의 체계적인 수학적 구조를 제공합니다.
2. 설계 기준: 하세 다이어그램의 최소 공통 상위 그래프에 대한 기준을 도입하여, 정의된 일련의 학습 태스크로부터 어텐션 마스크를 직접 도출하는 구성적 방법을 제공합니다.
3. 새로운 어텐션 메커니즘: 이 프레임워크를 적용하여 두 가지 새로운 어텐션 메커니즘을 설계했습니다.
   • 블록 투 스트림 어텐션(Block Two-Stream Attention): 근사 없이 학습-추론 일관성을 보장하는 블록 생성 태스크용 마스크입니다. 마스크 토큰을 블록 단위로 그룹화하여 블록 내 완전 연결성을 갖도록 합니다.
   • 버터플라이 어텐션(Butterfly Attention): 모든 위치가 자신의 토큰을 참조하지 않으면서도(주변 토큰의 집합된 표현을 입력으로 사용) 100% 감독 밀도를 달성하는 완전 지도 양방향 어텐션 마스크입니다.
4. 통합 도출 파이프라인: 본 논문은 표준화된 어텐션 메커니즘 설계 절차를 제시합니다: 태스크 군 정의 $\to$ 단일 태스크 하세 다이어그램 도출 $\to$ 노드 동치성 계산 $\to$ 최소 병합 태스크 구성 $\to$ 마스크 행렬 도출.

결과

본 논문은 세 가지 사례 연구를 통해 프레임워크를 검증합니다.

• 인과적 어텐션(Causal Attention): 프레임워크는 표준 인과적 어텐션 마스크를 재구성하며, 이것이 자기회귀(autoregressive) 태스크 군의 최소 병합 태스크임을 증명합니다.
• 블록 투 스트림 어텐션: 도출된 마스크는 마스크 토큰을 입력 시퀀스의 끝에 배치함으로써, 블록 간의 인과적 제약을 유지하면서도 단일 순방향 패스(forward pass)를 통해 다음 블록의 토큰을 예측할 수 있는 블록 단위 생성을 가능하게 합니다.
• 버터플라이 어텐션: 도출된 마스크는 모든 위치에 대해 양방향 컨텍스트를 가능하게 합니다. 타겟 토큰을 주변 이웃들의 집합된 표현으로 대체함으로써, 모델은 셀프 어텐션 누수(self-attention leakage) 없이 전체 양방향 컨텍스트를 사용하여 타겟 토큰을 예측할 수 있습니다.

의의 및 주장

본 논문은 경험적인 시행착오에서 벗어나 구성적이고 체계적인 방법론을 제공한다고 주장합니다. 하세 다이어그램을 통해 학습 태스크와 정보 흐름 사이의 관계를 공식화함으로써, 연구자들은 다음과 같은 작업을 수행할 수 있습니다.

• 임의의 어텐션 구성에 대한 정보 흐름 구조를 엄밀하게 분석합니다.
• 특정 태스크 군에 최적화된 학습 체계 및 어텐션 마스크를 체계적으로 도출합니다.
• 근사 없이 개념화하거나 구현하기 어려웠던 새로운 어텐션 구조(Butterfly 및 Block Two-Stream 등)를 발견합니다.

저자들은 현재 프레임워크가 주로 밀집 학습 태스크의 병합을 다루며, 희소 변형은 이를 정교화한 형태라고 언급합니다. 또한 선형 어텐션 메커니즘(예: KDA, GDN)으로 프레임워크를 확장하고, 마스크드 오토인코더(MAE)와 같은 시나리오를 위한 새로운 태스크 군을 탐구하는 것을 향후 과제로 꼽았습니다.