💬 NLP

Cubit: Token Mixer with Kernel Ridge Regression

본 논문은 더 강력한 수학적 기반을 제공하고 장기 시퀀스 모델링 능력을 향상시키는 것을 목표로 Transformer 의 어텐션 메커니즘을 커널 릿지 회귀로 대체하는 새로운 딥러닝 아키텍처인 Cubit 을 소개합니다.

원저자: Chuanyang Zheng, Jiankai Sun, Yihang Gao, Yuehao Wang, Liangchen Tan, Mac Schwager, Anderson Schneider, Yuriy Nevmyvaka, Xiaodong Liu

게시일 2026-05-08

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Chuanyang Zheng, Jiankai Sun, Yihang Gao, Yuehao Wang, Liangchen Tan, Mac Schwager, Anderson Schneider, Yuriy Nevmyvaka, Xiaodong Liu

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

이야기를 쓰려고 할 때, 다음에 어떤 단어가 올지 결정해야 한다고 상상해 보세요. 이를 위해 이미 쓴 모든 단어를 뒤로 돌아보며 확인합니다. 현재 이 작업을 수행하는 표준 방식 (Transformer 라고 함) 은 마치 매우 정돈된 사서처럼 과거 단어 전체 목록을 읽고 각 단어에 '관련성 점수'를 부여한 뒤, 그 점수들을 바탕으로 요약본을 만드는 것과 같습니다.

이 논문의 저자들인 Cubit은 이 사서가 Nadaraya-Watson 회귀라는 특정 유형의 수학을 수행하고 있다고 주장합니다. 이는 마치 "이 새로운 단어는 이전 단어들과 얼마나 유사한가?"라고 묻고, 그 유사도에 기반해 이전 단어들을 평균내는 것과 비슷합니다. 이는 잘 작동하지만, 저자들은 이 수학을 수행하는 더 나은 그리고 더 견고한 방법이 있다고 생각합니다.

다음은 그들이 제안하는 바에 대한 간단한 설명입니다:

1. 문제: "사서"는 훌륭하지만 완벽하지는 않습니다

현재의 Transformer 아키텍처는 유사한 단어를 찾는 데는 뛰어나지만, 때로는 노이즈에 혼란을 겪거나 이야기가 매우 길어지면 막히는 사서와 같습니다. 그 뒤의 수학은 오류나 경계 (문장의 시작이나 끝과 같은) 를 처리하는 데 있어 다소 "느슨한" 면이 있습니다.

2. 해결책: Cubit ("스마트 회계사")

저자들은 Cubit이라는 새로운 아키텍처를 제안합니다. 단순히 "이 단어들은 얼마나 유사한가?"라고 묻는 대신, Cubit 은 **Kernel Ridge Regression(KRR)**이라는 다른 수학적 도구를 사용합니다.

비유: 만약 기존 방법이 단순히 평균을 내는 사서라면, Cubit 은 유사성만 보는 것이 아니라 복잡한 방정식을 풀어 최적의 답을 찾는 스마트 회계사와 같습니다.
"Ridge" 부분: 수학에서 "Ridge"는 안전망과 같습니다. 이는 회계사가 하나의 이상한 숫자 (노이즈) 에 너무 흥분하지 않도록 막아주며, 최종 답이 안정적이고 균형 잡히도록 보장합니다. 이 논문은 이것이 모델을 수학적으로 더 견고하게 만들고, 이야기가 복잡해졌을 때 실수를 할 가능성을 줄여준다고 주장합니다.

3. 비법: Limited-Range Rescale(LRR)

저자들은 이 새로운 수학을 시도했을 때, 숫자가 너무 커지거나 너무 작아져 모델이 충돌하거나 학습이 느려질 수 있음을 발견했습니다.

비유: 스테레오의 볼륨 노브를 돌리는 상황을 상상해 보세요. 볼륨을 최대까지 올리면 스피커가 터질 수 있고, 최소까지 내리면 아무것도 들리지 않습니다.
해결책: Cubit 은 **Limited-Range Rescale(LRR)**이라는 기능을 도입합니다. 이는 볼륨 노브에 "리미터"를 설치하는 것과 같습니다. 이는 볼륨이 안전하고 통제된 범위 (낮은 한계와 높은 한계 사이) 에 머무르도록 강제합니다. 이로써 모델은 안정성을 유지하면서 깨지지 않고 더 빠르게 학습할 수 있습니다.

4. 테스트 결과는 어떨까요?

저자들은 그들의 "스마트 회계사"(Cubit) 가 "사서"(Transformer) 와 또 다른 경쟁자인 DeltaFormer 보다 더 나은지 확인하기 위해 실험을 수행했습니다.

긴 이야기: 가장 흥미로운 결과는 이야기 (시퀀스) 가 길어질수록 Cubit 이 다른 모델들에 비해 현저히 더 좋아진다는 것입니다. 기존 Transformer 가 8,000 단어 전의 내용을 기억하는 데 어려움을 겪는 동안, Cubit 은 훨씬 더 긴 컨텍스트를 효과적으로 처리하는 것으로 보입니다.
더 큰 모델: 모델을 더 크게 만들었을 때 (더 많은 "뇌력"), Cubit 은 계속 성능이 향상된 반면, 다른 모델들은 정체되거나 우위를 잃기 시작했습니다.
다른 작업: 과학 논문, 책, 인터넷 등 어떤 자료로 학습하든 Cubit 은 일관되게 더 나은 결과 (더 낮은 "손실", 즉 더 적은 실수를 의미하는 세련된 표현) 를 산출했습니다.

요약

이 논문은 오래된 수학 (Nadaraya-Watson) 을 더 견고한 수학 기법 (Kernel Ridge Regression) 으로 교체하고 볼륨 조절 메커니즘 (LRR) 을 추가함으로써 Cubit이라는 새로운 유형의 AI 두뇌를 구축했다고 주장합니다.

핵심 결론: Cubit 은 정보를 혼합하는 더 수학적으로 안정적인 방법입니다. 이는 유사성에 기반해 단순히 추측하는 것이 아니라, 안전망을 갖춘 채 최선의 답을 계산하며, 특히 매우 긴 텍스트 시퀀스를 다룰 때 그 빛을 발합니다.

기술 요약: 커널 릿지 회귀를 활용한 토큰 믹서 Cubit

문제 제기

2017 년 트랜스포머 아키텍처가 도입된 이후, 자연어 처리, 컴퓨터 비전, 멀티모달 작업을 위한 딥러닝의 지배적 패러다임이 되었습니다. 그러나 트랜스포머의 핵심인 토큰 믹싱 메커니즘은 자기 주의 (self-attention) 에 의존하는데, 이는 이론적으로 지수 커널 (소프트맥스를 통해) 과 L1 정규화를 수행하는 Nadaraya–Watson (NW) 회귀로 해석될 수 있습니다.

효과적이기는 하지만, 이 NW 기반 접근법은 고유한 한계를 지닙니다:

계산 복잡도: 표준 주의 메커니즘은 시퀀스 길이에 대해 이차적인 계산 비용 ( $O(N^2)$ ) 을 발생시켜 긴 시퀀스의 확장성을 제한합니다.
이론적 제약: NW 회귀는 편향 - 분산 트레이드오프를 위한 명시적 정규화가 부재하며, 노이즈와 경계 효과에 민감할 수 있습니다.
개선의 확장성: 희소 주의, 선형 주의, 전문가 혼합 (mixture-of-experts) 등 수많은 수정안이 제안되었으나, 근본적인 계산 패러다임은 여전히 NW 회귀에 뿌리를 두고 있습니다.

저자들은 NW 회귀 기반을 보다 견고한 통계적 프레임워크로 대체함으로써 수학적 기반이 더 튼튼하고 긴 시퀀스 모델링 능력이 향상된 아키텍처를 도출할 수 있다고 주장합니다.

방법론: Cubit 아키텍처

본 논문은 주의 모듈의 NW 회귀를 **커널 릿지 회귀 (KRR)**로 대체하는 새로운 아키텍처인 Cubit을 제안합니다.

1. 이론적 기반: NW 에서 KRR 로

Nadaraya–Watson (트랜스포머): 표준 주의에서 출력 토큰 $z_i$ 는 쿼리와 키 간의 유사도를 소프트맥스로 정규화하여 결정된 가중치로 가치 벡터 $v_j$ 의 가중합입니다. 이는 전역 시스템을 풀지 않고 지역 가중 평균에 기반하여 함수 값을 추정하는 NW 회귀에 해당합니다.
커널 릿지 회귀 (Cubit): KRR 은 정규화된 경험적 위험 최소화 문제를 풉니다. 새로운 입력에 대한 예측은 커널 행렬의 역행렬에 정규화 항을 더한 $(K + \lambda I)^{-1}$ $(K + λ I)^{- 1}$ 을 포함합니다.
- 주요 차이점: NW 와 달리 KRR 은 편향 - 분산 트레이드오프를 제어하기 위한 명시적 정규화 매개변수 $\lambda$ 를 도입하며, 그람 행렬의 역행렬을 활용하여 모든 데이터 샘플 간의 정보를 결합함으로써 노이즈와 경계 효과에 대한 더 큰 견고성을 제공합니다.

2. 아키텍처 설계

Cubit 은 KRR 의 폐형 해 (closed-form solution) 를 통합하여 고전적인 주의 계산을 수정합니다:

커널 구성: 커널 행렬은 키 임베딩에서 파생된 특징 맵을 사용하여 구성됩니다. 저자들은 학습 가능한 투영 행렬 $W_K$ 와 지정된 활성화 함수 $\phi$ 를 사용하여 특징 맵 $\psi(x)$ 를 매개변수화합니다.
역행렬을 통한 정규화: 행별 소프트맥스 정규화 대신, Cubit 은 정규화 행렬 $\Sigma^{(h)} = (K^{(h)}K^{(h)\top} + \lambda I)^{-1}$ 을 계산합니다. 이 행렬은 전역 결합 연산자 역할을 합니다.
출력 공식화: 출력은 $Z^{(h)} = A^{(h)} \Sigma^{(h)} V^{(h)}$ 로 계산되며, 여기서 $A^{(h)}$ 는 커널 유사도 행렬 (쿼리 - 키 상호작용) 을 나타내고 $\Sigma^{(h)}$ 는 KRR 에서 영감을 받은 정규화 행렬입니다.
인과적 제약: 자기회귀 생성을 처리하기 위해 커널 구성을 적응시켜 $i$ 번째 토큰이 자신과 이전 토큰에만 주의를 기울이도록 하여, 회귀가 사용 가능한 문맥에서 수행되도록 보장합니다.

3. 제한 범위 재조정 (Limited-Range Rescale, LRR)

큰 행렬을 역전하는 과정에서 내재된 학습 안정성 문제를 해결하기 위해 저자들은 **제한 범위 재조정 (LRR)**을 도입합니다:

학습 가능한 스케일링 벡터 $s^{(h)}$ 를 적용하여 커널 정밀도를 적응적으로 조정합니다.
스케일링 인자는 시그모이드 함수를 사용하여 유계 범위 $(\beta, \alpha + \beta)$ 내에서 제한됩니다.
이는 스케일링 행렬의 가역성을 보장하고 역행렬 요소를 제한하여 선형 시스템 해결 중 수치적 불안정성을 방지합니다.

4. 관련 연구와의 차별성

본 논문은 DeltaFormer(역행렬을 활용하는 모델) 와 Cubit 을 명시적으로 구분합니다:

동기: DeltaFormer 는 DeltaNet 과 트랜스포머를 결합하는 반면, Cubit 은 NW 에서 KRR 로의 이론적 전환에 동기를 부여받습니다.
구현: DeltaFormer 는 역행렬을 구성하기 위해 두 가지 다른 임베딩을 사용하고 대각선 값을 1 으로 강제합니다. 반면 Cubit은 단일 특정 임베딩 (예: $K$ 만) 을 사용하여 대각선 값을 1 로 강제하지 않고도 자연스럽게 가역성을 보장합니다.
안정성: Cubit 은 DeltaFormer 에는 없는 구성 요소인 LRR 을 제안하여 안정성을 개선합니다.

주요 기여

통합 이론적 프레임워크: 저자들은 토큰 믹싱 메커니즘과 고전적 회귀 방법 간의 연결을 확립하고, NW 회귀와 KRR 을 체계적으로 분석하여 모델 아키텍처 개발을 위한 새로운 관점을 제공합니다.
Cubit 아키텍처: 트랜스포머의 NW 추정기를 대체하는 KRR 기반 토큰 믹서를 제안하며, 학습 안정성을 강화하기 위한 제한 범위 재조정 (LRR) 메커니즘을 특징으로 합니다.
실증적 검증: Cubit 이 표준 트랜스포머와 DeltaFormer 보다 더 강력한 긴 시퀀스 모델링 능력을 보이며, 시퀀스 길이가 길어질수록 성능 향상이 증가함을 보여주는 포괄적인 실험을 수행했습니다.

실험 결과

저자들은 다양한 데이터셋 (Arxiv, Books3, FinWeb-Edu), 시퀀스 길이 (512 에서 8192), 모델 크기 (1 억 2 천 5 백만에서 13 억 파라미터) 에 걸쳐 Cubit 을 트랜스포머 및 DeltaFormer 와 비교 평가했습니다.

언어 모델링 성능: Cubit 은 베이스라인보다 일관되게 낮은 퍼플렉시티 (검증 손실) 를 달성했습니다. 예를 들어, Arxiv 데이터셋 (1 억 2 천 5 백만 모델, 512 토큰) 에서 Cubit 은 1.8802 의 손실을 기록한 반면 트랜스포머는 1.9003 을 기록했습니다.
긴 시퀀스 모델링: 훈련 시퀀스 길이가 증가함에 따라 Cubit 과 트랜스포머 간의 성능 격차가 벌어졌습니다. 8192 토큰에서 손실 격차는 0.05~0.06 으로 안정화되었으며, 이는 Cubit 의 커널 기반 메커니즘이 장거리 의존성을 더 효과적으로 포착함을 시사합니다.
확장 특성: Cubit 은 유리한 확장 행동을 보였습니다. 모델 크기가 증가함에 따라 (3 억 5 천만에서 13 억으로) DeltaFormer 의 트랜스포머에 대한 우위가 감소한 반면, Cubit 의 성능 향상은 지속되어 개선되었습니다.
하류 작업: 제로샷 평가 (ARC, HellaSwag, PIQA 등) 에서 Cubit 은 다양한 데이터 크기 (100 억에서 500 억 토큰) 와 모델 규모에서 트랜스포머와 DeltaFormer 를 일관되게 능가했습니다.
애블레이션 연구:
- 아키텍처 설계: 성능 향상은 매개변수 확장이 아닌 구조적 설계 (KRR 공식화) 에 기인했습니다.
- LRR 효능: LRR 포함은 측정 가능한 개선을 제공했으며, "Cubit-NoLRR"도 표준 트랜스포머보다 우수했으나, LRR 을 포함한 "Cubit"이 최상의 결과를 달성했습니다.

의의 및 주장

본 논문은 Cubit 이 커널 릿지 회귀의 이론적 이점 (명시적 정규화, 재생 커널 힐베르트 공간 (RKHS) 에서의 더 빠른 수렴 속도, 노이즈에 대한 견고성 등) 을 활용함으로써 바닐라 트랜스포머보다 더 강력한 수학적 기반을 제공한다고 주장합니다.

저자들은 Cubit 이 특히 긴 문맥 이해가 필요한 작업을 위한 차세대 기반 모델 아키텍처를 대표할 수 있다고 제안합니다. 성능 향상은 단순히 훈련 토큰 증가 때문이 아니라 긴 시퀀스에서의 아키텍처 효율성에서 비롯된다고 지적합니다. 논문은 더 강력한 모델의 남용 가능성이 있음을 인정하면서도, 향상된 모델링 능력의 잠재적 사회적 이점을 강조합니다.

이 연구는 경계 편향을 해결하지만 계산 속도가 느린 국소 선형 회귀 (Local Linear Regression) 및 기타 토큰 믹싱을 위한 회귀 기반 프레임워크 탐구와 같은 향후 방향을 암시하며, 주의에 대한 회귀 기반 관점이 추가적인 아키텍처 혁신을 위한 비옥한 토대임을 나타냅니다.