TurboPairFormer: 최적화된 삼각형 어텐션 커널을 사용한 빠르고 안정적인 단백질 폴딩 모델 학습

삼각형 어텐션(Triangular attention)은 AlphaFold3 스타일의 생체 분자 모델에서 핵심 계산이며, 토큰 수 측면에서 세제곱 비용이 발생합니다. 이 공유 쌍 편향(shared pair bias)은 일반적인 쿼리 및 키에 대한 감소 외에도 어텐션 슬라이스 전반에 걸쳐 그래디언트 감소를 추가합니다. 우리가 검토한 오픈 소스 백엔드들은 반복적인 확률 재계산, 부동소수점 원자적 연산(floating-point atomics), 또는 전체 점수-그래디언트 저장 방식을 통해 이러한 감소를 처리합니다. 별도로, BF16으로 반올림된 순방향 출력으로부터 소프트맥스 역방향 보정(softmax backward correction)을 계산하는 과정은 수치적 정밀도를 잃게 합니다. 우리는 이 두 가지 문제를 모두 해결하는 NVIDIA Hopper GPU용 삼각형 어텐션 구현체인 TurboPairFormer를 제시합니다. 우리의 핵심 타일 병렬 역방향 알고리즘(key-tile-parallel backward algorithm)은 순서화된 부분 감소(ordered partial reductions)를 사용하여 부동소수점 원자적 연산이나 전체 점수-그래디언트 저장 없이, 쿼리, 키, 값 및 쌍 편향 그래디언트 각각의 확률 타일을 한 번 재계산합니다. 출력 잔차 보상(Output-residual compensation)은 BF16 출력을 변경하지 않으면서도, 출력 반올림 잔차의 BF16 근사치를 유지하여 FP32에서 역방향 보정을 더 정확하게 계산합니다. 크롭 사이즈 384, 640, 768 및 헤드 차원 16과 32인 BF16 입력으로 TurboPairFormer는 본 논문에서 비교된 구현체 중 가장 낮은 평균 쿼리, 키 및 쌍 편향 그래디언트 RMSE를 FP64 참조값 대비 달성합니다. 잔차 보상은 제어된 제거(controlled ablations)에서 이러한 RMSE 값을 28-47% 감소시킵니다. 네 가지 모든 그래디언트는 고정 실행 조건 하의 600개 입력 사례에 걸쳐 다섯 번 반복 호출 시 비트 단위로 동일합니다. 우리의 삼각형 곱셈 커널과 통합된 OpenFold3는 16개의 H100 GPU에서 평가된 백엔드 구성 중 최저 GPU 계산 시간/최적화기 스텝을 달성하며, 크롭 사이즈 768에서 OpenFold3의 Triton 백엔드 대비 $1.73 imes$, cuEquivariance 대비 $1.13 imes$의 속도 향상을 보여줍니다.