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

CASTLE이란 무엇인가?

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

앞서 말씀드린 대로, 표준의 인과적 어텐션에서는 쿼리, 키, 값(QKV)이 한 번 생성되면 고정됩니다. CASTLE의 인과적 어텐션 접근법은 이 부분에서 변화구를 줘서, 더 많은 토큰이 처리되는 과정에서 이전에는 고정되어 있었던 키(K)를 업데이트합니다. 이렇게 업데이트된 키를 미래 예측 키(Lookahead Keys)라고 부릅니다:

• 이 키는 이전 토큰에 속합니다.

• 하지만 이후에 나타나는 토큰의 정보를 통합 - 일정 정도 반영 - 합니다.

예를 들어서, 모델이 4번째 토큰을 생성할 때, 1번째 토큰의 키는 이미 1~4번째 토큰의 정보를 알고 있고, 6번째 토큰을 생성할 때는 1~6번째 토큰의 정보를 포함하도록 확장됩니다.

맥락이 펼쳐져 갈수록, 초기 토큰의 표현은 뒤따르는 단어들의 힌트를 담아서 더 풍부해집니다. 이건, 토큰의 메모리가 끊임없이 업데이트되는 것과 비슷합니다.

하지만 중요한 점은 CASTLE이 여전히 자기회귀 규칙을 준수한다는 것입니다:

• (t+1) 단계에서, 토큰 s의 업데이트된 키는 오직 토큰 t까지의 정보만 포함합니다.

• 어떤 토큰도 미래(미래 토큰)의 정보를 절대 보지 않습니다.

• 쿼리와 값 역시 지금까지 본 맥락으로 제한됩니다.

이 쯤에서 궁금해 하실 수도 있을 것 같습니다 - 왜 Key만 업데이트하냐는 거죠. 거기에도 이유가 있습니다.

• 쿼리(Q)는 다음 토큰을 생성할 때 한 번만 사용됩니다. 따라서, 과거의 쿼리를 업데이트한다는 건 아무런 효과가 없습니다.

• 반면에, 키(K)는 이후의 모든 쿼리에 재사용되고, 따라서 키를 업데이트하면 이후 단계의 어텐션 스코어 품질이 향상됩니다.

• 값(V)을 업데이트하는 것도 도움이 될 수 있지만, 계산 비용이 더 많이 듭니다.

그래서 CASTLE은 키(Key)에 집중해서, 이걸 통해서 훈련의 효율성을 유지하는 깔끔한 수학적 기법을 활용하는 겁니다.

내부의 메커니즘

CASTLE은 혼합 어텐션 메커니즘(Hybrid Attention Mechanism)을 사용하는데, 여기서 각각의 토큰은 인과적 키(고정, Causal Key)와 미래 예측 키(동적, Lookahead Key)를 모두 가집니다. 인과적 키는 앞에서 설명한 표준의 인과적 어텐션과 동일한 것이고, 미래 예측 키는 이 과정에 새롭게 추가된 요소입니다. 이걸 계산하기 위해서 CASTLE은 다음의 어텐션 메커니즘을 따릅니다:

• 각 토큰 s는 자신보다 뒤에 오는 토큰(s+1 … t)의 정보를 모읍니다. 예를 들어서, 6번째 토큰을 생성할 때 3번째 토큰의 미래 예측 키는 4번과 5번 토큰의 정보도 포함하고, 1~5번 토큰 전체를 주시할 수 - 바라볼 수 - 있습니다.

• 마스크는 어떤 토큰도 t를 넘어선 미래를 보지 않도록 보장합니다.

• 모델이 항상 모든 후속 토큰을 흡수하지눈 않습니다(소프트맥스 함수처럼). CASTLE은 시그모이드 함수(Sigmoid Function)를 사용해서 어느 정도의 유연성을 제공해서, 토큰이 각 후속 토큰의 정보를 선택적으로 받아들일 수 있습니다.

• 시퀀스가 커질수록 미래 예측 키는 계속 갱신되기 때문에, 동일한 토큰의 키가 t 단계와 t+1 단계에서 달라집니다.

CASTLE이 인과적 키와 미래 예측 키를 결합하는 방식은 다음과 같습니다:

• 다음 토큰을 생성할 때, 현재 토큰의 쿼리는 인과적 키와 미래 예측 키(Lookahead Keys) 모두와 비교가 됩니다.

• 그래서 여기서 두 가지의 스코어 집합을 만듭니다. 하나는 고정된 과거 정보를 반영하고, 다른 하나는 업데이트된 맥락을 반영합니다.

• 미래 예측 키의 스코어가 인과적 키 스코어와 합쳐지기 전에, SiLU(시그모이드 선형 유닛, Sigmoid Linear Unit) 게이팅 함수를 거칩니다. 시끄럽거나 관련 없는 과거 토큰이 지배하지 않도록 막고, 각 과거 토큰이 잊혀지는 정도를 조절하기 위해서죠.

• 이후 미래 예측 키와 인과적 키의 스코어를 결합하고, 결과는 소프트맥스 함수(Softmax Function, 숫자 벡터를 확률 분포로 바꾸는 함수)를 거쳐서 모델이 유효한 어텐션 가중치를 생성하도록 합니다.

• 마지막으로, 이 어텐션 가중치가 값에 적용되어서 해당 단계의 출력을 만듭니다.

위에서 설명한 것은, CASTLE의 순환(Recurrent) 버전이기 때문에, 미래 예측 키를 단계별로 업데이트합니다.

하지만 CASTLE은 병렬 형태로도 재구성할 수 있어서, 동일한 결과를 내면서 대규모의 사전 훈련을 할 수 있게끔 해 줍니다. 이 형태는 순환 업데이트를 마스크된 행렬 곱셈으로 바꾸고, 저차원 분해(Low-rank Decomposition)를 활용해서 비용과 메모리 사용량을 줄입니다. 결과적으로 모델은 전체 시퀀스의 출력을 한 번에 계산하고, 매 토큰마다 키를 업데이트하고 저장할 필요가 없습니다.

추론을 효율적으로 하게끔 하기 위해서 연구자들이 UQ-KV 캐시(UQ-KV Cache)를 도입했는데, 이는 표준 KV 캐시와 비슷하지만 미래 예측 키에 필요한 정보를 저장합니다. 이것들은 처음부터 다시 계산되지 않고, CASTLE이 캐시된 결과를 재사용하며 순환적으로 업데이트합니다.

또, 일반 트랜스포머처럼 CASTLE도 멀티 헤드 어텐션(Multi-head Attention)으로 확장될 수 있습니다. 각 헤드는 자체적인 인과적 어텐션과 미래 예측 어텐션을 계산하고, 결과를 연결해서 모델의 숨겨진 차원으로 다시 투영합니다.

CASTLE이 보여주는 결과

언어 모델링 벤치마크로 테스트했을 때, CASTLE은 표준적인 인과적 어텐션보다 꾸준히 향상된 성능을 보입니다. 우선, 모든 규모에서 훈련 및 검증 손실과 퍼플렉시티(Perplexity)를 줄여 줍니다:

• 소규모 모델(1억 6천만 파라미터)에서는 개선 폭이 작지만, 어쨌든 퍼플렉시티가 16.41에서 16.32로 줄어듭니다.

• 모델 규모가 커질수록 퍼플렉시티의 개선이 두드러집니다. 중간 모델(3억 5천만)에서는 14.00에서 13.67로, 대규모 모델(7억 5천만)에서는 12.27에서 11.84로 떨어집니다.

• 초대형 규모(13억 파라미터)에서는 이 점이 가장 두드러져서, CASTLE은 퍼플렉시티를 11.31에서 10.92로 줄입니다.

다양한 다운스트림 활동 벤치마크에서도 CASTLE 모델은 기본 모델을 뛰어넘는 성능을 보여줍니다:

• ARC-C(추론 벤치마크)에서 점수는 제로-샷(Zero-shot)에서 33.79%에서 35.32%로, 5-샷(Five-shot)에서 35.58%에서 39.08%로 올라갑니다.

• BoolQ에서는 초대형(XL) 모델의 정확도가 61.07%에서 62.6%로 상승해서, 더 나은 읽기 이해 능력을 보여줍니다.

• 상식 추론을 테스트하는 Winogrande에서는 정확도가 제로-샷에서 54.06%에서 56.59%로, 5-샷에서 52.72%에서 58.33%로 향상됩니다.

CASTLE은 단순히 텍스트 예측 뿐 아니라 다양한 다운스트림 작업에서 성능을 향상시켜주고, 특히 추론과 일반화 능력을 뛰어나게 만들어 줍니다. 특히 더 큰 모델에서 효과가 두드러지는데, 이는 더 풍부하고 변화하는 맥락을 충분히 담아낼 수 있는 용량(Capacity) 덕분이겠죠.

CASTLE의 장단점

이제 CASTLE 어텐션 메커니즘이 돋보이는 이유를 한 번 정리해 보겠습니다.

• CASTLE은 모델이 자기회귀의 구조를 유지하면서도, 전체 맥락을 볼 수 있게 합니다.

• 고정된 인과적 키의 안정성과 진화하는 미래 예측 키의 적응성을 균형 있게 조화시킵니다.

• 퍼플렉시티(Perplexity)와 손실이 줄어들어서, 모델이 텍스트를 더 자신 있게 예측합니다.

• CASTLE 모델은 특히 추론과 상식 기반 작업에서 더 높은 정확도를 달성하게 해 줍니다.

• 병렬 형태로 훈련 복잡도가 표준 어텐션과 비슷해서, 대규모 사전 훈련에 실용적입니다.

• CASTLE 모델은 기본 트랜스포머보다 더 많은 파라미터를 필요로 하지 않습니다.

하지만 CASTLE 어텐션 메커니즘의 효과를 떨어뜨리는 요소도 분명히 있습니다:

• 표준 어텐션에 비하면 아무래도 추가적인 계산 비용이 듭니다.

• UQ-KV 캐시(UQ-KV Cache) 같은 특수 최적화가 필요해서 구현이 좀 더 어렵습니다.

• 소규모 모델에서는 이점이 크지 않습니다.

• 현재 최첨단 모델에서 사용되는 수조 규모 토큰 훈련에서의 동작은 아직 완전히 검증되지 않았습니다.

그럼에도 불구하고, CASTLE은 자기회귀 모델의 인과적 어텐션을 새로운 시각으로 확장하는 기법으로, 기존의 규칙을 깨지 않으면서도 미래의 정보를 주시하는 새로운 방식을 제시합니다.

하지만 CASTLE이 이 방향을 탐구한 최초의 접근은 아닙니다. 바이트댄스 이전에 시드니 대학교와 상하이 교통 대학교의 개발자들이 인과적 어텐션이 시각-언어 모델(Vision-Language Models, VLMs)에 어떻게 적용될 수 있는지 탐구한 흥미로운 연구가 있었거든요.

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

표준의 인과적 어텐션은 순차적인 텍스트 처리에는 잘 맞지만, VLM은 좀 다른 이야기일 거라는 생각이 드시죠?

VLM은 전체적인 시각 정보를 처리하기 때문에, 당연하게도 서로 다른 부분을 함께 봐야 하는 경우가 많습니다. 시드니 대학교와 상하이 교통 대학교의 연구자들은 인과적 어텐션과 시각 인코더 간의 텐션(Tension; 갈등) 관계에 우선 주목했습니다. 먼저, 시각 인코더는 전체 이미지를 토큰 집합으로 압축하고, 각각의 시각 토큰은 이미 전역의(글로벌한) 이미지 맥락을 담고 있죠. 그런데 인과적 어텐션은 이 시각 토큰을 텍스트 토큰처럼 행동하도록 강제해서, 이미지에서 토큰 간 공유할 수 있는 유용한 정보를 제한하게 되죠.

연구자들은 시각 토큰에 대한 인과 마스크를 완화해서 미래 이미지 영역의 정보를 활용하면, 시각-언어 추론 작업에서 모델 성능이 더 좋아진다는 점을 발견했습니다. 동시에 텍스트는 여전히 엄격한 인과성을 유지해야 한다고 주장했습니다. 그래서 이들은 두 가지를 고려한 미래 인식 인과 마스크(Future-aware Causal Masks)를 제안했습니다.

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

연구자들은 세 가지 특화된 마스크를 설계했습니다:

• 전체 미래 마스크(Mf, Full Future Mask): 시각 토큰이 모든 미래 토큰(시각+텍스트)을 볼 수 있습니다. 이건 네비게이션이나 이미지 시퀀스의 상태 변화 같은 시간적 다중 이미지 작업에 유용합니다.

• 시각-시각 마스크(Mv2v, Visual-to-Visual Mask): 시각 토큰이 미래 시각 토큰만 볼 수 있고, 텍스트는 제외됩니다. 이건 객체 변화나 프레임 간 관계를 탐지하는 시각 관계 작업을 향상시킵니다.

• 시각-텍스트 마스크(Mv2t, Visual-to-Textual Mask): 시각 토큰이 미래 텍스트 토큰만 볼 수 있습니다. 이는 OCR-VQA, TextVQA 같은 텍스트가 풍부한 이미지 작업에서 시각 패치가 이미지에 포함된 단어와 연결되도록 도와줍니다.

시각 토큰이 전체 미래 맥락을 모두 볼 수 있게 하는 건 꿈같은 일이긴 하지만, 실제로는 토큰이 하나씩 생성되는 자기회귀 디코딩에서 추론 속도가 너무 크게 느려집니다. 그래서 연구자들은 이에 대한 타협안으로 ‘경량 미래 인식 어텐션(Light Future Aware Attention)’을 제안했습니다.

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

VLM은 보통 두 단계로 실행됩니다:

• 프리필 단계(Prefill Stage), 전체 입력을 한 번에 처리하는 단계입니다.

• 디코딩 단계(Decoding Stage), 모델이 출력 토큰을 단계별로 생성하는 단계입니다.

경량 미래 인식 어텐션(Light Future Aware Attention)에서 ‘미래의 정보’에 대한 접근은 프리필 단계에서만 일어납니다. 즉, 미래의 정보가 초기 토큰으로 압축된다는 의미이고, 디코딩이 시작될 때 모델은 추가적인 부담 없이 일반적인 인과 마스크(Causal Mask)를 유지할 수 있습니다.

단계별로 어떤 일이 어떻게 일어나는지 살펴보겠습니다:

• 자기회귀 모델에서는 맨 처음 토큰들이 많은 어텐션을 모으는 “싱크” 역할을 합니다(어텐션 싱크 현상, Attention Sink Phenomenon). 이 특징이 이 기법의 장점으로 활용됩니다.

• 프리필 단계(Prefill Stage)에서 모델은 미래 맥락을 작은 표현으로 요약하는 커널 풀링 작업(Kernel Pooling Operation)을 적용합니다.

• 이렇게 압축된 미래 표현은 초기 토큰(프리픽스, Prefix)에 다시 통합되고, 이제 과거 토큰들은 미래에 대한 힌트를 담고 있게 됩니다.

• 모델이 출력을 생성할 때는 여전히 엄격하게 왼쪽에서 오른쪽으로의 인과 마스크(Causal Mask)를 사용합니다. 과거 토큰이 이미 미래 맥락을 포함하고 있으니, 모델은 자기회귀를 깨지 않으면서 이점을 얻습니다.

그럼, 경량 미래 인식 어텐션(Light Future Aware Attention)의 효과에 대해 알아보겠습니다.

이점, 그리고 잠재적 이슈

• 전체 미래 마스크(Full Future Mask)에 비해서, 경량 미래 인식 어텐션은 ‘정확도’ 관점에서의 이점을 대부분 유지하면서도 2~3배 속도를 높입니다(아래 결과를 확인해 보세요).

• 일반적으로 경량 미래 인식 어텐션(Light Future Aware Attention)은 추론 속도를 빠르게 유지하면서 풍부한 맥락의 이점을 보존합니다.

• 모델은 시간적 추론, 시각 관계 추론, 텍스트가 풍부한 질의응답 같은 비전 작업에서 더 나은 성능을 보입니다.

• 최종 마스크는 엄격히 인과적(Causal)으로 유지되어서 일관성을 보장합니다.

그렇지만, 미래 인식 어텐션의 한계로 몇 가지 중요한 이슈가 있습니다:

• CASTLE 방식과 비슷하게, 커널 풀링(Kernel Pooling)이나 프리픽스 병합(Prefix Merging) 같은 추가 기법이 필요한 복잡한 접근법입니다.

• 적절한 전략(전체, 시각-시각, 시각-텍스트)을 선택하고 병합 여부를 결정하는 과정은 수작업의 복잡성을 더합니다.

• 텍스트 중심이나 검색 중심 작업에서는 개선이 별로 없습니다. 하지만, 미래 인식 마스크가 애초에 주로 비전 작업의 성능 향상을 위해 고안되었다는 점을 고려하면, 이건 큰 문제는 아닐 수도 있겠습니다.

• 미래 인식 어텐션은 여전히 근사적입니다. 미래 맥락을 프리픽스 토큰으로 압축하면 모델은 원시 미래 토큰에 완전히 접근하지는 못합니다.

이러한 이슈에도 불구하고, 이 방식은 VLM이 표준적인 인과적 어텐션을 미래 인식 어텐션으로 확장하면 큰 이점을 얻을 수 있다는 점을 보여줍니다. 결국, 핵심은 균형에 있는 것이죠.

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

오늘 에피소드에서 ‘인과적 어텐션’에 대한 이해를 해 가면서, “인과”가 항상 “하나씩”, “왼쪽에서 오른쪽”, “단계별”만을 의미하지는 않는다는 점을 꼭 덧붙여야 할 것 같습니다.

더 넓게 보면, “인과”는 데이터에서 인과-효과 관계(Cause-Effect Relationships)를 포착하는 데 중요한 Causal AI라는 주제와도 관련됩니다. 즉, 또 다른 하나의 연구 흐름에서는, 인과적 어텐션을 문자 그대로 올바른 인과-효과 패턴에 주목하도록 업그레이드하는 방법을 탐구하기도 합니다.

이 흐름의 기법들을 간단히 살펴보겠습니다:

• 인과 어텐션 튜닝(Causal Attention Tuning, CAT)은 LLM이 표면적 패턴을 따르는 대신 무엇이 원인이 되고 무엇이 결과인지 주목하고, 데이터에서 잘못된 상관관계를 막도록 훈련하는 방식입니다. CAT은 훈련 중 어텐션 맵에 인과 사전 정보(Causal Priors)를 직접 주입해서 원인에 주목하도록 유도합니다. 사람과 LLM이 생성한 신호를 사용해 대규모 “인과 힌트”로 이루어진 인과 그래프(Causal Graph)를 구축합니다.
  두 번째 부분은 인과 제약 어텐션 훈련(Causal Constraint Attention Training)입니다. CAT은 인과 토큰이 비인과 토큰만큼 어텐션을 받도록 추가 재어텐션 손실(Re-attention Loss)을 더합니다. 모든 레이어와 헤드에서 평균 어텐션 가중치를 비교한 뒤에, 이를 인과 인접 행렬(Causal Adjacency Matrix)로 조정합니다. 최종 훈련 손실은 일반적인 다음 토큰 예측 손실과 이 인과 손실을 결합해서 모델이 올바른 인과 구조에 집중하도록 합니다.

• 포커스 학습 방법(Learning to Focus, LeaF)은 교사 모델과 학생 모델을 비교해서 ‘혼란 토큰(Confounding Tokens)’—인과적이지 않지만 주의를 분산시키는 토큰—을 찾아냅니다. 이런 혼란 토큰은 증류(Distillation) 과정에서 제거되어서 학생 모델이 의미 있는 신호에만 집중하도록 학습합니다.

• AutoHFormer는 시계열 예측을 위한 트랜스포머로, 인과성, 효율성, 다중 스케일 패턴 인식을 균형 있게 설계했습니다. 시퀀스를 블록으로 나누어 병렬 예측을 수행한 뒤에, 각 블록 내에서 이를 정제합니다. 주요 기법은 다음과 같습니다:

  • 다이나믹 윈도우 어텐션(Dynamic Windowed Attention): 학습 가능한 슬라이딩 윈도우와 지수 감쇠(Exponential Decay)를 사용해서 계산 복잡도를 2차 미만으로 유지하며 엄격한 인과성을 보장합니다.

  • 적응형 시간 인코딩(Adaptive Temporal Encoding): 단기 역학을 위한 사인파 패턴과 장기 추세를 위한 학습 가능한 감쇠를 조합합니다.

• “인과성 변환: 사전 지식 통합을 통한 트랜스포머 기반 시간적 인과 발견” 논문은 다층 트랜스포머를 사용해서 시계열 데이터에서 인과 그래프(Causal Graph)와 시간 지연(Time Lags)을 학습하는 프레임워크를 제안합니다. 두 가지 주요 구성 요소가 있습니다:

  • 인과 인식 예측기(Causality-aware Forecaster): 미래 값을 예측하도록 훈련된 트랜스포머 모델로, 변수 간의 상호 영향을 파악합니다. 전문가가 특정 연결이 존재하지 않는다고 이미 알고 있으면, 어텐션 마스크를 적용해서 모델이 잘못된 연결을 학습하지 않도록 합니다.

  • 인과 그래프 추출기(Causal Graph Extractor): 훈련 후에, 입력에 대한 예측의 민감도를 기울기(Gradients)를 통해서 분석, 변수 간 관계와 시간 지연을 보여주는 인과 그래프를 만듭니다. 잘못된 연결을 줄이기 위해서 HITL(Human-in-the-loop)를 추가, 사용자가 그래프를 검토하고 오류를 수정하고, 이를 모델에 다시 반영할 수 있습니다.

맺으며

인과적 어텐션은 자기회귀 트랜스포머를 구성하는 핵심 기반이지만, 그렇다 해도 원래의 의도나 형태에 매몰될 필요는 없겠죠.

CASTLE의 미래 예측 키(Lookahead Keys)부터 VLM의 미래 인식 마스크(Future-aware Masks)까지, 모델의 전역적인 맥락을 확장하면서 구조를 유지하는 창의적인 방법들이 등장하고 있습니다. 모델이 미래를 살짝 엿보게 하면 추론, 이해, 다양한 작업에서의 성능이 향상된다는 것도 증명되고 있구요.

우리가 알고 있는, 일반적으로 받아들여지는 규칙을 신중하게 완화함으로써, 연구자들은 모델이 원래 강력하게 보여주었던 신뢰성을 잃지 않으면서 더 주어진 맥락에 민감하게, 심도있게 반응할 수 있다는 점을 발견했습니다. 결국, 이 어텐션은 모델과 함께 계속 진화하고 있는 셈이고, 이런 작은 변화가 언어 및 다중 모달 모델의 다음 돌파구로 이어질 수 있을지도 모르겠습니다.

더불어, 인과적 어텐션의 또 다른 측면인 인과-효과 추론(Cause-Effect Reasoning)을 잊지 마시기 바랍니다. 이 방향의 연구들은 모델이 무엇이 일어나고 있는지, 왜 그런지 인식하게 해서 마찬가지로 전역적인 시야를 넓힙니다.

정말, 어렵지만 매력적인 필드 아닌가요?

보너스: 참고자료

• Causal Attention with Lookahead Keys

• Rethinking Causal Mask Attention for Vision-Language Inference

• CAT: Causal Attention Tuning For Injecting Fine-grained Causal Knowledge into Large Language Models

• Learning to Focus: Causal Attention Distillation via Gradient-Guided Token Pruning

• AutoHFormer: Efficient Hierarchical Autoregressive Transformer for Time Series Prediction

• Transforming Causality: Transformer-Based Temporal Causal Discovery with Prior Knowledge Integration