하나의 QK 채널, 다양한 소스: 저정밀 Attention 모델의 급격한 성능 저하 방지

One QK Channel, Many Sources: Guarding Low-Precision Attention Collapse

bfloat16 트랜스포머는 많은 단계 동안 정상적으로 학습될 수 있지만, 갑자기 성능이 급락하는 현상이 발생할 수 있습니다. 서로 다른 종류의 저정밀 오류가 동일한 실패를 유발할 수 있으며, 각 소스에 대한 개별적인 수정이 필요한지 또는 단일 공유 경로를 차단하는 것이 더 효과적인지에 대한 불확실성이 존재합니다. 본 연구에서는 GPT-2 수준의 성능 저하 현상이 스트리밍-소프트맥스 누적기에 의해 유발되는 것을 확인하고, fp32 누적이 이를 복구할 수 있다는 사실을 밝혀냈습니다. 또한, 이 오류를 이용하여 다양한 소스로부터 제어된 오류를 전파하는 실험을 수행했습니다. Attention 외부에 위치한 오류도 동일한 query-key (QK) 스펙트럼의 폭주 현상을 유발하며, QK 부분만 수정하면 소스 오류가 활성화되어도 학습이 안정적으로 유지됩니다. 이러한 소스와 채널 간의 분리 현상은 오류의 근원이 실패 경로가 아니라는 것을 보여줍니다. 본 연구 결과는 다양한 아키텍처 및 크기에서 동일하게 나타나며, 다른 GPU 아키텍처에서도 재현 가능했습니다. 인과적 분석을 통해 각 업데이트가 현재 QK 가중치의 주요 세 가지 특이 방향으로부터 얼마나 벗어나는지 확인한 결과, query 투영의 가장 큰 특이값은 11.1로 유지되는 반면, 동일한 에너지를 다른 곳에 분배하면 237로 증가합니다. 따라서 QK 채널은 단순히 폭주 현상을 추적하는 것이 아니라, 초기 단계에서 폭주를 유발하는 역할을 합니다. 이러한 현상은 시간적으로 일관된 부호 변화에 의존하며, 전체적인 편차에는 영향을 받지 않습니다. QK-Guard는 Attention 로짓의 포화가 시작되면 작동하는 파라미터 없는 QK 정규화를 수행하는 제어기를 사용하여 QK 채널을 차단합니다. 이 방법은 테스트된 모든 폭주 현상을 완화하며, 6만 단계 동안 항상 활성화된 QK 정규화와 유사한 성능을 보입니다. 반면, QK 영역 외의 다른 조치는 동일한 조건에서도 실패했습니다. 이러한 결과는 각 오류 소스에 대한 개별적인 수정보다는 공유된 QK 지점에서 직접적인 개입이 더 효과적임을 시사합니다.

A bfloat16 transformer can train normally for many steps and then collapse abruptly. Distinct low-precision errors can trigger the same failure, leaving unclear whether each source needs its own repair or one shared route can be blocked. We isolate a reproduced GPT-2-class collapse to the streaming-softmax accumulator, where fp32 accumulation repairs it, and use the fault as an assay for moving controlled errors across sources. Errors placed outside attention still drive the same query-key (QK) spectral runaway, while correcting only QK keeps training stable with the source fault active. This source-channel dissociation shows that fault source is not failure channel. It holds across the tested architectures and scales and reproduces on a second GPU architecture. A causal probe projects each update off the current QK weights' leading three singular directions: the query projection's largest singular value stays at 11.1, whereas removing equal energy elsewhere leaves it at 237. The QK channel therefore drives the early runaway rather than merely tracking it. Entry depends on temporal sign-coherence across steps, not aggregate deviation. QK-Guard closes the channel with a dormant controller that switches on parameter-free QK normalization when attention-logit saturation begins. It contains every tested runaway and matches always-on QK normalization over 60k steps, while non-QK actions at the same trigger fail. The results support intervention at the shared QK locus rather than separate repair at each fault source.