'인과적 어텐션(Causal Attention)' 다시 한 번 들여다보기

들어가며

인과적 어텐션. 오늘날 아주 큰 영향력이 있는 많은 AI 모델을 구동하는 핵심에 있는 기술입니다.

자기회귀 트랜스포머(Autoregressive Transformer) 구조를 가진 모델이라면, 사실상 인과적 어텐션은 필수적으로 따라오는 요소이기도 합니다. 과거의 데이터를 바탕으로 미래를 한 단계씩 예측하는 방식인 인과적 어텐션 메커니즘(Causal Attention Mechanism)에서는 각 토큰이 이전 토큰들만 살펴보고, 절대 미래를 내다보지 않습니다.

언뜻 보기에, 이런 방식이 완벽한 구조를 갖춘 것처럼 느껴지긴 하지만, 세상사가 그렇게 녹록치는 않죠. 만약에, 미래의 토큰이 현재의 순간에 대한 중요한 정보를 줄 수 있다면 어떨까요? 즉, 더 높은 수준의 추론을 수행할 수 있는 모델로 진화하려면, 전체의, 즉 글로벌한 맥락을 포착할 필요가 생기게 됩니다.

오늘은 이 인과적 어텐션을 확장할 수 있는, 토큰이 과거 뿐 아니라 미래의 토큰에 숨은 정보에도 어텐션을 기울일 기회를 주는, 흥미로운 아이디어를 탐구해보려고 합니다.

바이트댄스 팀에서 연구한 Causal Attention with Lookahead Keys (줄여서 CASTLE)은, 토큰이 처리될 때마다 키를 다이나믹하게 업데이트하는 방식으로 이런 아이디어를 구현합니다. 시드니 대학교와 상하이 교통 대학교에서 개발한, 또 다른 흥미로운 접근법 ‘미래 인식 인과 마스크 (Future-aware Causal Masks)’는 시각-언어 모델 (VLM)의 비전 작업에서 미래의 맥락에 접근하는 것이 아주 효과적이라는 점을 보여주기도 합니다.

이 두 가지 접근법이 가져오는 변화와 함께, '인과적'이라는 개념이 다른 관점에서 어떻게 재정의되고 있는지—즉, 실제적으로 인과-효과 추론 (Cause-Effect Reasoning)을 도입하는 방식도 살짝 알아볼까 합니다.

오늘 에피소드에서는 다음과 같은 내용을 다룹니다:

• 인과적 어텐션의 기초

• CASTLE이란 무엇인가?

  • 기본적인 인과적 어텐션 플로우의 변화

  • 내부의 메커니즘

  • CASTLE이 보여주는 결과

  • CASTLE의 장단점

• 인과적 어텐션과 VLM(Vision-Language Model)

  • 미래를 인식하는 인과적 마스크(Future-Aware Causal Mask)

  • 경량 미래 인식 어텐션(Light Future Aware Attention)의 작동 방식

  • 이점, 그리고 잠재적 이슈

• 다른 방법들: 인과에서 인과 효과(Cause-Effect) 관계로

• 맺으며

• 보너스: 참고자료

인과적 어텐션의 기초

현대적인 생성형 AI를 떠받치고 있는 튼튼한 두 개의 기둥이 있다면, 아마 첫 번째는 트랜스포머가 제공하는 핵심 아키텍처, 두 번째는 텍스트의 생성 전략을 잡아주는 자기회귀 기법(Autoregressive Technique) 두 가지일 겁니다.

자기회귀적 생성은, 모델이 텍스트를 한 토큰씩, 왼쪽에서 오른쪽으로 만들어내면서, 새로운 토큰을 예측할 때 이전의 토큰들만 참고하는 방식이죠. 자기회귀 모델이 효과적이고 신뢰할 만한 이유는 다음과 같습니다:

• 자기회귀는 사람이 문장을 하나씩 자연스럽게 만들고 이해하는 모습과 닮았습니다.

• 모델은 훈련 중에 다음 토큰을 예측하면서 훈련하는데, 이 과정은 실제의 텍스트 생성과 동일한 과정이라서 훈련과 추론이 안정적이고 일관되게 이루어질 수 있습니다.

• 지금까지 생성된 모든 내용을 바탕으로 조건을 설정해서 모델은 출력이 맥락과 조화를 이루게끔 합니다.

• 자기회귀 모드의 트랜스포머는 병렬로 효율적으로 훈련할 수 있어서, 배치 안의 모든 다음 토큰을 한 번에 예측합니다.

• 무엇보다, 수많은 AI 혁신을 이끌어 온 핵심적인 강점은, 자기회귀 모델이 데이터와 파라미터가 늘어날수록 놀라울 정도로 잘 확장된다는 점이죠.

이 자기회귀 전략은 내부적으로 어떻게 작동할까요? 자기회귀 트랜스포머는 기본적인 인과적 어텐션 메커니즘(Causal Attention Mechanism)을 사용해서, 각 토큰이 과거와 현재 토큰만 바라보고, 미래 토큰은 절대 보지 않게끔 합니다. 예를 들어서, 5번째 단어를 예측할 때는 1~4번째 단어만 참고하는 식이죠. 시퀀스에서 뒤에 오는 토큰을 차단하는 인과 마스크(Causal Mask)로 엄격하게 관리를 하구요.

표준적인 인과적 어텐션 메커니즘의 작동 방식을 단계별로 살펴보면 다음과 같습니다:

• 각 토큰은 쿼리, 키, 값(QKV)을 가집니다.

• 다음 토큰을 생성할 때, 모델은 새로운 쿼리를 사용해서 지금까지의 모든 키와 비교합니다.

• 이 비교는 어텐션 스코어를 만들어내고, 다시 가중치로 변환됩니다.

• 이 가중치가 값에 적용되어 다음 토큰의 출력값을 생성합니다.

• 각 키는 해당 토큰의 위치까지의 정보만 나타냅니다. 키는 나중에 절대 업데이트되지 않습니다.

• 추론 시, 모델은 한 번에 한 토큰을 생성하고, 새 토큰을 다시 자신에게 입력합니다.

이 구조에서 각 토큰의 쿼리, 키, 값(QKV)은 토큰의 표현에서 한 번 계산해서 고정되고, 오직 이전 토큰들의 정보만 담게 됩니다.

이런 구조는 꽤 오랫동안 잘 작동해 온 게 사실입니다. 그렇지만, 문장에서 나중에 나타나는 중요한 정보를 모델이 놓칠 수도 있겠죠 - 경우에 따라서는, 인과 마스크 때문에 모델이 문장이나 글 전체의 큰 그림을 이해하는 능력이 제한될 수도 있는 겁니다.

BERT(Bidirectional Encoder Representations from Transformers) 모델을 기억하시나요? 이 모델은 토큰을 양방향으로—왼쪽에서 오른쪽, 오른쪽에서 왼쪽으로—처리하는데, 이런 방식으로 맥락을 더 깊게 이해해서 단어의 의미를 더 풍부하게 이해하고, 모델이 결과적으로 맥락에 더 민감하게 작동하게끔 해 줍니다.

그렇다면, 자기회귀 모델에서 미래의 토큰을 살짝 엿볼 수 있게 한다면 어떨까요?

이 아이디어를 가지고, 여러 연구자들이 표준적인 인과적 어텐션 기법을 한 번 다시 검토해 볼 필요가 있겠다는 생각을 했고, 결과적으로 모델이 미래 토큰과 맥락에 주목할 기회를 주는 것이 나쁜 선택이 아니라는 점을 발견한 겁니다.

바이트댄스 시드(ByteDance Seed)의 최근 논문에서 바로 이 아이디어를 더 깊이 탐구하고, 이 개념이 성공적으로 적용된 다른 사례를 살펴봤습니다. 일단은, 바이트댄스의 새로운 기법, CASTLE(Causal Attention with Lookahead Keys)에서부터 시작해 볼까 합니다.

읽어주셔서 감사합니다. 친구와 동료 분들에게도 뉴스레터 추천해 주세요!