🤖 AI

Mask-Based Priors Are More Persistent than Query-Key Initializations

이 논문은 가산적 어텐션 마스크(additive attention masks)를 태스크 수준의 상호작용 구조로 직접 초기화하는 것이 쿼리-키 초기화보다 더 지속적이고 효과적인 귀납적 편향을 제공하며, 이를 통해 트랜스포머가 불리언 추론 작업에서 완벽에 가까운 외삽 성능을 달 수 있고 저데이터 산술 시나리오에서 개선된 성능을 달성할 수 있음을 입증한다.

원저자: Mingze Ma, Hemanth Saratchandran, Cameron Gordon, Simon Lucey

게시일 2026-08-04

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Mingze Ma, Hemanth Saratchandran, Cameron Gordon, Simon Lucey

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

당신이 로봇에게 퍼즐을 푸는 법을 가르치고 있다고 상상해 보세요. 당신은 로봇에게 게임을 플레이하는 백만 개의 사례를 보여주고, 로봇은 그 게임에 매우 능숙해집니다. 하지만 그러고 나서 당신이 로봇에게 한 번도 본 적 없는 약간 새로운 버전의 게임을 수행하라고 요청합니다. 그러면 로봇은 새로운 규칙을 파악하는 대신, '지름길'을 찾아냅니다. 즉, 이전의 예시들에서는 통했지만 새로운 상황에서는 완전히 틀린 단순한 요령을 사용하는 것입니다. 이것은 인공지능 세계, 특히 **트랜스포머(Transformer)**라고 불리는 일종의 뇌를 닮은 컴퓨터 모델에서 발생하는 유명한 문제입니다. 트랜스포머는 이야기를 쓰고, 수학 문제를 풀고, 얼굴을 인식하는 많은 현대적 AI 도구들의 엔진입니다. 과학자들은 이 로봇들이 본 것을 기억하는 데는 천재적이지만, 지식을 새로운 상황에 적용하는 '외삽(extrapolate)'에는 자주 실패한다는 사실을 발견했습니다. 왜냐하면 로봇이 진정한 정답보다는 가장 단순한 답에 매몰되기 때문입니다.

연구자들이 던지는 핵심적인 질문은 이것입니다: 어떻게 하면 로봇이 지름길로 빠지는 것을 막을 수 있을까? 한 가지 아이디어는 로봇이 학습을 시작하기도 전에 규칙을 가르쳐서 '선행 지식'을 주는 것입니다. 트랜스포머의 세계에서 이것은 **귀납적 편향(inductive bias)**이라고 불립니다. 이것은 탐정에게 사건을 해결하기 전에 지도를 건네주는 것과 같습니다. 만약 지도가 정확하다면 탐정은 범인을 빠르게 찾아낼 것입니다. 하지만 만약 지도가 틀렸거나 탐정이 지도를 무시한다면, 탐정은 길을 잃게 될 것입니다. 오랫동안 과학자들은 로봇의 내부 '센서'(쿼리 및 키 투영, Query and Key projections라고 불림)를 미세하게 조정하여 로봇이 자연스럽게 올바른 단서를 찾도록 지도를 그려주려 노력했습니다. 하지만 결과적으로 로봇은 데이터를 배우려는 의지가 너무 강해서, 당신이 준 지도를 빠르게 잊어버리고 다시 지름길을 찾는 습관으로 돌아가 버린다는 것이 밝혀졌습니다.

"Mask-Based Priors Are More Persistent than Query-Key Initializations"라는 제목의 이 논문은 바로 이 문제에 대해 깊이 파고듭니다. 호주 머신러닝 연구소(Australian Institute for Machine Learning)의 밍제 마(Mingze Ma)와 그의 팀은 단순하지만 심오한 질문을 던졌습니다. "로봇이 결코 잊을 수 없는 지도를 주는 방법이 있을까?" 그들은 지도를 그리는 두 가지 방법을 테스트했습니다. 첫 번째 방법은 기존 방식인 센서 미세 조정이었습니다. 두 번째 방법은 더 단순하고 새로운 접근법으로, 특수한 **마스크(mask)**를 사용하여 로봇의 '주의력(attention)' 메커니즘에 직접 지도를 그려 넣는 것이었습니다.

그들이 발견한 결과는 다음과 같습니다. 기존 방식(센서 미세 조정)을 사용했을 때, 로봇은 처음에 올바른 지도를 가지고 시작했지만, 학습이 시작되자마자 로봇은 지도를 지워버리고 다시 잘못된 지름길로 돌아갔습니다. 그것은 마치 화이트보드에 글을 써 놓으면 로봇이 즉시 지워버리는 것과 같았습니다. 그러나 새로운 마스크 기반 방식을 사용했을 때는 지도가 그대로 유지되었습니다. 그들은 센서를 변경하지 않았습니다. 대신, 주의력 점수(attention scores)에 직접적으로 학습 가능한 편향, 즉 '넛지(nudge)'를 추가했습니다. 이 넛지는 매우 끈질겨서, 수천 번의 학습 단계가 지난 후에도 로봇은 여전히 규칙을 기억하고 있었습니다.

결과는 극적이었습니다. 로봇이 보통 실패하는 까다로운 논리 퍼즐(불리언 태스크, Boolean tasks)에서, 이 새로운 방식은 로봇이 보지 못한 문제에서도 거의 100%의 정답률을 기록한 반면, 기존 방식들은 75%에 머물렀습니다. 이는 마치 로봇이 단순히 요령 피우는 암기자가 아니라, 마침내 인간처럼 생각하는 법을 배운 것과 같습니다. 이 성공은 단지 논리 퍼즐에만 국한되지 않았습니다. 연구팀은 또한 수학 문제(큰 숫자의 덧셈과 곱셈 등)와 이미지 인식에 대해서도 테스트했습니다. 모든 경우에서 로봇은 더 빠르게 학습했고, 특히 공부할 데이터가 많지 않을 때 실수를 훨씬 적게 했습니다.

이 논문은 비결이 단순히 로봇에게 '무엇을' 가르치느냐가 아니라, '어디에' 가르치느냐에 있다는 점을 시사합니다. 구조적 규칙을 복잡한 센서 계산 속에 숨기는 대신 주의력 마스크에 직접 배치함으로써, 로봇은 그 규칙을 마음속에 간직하게 됩니다. 저자들은 이 단순한 기술이 마치 튼튼한 닻처럼 작용하여, 새로운 데이터의 파도가 로봇을 경로에서 벗어나도록 밀어내려 할 때도 학습의 방향을 올바르게 잡아준다는 것을 보여줍니다. 비록 이것이 모든 AI 문제를 해결해주지는 못하며(연구진은 지도를 그리기 위해 사전에 규칙을 알고 있어야 한다는 점을 인정합니다), 경험으로부터 단순히 암기하는 것이 아니라 경험으로부터 실제로 배울 수 있는, 더 똑똑하고 신뢰할 수 있는 AI를 구축하기 위한 강력한 새로운 도구를 제공합니다.

기술 요약: 마스크 기반 사전 지식(Priors)이 쿼리-키 초기화보다 더 지속적이다

문제 정의

트랜스포머는 특히 불리언(Boolean) 외삽 태스크에서 미관측 데이터(GOTU, Generalization On the Unseen)에 대한 체계적인 실패 모드를 보입니다. 이러한 모델들은 관찰된 훈련 데이터에는 완벽하게 적합할 수 있지만, 실제 고차 타겟 함수보다는 훈련 도메인에는 부합하지만 미관측 입력에서는 실패하는 더 단순하고 낮은 차수의 함수인 "최소 차수 보간기(Minimum-Degree Interpolator, MDI)"에 따라 외삽하는 경tendency가 있습니다. 이는 트랜스포머가 저차수 구조적 솔루션을 선호한다는 특정 귀납적 편향(Inductive Bias)을 나타냅니다.

이를 교정하기 위해 쿼리-키(QK) 초기화(예: $W_Q$ 와 $W_K$ 를 설정하여 $QK^\top$ 가 원하는 어텐션 패턴을 근사하도록 함)를 통해 구조적 사전 지식을 주입하려는 기존의 시도들은 불충분한 것으로 판명되었습니다. 본 논문은 QK 기반 방식이 초기화 단계에서는 원하는 어텐션 패턴을 생성할 수 있지만, 이 구조는 간접적이며 콘텐츠 의존적(content-dependent)이라는 점을 지적합니다. 결과적으로, 주입된 사전 지식은 모델이 태스크 의존적 표현을 학습하기 위해 쿼리 및 키 프로젝션을 재목적화함에 따라 빠르게 덮어씌워지거나 "씻겨 내려가며(washed out)", 최종적인 외삽 규칙을 변경하는 데 실패합니다.

방법론

저자들은 **마스크 기반 구조적 초기화(Mask-Based Structured Initialization)**를 더 지속적인 대안으로 제안합니다. 이 방법은 사전 지식을 콘텐츠 의존적인 쿼리 및 키 프로젝션의 유사도 점수를 통해 간접적으로 인코딩하는 대신, 이를 유한하고 학습 가능한 가산 바이어스(additive bias)로서 어텐션 로짓(attention logits)에 직접 주입합니다.

핵심 메커니즘

표준 셀프 어텐션 메커니즘은 다음과 같이 수정됩니다:
$\text{Attention}(X) = \text{softmax}\left( \frac{XW_Q W_K^\top X^\top}{\sqrt{d}} + \log \sigma(M) \right) XW_V$
여기서:

$M \in \mathbb{R}^{N \times N}$ 은 태스크 수준의 상호작용 구조로부터 초기화된 학습 가능한 마스크 행렬입니다.
$\sigma(\cdot)$ 는 시그모이드 함수로, $\log \sigma(M) \leq 0$ 임을 보장합니다.
$\log \sigma(M)$ 항은 원하는 상호작용은 유지하면서 원치 않는 상호작용을 선택적으로 억제하는 가산 바이어스로 작용합니다.

설계 원칙

마스크는 타겟 태스크의 기능적 분해(functional decomposition)를 기반으로 초기화됩니다:

태스크 정렬 상호작용(Task-Aligned Interaction): 타겟 규칙에 따라 상호작용해야 하는 토큰들로 어텐션을 제한합니다 (예: 패리티를 위한 쌍별 상호작용, 덧셈을 위한 자릿수 정렬).
지름길 억제(Suppression of Shortcuts): MDI 솔루션으로 이어지는 상호작용을 억제합니다 (예: 다수결 투표에서의 전역 합계 지름길).
학습 가능성(Learnability): 하드한 인과적(causal) 또는 패딩 마스크와 달리, 이 마스크들은 유한하고 학습 가능하여 모델이 핵심 구조적 제약을 유지하면서 최적화 과정 중에 사전 지식을 정교화할 수 있도록 합니다.

이론적 정당화

본 논문은 마스크 기반 사전 지식의 지속성을 설명하기 위해 그래디언트 수준의 분석(Theorem 1 및 Corollary 1)을 제공합니다:

QK 초기화: $Q$ 와 $K$ 에 대한 그래디언트는 노름 $\|q\|$ 및 $\|k\|$ 에 비례합니다. 실제로 큰 피처 노름은 초기화된 구조를 빠르게 덮어쓸 수 있는 큰 업데이트를 유도할 수 있습니다.
마스크 초기화: 마스크 파라미터 $m_j$ 에 대한 그래디언트는 $p_i \sigma(-m_j) \leq p_i$ 에 의해 유계(bounded)됩니다. 마스크는 콘텐츠 의존적인 피처 스케일로부터 분리되어 시그모이드 인자에 의해 조절되므로, 더 느리고 안정적으로 업데이트되어 최적화 과정 내내 구조적 편향을 보존합니다.

주요 기여

사전 지식 지속성 실패 식별: QK 기반 구조적 초기화가 초기화 시에는 효과적이지만, 불리언 GOTU 태스크에서 훈련될 때 주입된 사전 지식을 유지하지 못하고 모델을 기본 MDI 편향에 갇히게 만든다는 것을 입증했습니다.
마스크 기반 초기화 제안: 구조적 제약을 콘텐츠 의존적인 유사도 점점으로부터 분리하여, 가산 어텐션 로짓 바이어스로서 구조적 사전 지식을 직접 주입하는 새로운 방법을 제시했습니다.
그래디언트 수준의 설명: 가산 마스크가 피처 스케일로부터 구조적 편향을 어떻게 디커플링(decouple)하여 QK 기반 방식보다 더 안정적인 업데이트와 더 높은 지속성을 이끌어내는지에 대한 이론적 증명을 제공합니다.
실험적 검증: 불리언 추론, 산술, 비전 벤치마크에 대한 포괄적인 평가를 통해 마스크 기반 초기화가 MDI 솔루션으로부터 외삽 동작을 성공적으로 전환함을 보여줍니다.

실험 결과

불리언 외삽 (GOTU)

트랜스포머가 일반적으로 외삽에 실패하는 2Parity, 3-Bit Majority, Cyclic3와 같은 태스크에서:

Vanilla 및 QK 초기화 모델: MDI 편향에서 벗어나지 못하며, 2Parity에서 약 75%, 3-Bit Majority에서 약 87%의 정확도를 유지합니다.
마스크 초기화 모델: 거의 완벽한 외삽을 달者합니다 (2Parity 및 3-Bit Majority에서 100% 정확도, Cyclic3에서 99.98%). 손실값 또한 크게 감소합니다 (예: 2Parity에서 1.379에서 0.0029로 감소).

산술 태스크

낮은 데이터 비율에서의 3자리 덧셈 및 곱셈 태스크에서:

마스크 초기화는 Vanilla 및 QK 기반 방식(SVD 및 Opt)보다 현저히 우수한 성능을 보입니다.
단 0.4%의 훈련 데이터만 사용한 덧셈 태스크에서, Mask Init은 기본 모델의 1.19%와 비교하여 **77.49%**의 시퀀스 정확도를 달성했습니다.
3%의 데이터로 수행된 곱셈 태스크에서, Mask Init은 5.01%의 기본 모델 대비 **69.49%**의 시퀀스 정확도에 도달했습니다.

비전 벤치마크

ViT-Tiny를 사용한 CIFAR10, CIFAR100, ImageNet100에서:

마스크 초기화는 기본 초기화 및 기존의 구조적 방식(Mimetic 및 Impulse)보다 일관되게 높은 Top-1 정확도를 달성했습니다.
예를 들어, CIFAR10에서 Mask Init은 94.20%(Default) 및 95.12%(Impulse) 대비 **95.66%**를 달성했습니다.

NLP 태스크

명시적인 구조적 사전 지식이 덜 정의된 WikiText-2 및 TinyStories에서:

마스크 초기화(인과적 상대 거리 사전 지식 사용)는 퍼플렉시티(perplexity)와 수렴 속도에서 완만하지만 일관된 개선을 보였으며, 이는 "진정한" 사전 지식이 엄격히 알려지지 않은 경우에도 적응성을 보여줍니다.

의의 및 주장

본 논문은 **"사전 지식이 어디에 존재하는가가 외삽을 변화시킬지를 결정한다"**고 주장합니다. 구조적 사전 지식을 콘텐츠 의존적인 $QK^\top$ 프로젝션에서 가산 어텐션 로짓으로 이동시킴으로써, 저자들은 귀납적 편향을 지속 가능하게 만들 수 있음을 보여주었습니다.

이 연구의 의의는 어텐션 마스크가 단순히 인과성이나 국소성을 위한 건축적 제약이 아니라, 지속적인 귀납적 편향을 인코딩하기 위한 단순한 기질(substrate)로 기능할 수 있음을 입증했다는 데 있습니다. 이 접근 방식은 아키텍처를 수정하거나 새로운 연산을 도입하지 않고도 트랜스포머가 논리적 및 조합적 태스크에서 기본 성향인 최소 차수 보간기로 향하는 경향을 극복할 수 있게 합니다. 저자들은 이 방법이 의미 있는 상호작용 사전 지식에 대한 접근(즉, "오라클" 가정)을 전제로 하지만, 이러한 사전 지식이 마스크를 통해 올바르게 인코딩될 때 트랜스포머의 학습 역학 및 일반화 능력을 근본적으로 변화시킬 수 있음을 증명한다고 언급합니다.