Transformer 및 LLM의 중심인 attention 공식의 '오프 바이 원' 버그를 파고들며, AI 모델 경량화 및 정량화 문제와 그 개선 방안을 제시하는 의견 글.
By Evan Miller
2023년 7월 24일
말할 수 없는 것에 대해선 침묵해야 한다. –비트겐슈타인
이 수식에서 오프 바이 원 오류를 보신 적 있나요?
Attention(Q,K,V)=softmax(Q K^T / d)V
attention 공식은 현대 AI의 중심 공식입니다. 하지만 지난 한 주 동안 저를 괴롭힌 버그가 여기 숨어 있습니다. 저는 이 버그와 제 제안을 정식 논문으로 써보려 했지만, Pytorch와 biblatex와의 몇 차례 전투에서 패배하는 바람에 블로그 포스트로 대체하기로 했습니다. (역사는 승자가 쓰지만, 블로그는…)
이번 글에서는 (제가 보기엔) 현세대 AI 모델들이 결정적인 부분에서 오프 바이 원 오류를 갖고 있고, 이것 때문에 Transformer 모델들이 불필요하게 압축 및 배포가 어렵게 되고 있음을 설명하겠습니다. 일종의 의견 글이라 생각하시고, 만약 이 글을 읽는 분 중 실험을 통해 증명해줄 의향이 있다면 함께 협업해서 과학자라고 불러봅시다.
이 오프 바이 원 오류가 왜 문제인지부터 얘기해야겠죠. “ChatGPT 잘 돌아가는데 뭐가 문제야?” 저는 양자화(quantization)에 관한 논문을 훑던 중 뭔가 이상한 낌새를 처음 맡았습니다. LLM 최적화 고수들이 엄청 큰 모델을 Mac Mini, 라즈베리 파이, 해킹한 온도 조절기에 욱여넣을 때 쓰는 기법이죠. AI 연구자 모두 RAM 한계에 시달리는 터라, 더 적은 RAM으로 더 쿨한 걸 할수록 유리하기 때문입니다. LLM은 수십억 개의 파라미터를 갖고 있고, 이 가중치들을 최대한 효율적으로 줄이거나 압축하면 더 많은 걸 할 수 있습니다. 설령 그 목적이 더 좋은 산문을 쓰거나, 더 뛰어난 표절, 혹은 종말 가속이든 말이죠.
RAM은 정보를 저장합니다. (괜히 해설 같지만…) 정보란 음의 로그 확률입니다. 즉, 수열이 예측 가능할수록 필요한 비트가 적고, 아주 가끔 큰 수가 튀어나오는(예측 불가) 경우 더 많은 비트가 필요하게 됩니다.
이상치(outlier)가 LLM에서 계속해서 발견됩니다. Transformer 모델들은 엄청나게, 말 그대로 순서가 다를 정도로 큰 값을 가진 outlier 가중치와 Black Swan처럼 큰 활성값을 내뱉고 있습니다. 그런데 아무도 이걸 제거하지 못하고 있습니다! 이런 거대 outlier들은 모델 동작에 치명적인 역할을 하며, 기존 신경망 관련 상식과도 배치됩니다. 다양한 논문들이 이를 분석하고, 여러 가지 비트 저장 기법들이 제안됐지만, vanilla scale-and-bias 정수 양자화 방식으로는 상당한 성능 저하를 피할 수 없습니다. Georgi가 더 잘 설명해줄 수 있지만, 이 부분은 넘어가죠.
이 현상에 가장 적확한 분석은 Qualcomm AI Research 팀의 논문 **Quantizable Transformers: Removing Outliers by Helping Attention Heads Do Nothing**에서 볼 수 있습니다. 그들은 이런 outlier의 원인을 attention의 softmax 함수로 추적했고, 그 innocent해 보이는 softmax가 엄청난 'kurtosis'의 주범임을 밝혔죠. 연구진은 아주 근접하게, 거의 정답 직전까지 이 오프 바이 원 버그에 다가갔습니다. (하지만 여름 휴가 중인지 이메일 답변이 없기에 이렇게 공개적으로 호소합니다.)
위 논문을 읽으실 분은 그들의 제안은 그냥 건너뛰세요. 냉정해 보여도 이유 있습니다. clipped softmax는 gradient가 0이라는 허점이 있고, gated attention은 실현 가능하지만 엄청난 파라미터 증설을 동반하는, 사실상 인덱스 증가의 실패에서 비롯된 문제를 해결하려는 제안이니까요. 여기서 간과된, 너무도 단순한 수정이 있습니다—아직 아무도 시도조차 안 했던.
이제 이 오류가 무엇인지, 그리고 왜 softmax가 attention에 결코 적합하지 않은 도구인지 살펴보겠습니다.
attention 메커니즘이 뭘 하려는지 알지 못하면 이 버그도 이해하기 어렵습니다. 대부분의 수치 버그는 잘못된 코드 구현에서 오는데, 이번 경우는 “코드는 문제 없다, 수식 자체가 문제다”라는 얘기이니, 수식의 출처와 목적을 아는 게 치유의 출발입니다.
저도 이걸 이해하려고 arXiv 논문 50편은 읽었습니다. 원래는 Udemy 강좌나 Karpathy의 유튜브를 보는 게 더 나았을지도 모르지만, 일단 input embedding(입력 임베딩)부터 설명하죠. 이는 문자열 내 단어를 부동소수점 벡터로 나타내는 방식입니다.
이 벡터는 매년 더 커져가고 있습니다. 예를 들어 Meta의 LLaMA 2 모델은 길이 3,204의 임베딩 벡터를 사용하며, half-precision float로 단어 하나당 6KB 이상이 필요합니다. 어휘가 3만~5만 개쯤 된다면 단 1개 단어 정보에 이만큼의 메모리가 쓰입니다.
만약 저처럼 메모리를 아끼는 C 프로그래머라면 “단어 하나 표현하는 데 6KB 쓸 거면, 2바이트도 아깝지 않냐?”고 의문가질 수 있습니다. 어휘가 2^16=65,384개 미만이라면, 16비트면 충분할 터인데 말이죠.
하지만, Transformer는 입력 벡터를 같은 크기의 출력 벡터로 "변환"합니다. 그리고 최종 결과 6KB짜리 벡터가 다음 토큰 예측에 필요한 모든 정보를 담아야 합니다. 각 Transformer 계층의 역할은 원래의 단어 벡터에 정보를 추가하는 것. 이게 바로 skip(residual) connection의 정체죠. attention으로 맥락 분석을 하며, 예를 들어 _pupil_이 학생을 의미하는지, 아니면 눈동자를 뜻하는지 구분할 단서를 추가합니다. 이 과정을 여러 번 반복하면 결국 영어 전체 맥락을 익히는 셈입니다.
마지막 단계에서는 이 출력 벡터를 직사각행렬과 곱하고, softmax에 넣어 결과를 다음 토큰 확률로 취급합니다. 이건 합리적인 설계지만, 누구도 이 확률을 절대적으로 신뢰하지 않죠. 대부분의 구현에서는 low probability가 과다 표현되는 softmax 특성을 숨기려 샘플링 메커니즘을 별도로 둡니다.
이 점은 나쁘지 않습니다. softmax를 출력에 쓰면 어휘 내 모든 단어에 대해서 기울기가 들어오니, 차선책이 없는 한 합리적 선택입니다.
하지만 제가 주장하고 싶은 건 gander(거위) 소스와 goose(기러기) 소스는 달라야 한다는 점입니다. Transformer의 output softmax와 attention 내부 softmax는 목적 자체가 다릅니다. 우리는 후자를 과감히 폐기하거나, 최소한 분모에 무언가를 추가해야 합니다. ⛱️
softmax는 통계물리학에서 에너지 분포 예측에 처음 쓰였습니다:
p_i ∝ exp(-ε_i / kT)
경제학자들이 이걸 가져와서, 사람의 효용함수 상 잡음항이 Gumbel 분포를 따른다면, 어떤 선택을 할 확률도 input utility의 지수에 비례한다고 설명했습니다:
Pr(Y_i = k) = exp(X_i β_k) / ∑_j exp(X_i β_j)
softmax는 다항 로지스틱 회귀 등 다양한 모델로 진화했으며, softmax의 핵심은 연속 수치를 확률 분포로 바꾸는 "꼼수"입니다. 물리학에선 잘 맞고, 경제학에선 약간 왜곡몫이 있지만, 수학적으로 discrete choice가 있을 때 잘 동작합니다. softmax(x)_i = exp(x_i) / ∑_j exp(x_j).
softmax의 본질: _경쟁 대안 중 하나를 반드시 선택_하도록 강요합니다. 입자든, 소비자든, 뭔가 하나는 꼭 뽑히게 만들죠. 즉, softmax가 "아무것도 선택하지 않음"이 가능해야 할 때 강제로 분배를 해버립니다. 그래서 실제 데이터에선 왜곡이 생길 수 있죠.
LLM에선 그 왜곡이 문장부호/공백 같은 비의미 토큰을 과도하게 강조하면서, 압축이 어려운 outlier activation을 유발합니다. Qualcomm 연구팀은 LLM outlier의 97% 이상이 공백/구두점 위치에서 나온다고 밝혔습니다. 뭔가 수상하죠… (난 닭을 시켰는데 어째서 생선 냄새지!)
attention 내부에서 softmax가 어떻게 쓰이는지 자세히 들여다보죠:
Attention(Q,K,V)=softmax(Q K^T / d)V
디코더 전용 모델(즉, ChatGPT 이후 모든 모델)에서 Q,K,V 모두 같은 입력 시퀀스에서 시작합니다. 완전히 같진 않지만, 각 계층마다 동일한 주석 추가 임베딩에서 출발합니다.
Q K^T는 서로 다른 위치 토큰 벡터들 간 상관 관계를 보는 단계입니다. 각 행(토큰)의 확률 분포(softmax)가 V 행렬의 값 벡터에 섞여서, 입력 임베딩에 더해지고, 신경망 아래로 전달되죠.
multi-head attention에선 임베딩을 여러 부분(head)로 분할해서, 각 head가 전체 임베딩 정보를 한 segment에만 보강합니다. 원 Transformer 논문의 Concatenation 연산이 바로 이것입니다: Head 1이 Segment 1에, Head 2가 Segment 2에, 이런 식으로 정보를 더하는 구조.
문제는 softmax가 각 attention head로 하여금 반드시 주석(정보 추가)을 하게 강제한다는 점입니다. discrete 대안에서 softmax를 쓰는 건 좋지만, "주석은 옵션이다(입력에 덧셈)" 같은 상황에는 softmax는 적절치 않습니다. multi-head 구조에서는 각 head가 pass 하고 싶어도 그럴 방법이 없습니다. 안 그래도 노이즈가 많은데, abstain(기권)이 불가한, 소음 넘치는 민주주의인 꼴이죠.
softmax를 아예 버려야 한다고 주장할 순 있지만, 대체로 잘 돌아가니 간단한 고칠 점만 나왔으면 합니다. 이 버그는 attention이 등장할 때부터 존재한 문제죠.
준비되셨나요?
정말 준비되셨나요?
잘 안 들리네요.
드디어 공개합니다. 모든 LLM 해커를 들끓게 할 신형 Softmax 슈퍼모드:
(softmax₁(x))_i = exp(x_i) / (1 + ∑_j exp(x_j))
좀 허무하죠? 분모에 1만 더했을 뿐입니다. 이제 전체 벡터가 원한다면 0에 가까워질 수 있습니다. 대개는 아주 약간 shrink하는 효과지만, normalization 덕분에 그 정도는 이후에 무난히 보상됩니다.
핵심 차이는 음의 극한입니다. x 원소가 크게 음수(정보 보강을 완전히 피하려는 상황)일 때의 극한을 보세요.
기존 softmax: lim_{x₁→-∞,…,x_k→-∞} softmax(x)_i = 1/k > 0
새 softmax₁: lim_{x₁→-∞,…,x_k→-∞} softmax₁(x)_i = 0
vanilla softmax는 항상 동일 가중치 합을 내보냅니다. softmax₁도 거의 비슷하지만 음수 직교 방향으로 탈출구가 생깁니다. _이건 수치적 문제가 아니고, 수학적 오류_입니다. 정밀도를 아무리 높여도 softmax의 구조적 문제라, 모든 Transformer가 영향 받습니다.
softmax₁의 다른 특징: 도함수는 항상 양수(즉, 항상 gradient가 들어온다), 출력 합은 0~1 사이, 즉 값 폭주 없음. 속성은 유지됩니다.
(softmax₁(x))_i/(softmax₁(x))_j = (softmax(x))_i/(softmax(x))_j = exp(x_i)/exp(x_j) ∀i,j
즉, 출력 벡터의 상대적 값은 변동 없음.
원래 이 함수 이름을 ghostmax라 부르고 싶었지만(벡터 x에 0값 추가라고 볼 수 있고, exp(0)=1이니까), 너무 무서울까봐 안 지었습니다.
기능적으로 매우 지루해 보여도, quantization에 수많은 연구가 몰리는 outlier 피드백 문제의 근본 해결책이라고 99.44% 확신합니다. 같이 실험해보고 싶은 분은 트위터로 DM주세요.
이제부터 개선된 attention 메커니즘은 _QuietAttention_이라고 부르겠습니다. attention head가 "아무 말도 안 하기"가 가능한 구조입니다.
QuietAttention(Q,K,V):=softmax₁(Q K^T / d)V
아주 간단하게 실험할 방법도 있을 듯합니다: 매 input context 맨 앞에 zero vector를 prefix로 붙이고, 신경망이 bias(포지셔널 인코딩 포함)를 더하지 않게 하면, zero가 그대로 통과해 모든 이후 softmax 분모에 1이 더해질 것입니다. gradient 코드 건드릴 필요 없이 sanity check가 되죠. LLaMA 모델에서 fixed embedding + 특수 prefix 토큰으로 실험이 가능할 듯한데, 제 취미 생활을 방해하고 있으니 여기까지만 파보고 접겠습니다.
모델 재학습은 필수이니, 아직 라즈베리파이에서 바로 쓸 수는 없습니다. 다만, weight kurtosis와 activation infinity norm 등이 실험 후 어떻게 변하는지 꼭 알려 부탁드립니다. 실험표 하나 멋지게 만들어, 곧 arXiv 논문으로 변신시켜 보렵니다. Qualcomm 연구진이 여행에서 돌아오든, LLM 해커가 biblatex를 정복하든, 누가 먼저든요.
지금 여러분은 evanmiller.org의 수학, 기술, 단상 모음글을 읽으셨습니다. 이 글이 마음에 든다면 다음 글도 추천합니다:
새로운 글을 받아보고 싶으신가요? LinkedIn, Twitter, RSS로 구독하세요.
MySQL, PostgreSQL, SQLite 통계 패턴을 찾고 싶으신가요? 데스크탑 통계 소프트웨어 **Wizard**로 더 많은 데이터를 더 빠르게 시각화 및 분석하세요. 복잡한 명령어 없이 쉽고 빠르게 결과를 얻을 수 있습니다!
Mac 방식의 통계 소프트웨어
에반 밀러의 홈으로 돌아가기 – RSS 구독 – LinkedIn – Twitter
