Attention Rollout은 Transformer (architecture) 기반 모델, 주로 대규모 언어 모델을 해석하는 방법으로, 네트워크의 모든 층을 통해 각 입력 토큰에서 다른 모든 토큰으로 흐르는 누적 주의(attention)를 계산한다. 개별 층의 주의 가중치가 종종 노이즈가 많고 모델의 최종 결정을 직접 반영하지 않는다는 관찰에 대한 응답으로 도입된 Attention Rollout은 층 간 주의 행렬을 집계하여 단일하고 일관된 중요도 맵을 생성한다. 이 기법은 설명 가능한 AI 분야에서 입력의 어떤 부분이 모델의 예측에 가장 많이 기여하는지 시각화하는 데 널리 사용되며, 디버깅, 편향 탐지, 모델 동작 이해를 돕는다.
핵심 아이디어는 transformer에서 각 층의 주의 메커니즘이 이전 층의 표현에 대한 가중 합을 계산한다는 개념에 기반한다. 연속된 층의 주의 행렬을 곱함으로써 초기 토큰의 정보가 네트워크를 통해 최종 출력으로 어떻게 전파되는지 추적할 수 있다. 이는 노드가 토큰이고 간선이 주의 가중치인 방향 그래프에서 경로를 따르는 것과 유사하다. 결과적인 롤아웃 행렬은 토큰 영향에 대한 전역적 관점을 제공하며, 이는 단일 층의 주의 맵을 보는 것보다 종종 더 해석 가능하다.
배경 및 동기
2017년 논문 "Attention Is All You Need"에서 Google DeepMind와 University of Toronto의 연구자들이 도입한 Transformer는 자기 주의 메커니즘에 의존하여 토큰 간 의존성을 포착한다. 각 층은 각 토큰이 다른 토큰에 얼마나 주의를 기울이는지 나타내는 주의 점수를 계산한다. 그러나 이러한 점수는 중요도로 직접 해석할 수 없다. 예를 들어, 토큰이 의미적 가중치가 거의 없는 구두점에 주의를 기울이거나, 주의가 많은 토큰에 분산될 수 있다. 또한, 이후 층의 주의는 이전 층에 의해 이미 변환된 표현을 기반으로 계산되므로, 원시 가중치는 이러한 복합 효과를 고려하지 않는다.
초기 설명 가능성 시도는 개별 헤드나 층의 주의 맵을 시각화하는 데 초점을 맞췄지만, 종종 상충되거나 혼란스러운 결과를 낳았다. Anima Anandkumar와 Jakob Uszkoreit 같은 연구자들은 주의 패턴이 층과 헤드에 걸쳐 매우 다양할 수 있어 결론을 도출하기 어렵다고 지적했다. 이는 네트워크 전체에 걸쳐 정보를 집계하는 방법의 개발을 촉진했다.
알고리즘
Attention Rollout은 주의 가중치가 정보 라우팅의 한 형태를 나타낸다는 가정 하에 작동한다. 알고리즘은 다음과 같이 진행된다:
- 각 층 \( l \)에 대해 (sequence_length, sequence_length) 형태의 주의 행렬 \( A_l \)을 얻는다. 여기서 \( A_l[i][j] \)는 토큰 \( i \)에서 토큰 \( j \)로의 주의 가중치이다. 이 행렬은 일반적으로 해당 층의 모든 주의 헤드에 대해 평균화된다.
- 각 주의 행렬에 항등 행렬을 추가하여 잔차 연결을 고려한다. 이를 통해 각 토큰이 자신의 정보를 유지할 수 있다. 결과 행렬은 \( \tilde{A}_l = A_l + I \)이다.
- \( \tilde{A}_l \)의 각 행을 정규화하여 합이 1이 되도록 하여 행렬이 확률적(stochastic)으로 유지되도록 한다.
- 모든 층에 걸쳐 정규화된 행렬을 곱하여 롤아웃 행렬 \( R \)을 계산한다: \( R = \tilde{A}_1 \cdot \tilde{A}_2 \cdot ... \cdot \tilde{A}_L \). 여기서 \( L \)은 층 수이다.
결과 행렬 \( R \)은 네트워크의 모든 경로를 고려하여 각 토큰에서 다른 모든 토큰으로의 총 주의 흐름을 제공한다. 항목 \( R[i][j] \)는 출력에서 토큰 \( i \)의 표현에 영향을 미치는 토큰 \( j \)의 정보 비율로 해석할 수 있다.
이 방법은 2020년 OpenAI와 Stanford AI Lab의 연구자들이 제안한 논문에서 처음 소개되었으며, GPT-2 및 기타 모델에서 효과를 입증했다. 그들은 Attention Rollout이 감정 분석 및 대명사 해결과 같은 작업에서 중요한 토큰을 식별할 수 있으며, 종종 인간 판단과의 정렬 측면에서 단일 층 주의보다 우수함을 보여주었다.
모델 해석에서의 응용
Attention Rollout은 자연어 처리 및 그 이상의 다양한 작업에 적용되었다. Machine learning 연구에서 다음과 같이 사용된다:
- 핵심 토큰 식별: 롤아웃 행렬을 검사함으로써 실무자는 어떤 입력 단어나 하위 단어가 모델 출력에 가장 강하게 영향을 미치는지 확인할 수 있다. 예를 들어, 감정 분류 작업에서 롤아웃은 "terrible"이나 "amazing"과 같은 단어를 매우 영향력 있는 것으로 강조할 수 있다.
- 편향 탐지: 모델이 편향된 예측을 초래하는 방식으로 인구 통계적 용어에 일관되게 주의를 기울이는 경우, 롤아웃은 이러한 패턴을 드러내어 연구자가 불공정한 동작을 완화하는 데 도움을 줄 수 있다.
- 모델 오류 디버깅: 모델이 잘못된 예측을 할 때, 롤아웃은 입력의 무관한 부분에 집중했는지 보여주어 훈련 데이터나 아키텍처 개선을 안내할 수 있다.
- 모델 비교: 다른 모델에 대해 롤아웃 행렬을 계산함으로써 내부 정보 흐름을 비교할 수 있으며, 이는 모델 선택 및 아키텍처 차이 이해에 유용하다.
텍스트를 넘어, Attention Rollout은 컴퓨터 비전 작업에 사용되는 비전 transformer(ViT)에도 적용되었다. 예를 들어, Google Cloud와 Amazon Web Services의 연구자들은 이미지 분류에 이를 적용하여 이미지의 어떤 영역이 예측에 기여하는지 시각화하고, 의료 영상 분석 및 자율 주행 시스템을 지원했다.
한계 및 비판
널리 사용됨에도 불구하고 Attention Rollout에는 여러 한계가 있다. 첫째, 주의 가중치가 정보 흐름의 확률 분포로 취급될 수 있다는 가정이 항상 유효한 것은 아니다. 주의 가중치는 학습된 표현을 기반으로 계산되며, 반드시 인과적 영향을 반영하지는 않는다. Melanie Mitchell과 Aleksander Madry 같은 비평가들은 주의가 설명이 아니며, 롤아웃과 같은 방법이 오해를 불러일으키는 해석을 생성할 수 있다고 주장했다.
둘째, 층 간 주의 행렬의 곱셈은 많은 층을 가진 깊은 네트워크에서 수치적 문제를 초래할 수 있다. 롤아웃 행렬이 지나치게 확산되어 모든 토큰이 유사한 중요도를 가지거나, 반복된 곱셈으로 인해 소수의 토큰에 집중될 수 있다. 이는 실제 사용에서 결과의 유용성을 떨어뜨릴 수 있다.
셋째, Attention Rollout은 주의 층 사이에 발생하는 비선형 변환(피드포워드 네트워크, 층 정규화)을 고려하지 않는다. 이러한 변환은 정보 흐름을 크게 변경할 수 있으므로, 단순한 곱셈 모델은 중요한 효과를 놓칠 수 있다.
넷째, 이 방법은 헤드에 대해 주의를 평균화하여 서로 다른 헤드의 고유한 역할을 모호하게 만들 수 있다. 일부 헤드는 구문 관계를 포착하고 다른 헤드는 의미 관계를 포착할 수 있으며, 평균화는 이러한 뉘앙스를 잃게 한다.
변형 및 확장
이러한 한계를 해결하기 위해 연구자들은 여러 변형을 제안했다:
- Attention Flow: MIT CSAIL의 연구자들이 도입한 이 방법은 주의를 흐름 문제로 취급하고 최대 흐름 알고리즘을 사용하여 토큰 중요도를 계산하며, 곱셈 가정을 피한다.
- 잔차 가중치를 사용한 롤아웃: 일부 구현은 항등 행렬을 균일하게 추가하는 대신 잔차 연결의 강도를 반영하는 요소로 가중치를 부여한다.
- 층별 관련성 전파(LRP): Deep learning 커뮤니티의 이 기법은 관련성 점수를 네트워크를 통해 역방향으로 전파하여 주의 기반 방법에 대한 대안을 제공한다.
- 통합 그래디언트 및 SHAP: 이들은 주의 가중치에 의존하지 않고 그래디언트 정보나 게임 이론적 개념에 기반한 특성 기여도 방법이다. 주의 기반 방법에 대한 건전성 검사로 자주 사용된다.
이러한 대안에도 불구하고 Attention Rollout은 단순성과 계산 효율성으로 인해 인기 있는 기준선으로 남아 있다. 추가 훈련이나 그래디언트 계산이 필요 없어 사전 훈련된 transformer에 쉽게 적용할 수 있다.
구현 고려 사항
Attention Rollout 구현은 현대 딥러닝 프레임워크에서 간단하다. PyTorch 또는 TensorFlow에서 순방향 패스에 훅을 걸어 주의 행렬을 포착할 수 있다. 주요 단계는 다음과 같다:
- 각 주의 층에 순방향 훅을 등록하여 주의 가중치를 저장한다.
- 순방향 패스 후 주의 가중치를 헤드에 대해 평균화한다.
- 항등 행렬을 추가하고 행을 정규화한다.
- 행렬을 순차적으로 곱한다.
한 가지 실용적 고려 사항은 긴 시퀀스의 경우 주의 행렬이 커져 메모리 오버헤드가 발생할 수 있다는 점이다. 길이 1024의 시퀀스의 경우 각 행렬은 1024x1024이며, 많은 행렬을 곱하는 것은 계산 비용이 많이 들 수 있다. 그러나 일반적인 입력에서는 관리 가능하다.
또 다른 고려 사항은 이 방법이 표준 transformer 아키텍처를 가정한다는 것이다. Cross-Attention이 있는 모델(예: 인코더-디코더 모델)의 경우, 롤아웃은 인코더와 디코더 주의 간의 상호 작용을 처리하도록 조정되어야 한다. 이러한 경우 인코더와 디코더에 대해 별도의 롤아웃을 계산하거나 더 복잡한 방식으로 결합할 수 있다.
다른 설명 가능성 방법과의 관계
Attention Rollout은 Artificial intelligence 시스템을 위한 더 넓은 설명 가능성 기법의 지형 내에 위치한다. 종종 다음과 비교된다:
- 그래디언트 기반 방법: 출력의 입력 임베딩에 대한 그래디언트를 계산하여 민감도 측정을 제공한다. 이론적으로 더 견고하지만 역전파가 필요하다.
- 섭동 기반 방법: 입력 토큰을 수정(예: 제거 또는 마스킹)하고 출력 변화를 관찰한다. 모델에 구애받지 않지만 계산 비용이 많이 든다.
- 대리 모델: LIME이나 SHAP 같은 기법은 블랙박스 모델을 근사하기 위해 로컬에서 해석 가능한 모델을 훈련한다. 표 형식 데이터에 유용하지만 텍스트에는 덜 적합하다.
Attention Rollout은 순방향 패스 외에 추가 계산이 필요 없는 모델의 내부 주의 가중치에만 기반한다는 점에서 독특하다. 이는 대화형 디버깅 도구와 같은 실시간 응용 프로그램에 특히 매력적이다.
미래 방향
Transformer 모델이 계속 커지고 복잡해짐에 따라 견고한 설명 가능성 방법의 필요성이 증가한다. Attention Rollout은 여러 방식으로 진화할 가능성이 있다:
- 인과 방법과의 통합: 롤아웃을 인과 추론 기법과 결합하면 더 정확한 기여도를 제공할 수 있다.
- 다중 모달 모델 처리: OpenAI와 Anthropic이 개발한 것과 같이 텍스트, 이미지, 오디오를 동시에 처리하는 모델로 롤아웃을 확장한다.
- 자동 분석: 롤아웃 출력을 사용하여 모델 결정에 대한 자연어 설명을 자동 생성하며, 이는 규정 준수 및 감사에 유용할 수 있다.
- 벤치마킹: Carnegie Mellon University와 BAIR (Berkeley AI Research)의 연구자들이 제안한 것처럼 설명 가능성 방법의 충실도를 평가하기 위한 표준화된 벤치마크를 개발한다.
한계에도 불구하고 Attention Rollout은 설명 가능성 도구 키트의 기본 도구가 되었다. 단순성과 효과성 덕분에 많은 연구 논문과 실제 응용 프로그램에서 표준 기준선이 되었다. 분야가 발전함에 따라, 이는 더 정제되고 다른 기법과 결합되어 신경망의 내부 작동에 대한 더 깊은 통찰력을 제공할 것이다.
결론
Attention Rollout은 transformer가 정보를 처리하는 방식을 이해하는 실용적이고 직관적인 방법을 제공한다. 주의 가중치를 층 간에 전파함으로써 개별 층을 검사하는 것보다 종종 더 유용한 토큰 영향에 대한 전역적 관점을 제공한다. 결함이 없지는 않지만, 사용 용이성과 해석 가능성 덕분에 설명 가능성 문헌에서 확고한 위치를 차지했다. Transformer 모델로 작업하는 모든 사람에게 Attention Rollout을 이해하는 것은 강력하지만 불투명한 이러한 시스템을 해명하는 데 중요한 단계이다.