Towards Foundation Models on Hardware Accelerators for Particle Physics

기술 요약: 하드웨어 가속기를 위한 입자 물리학용 파운데이션 모델을 향하여

문제 정의
입자 물리학 실험, 특히 거대 강입자 충돌기(LHC)는 맞춤형 하드웨어(예: FPGA)에서 실시간 의사결정을 수행해야 하는 심각한 대역폭 제약에 직면해 있습니다. 오프라인 분석에서는 제트 태깅(jet tagging)과 같은 작업에서 최첨단 성능을 달ert하기 위해 수억 개의 파라미터를 가진 거대 파운데이션 모델을 활용할 수 있지만, 이러한 모델들은 마이크로초 단위 내에 결정을 내려야 하는 하드웨어 트리거에는 계산 비용이 너무 높습니다. 현재의 하드웨어 제약 기반 솔루션은 종종 단순화된 알고리즘이나 처음부터 학습된 더 작은 모델에 의존하며, 이 과정에서 대규모 사전 학습(pretraining)이 제공하는 정확도 이득을 희생할 가능성이 있습니다. 과제는 대규모 파운데이션 모델의 고성능 지식을 상당한 정확도 손실 없이 하드웨어 배포에 적합한 압축되고 효율적인 아키텍처로 전이하는 것입니다.

방법론
저자들은 거대한 "교사(teacher)" 모델로부터 압축된 "학생(student)" 네트워크로 지식을 전이하기 위해 **지식 증류(knowledge distillation)**를 채택했습니다.

• 교사 모델: 교사는 OmniLearned 파운데이션 모델(423M 파라미터)로, 28개의 블록과 32개의 어텐션 헤드를 가진 트랜스포머 기반 아키텍처입니다. 이 모델은 210개의 클래스에 걸친 10억 개 이상의 제트 데이터셋을 통해 사전 학습되었으며, 이후 탑 쿼크(top-quark) 제트 태깅을 위해 미세 조정되었습니다. 학생 학습 동안 교사는 동결되며, 교사의 소프트 예측(logits)은 사전 계산되어 캐싱됩니다.
• 학생 아키텍처: 학생은 순열 불변성(permutation-invariant)을 가지며 어텐션이 없는 Deep Sets 네트워크입니다. 이는 개별 입자를 임베딩하는 공유 MLP($\phi$), 평균화 연산, 그리고 풀링된 벡터를 클래스 로짓으로 매핑하는 두 번째 MLP($\rho$)로 구성됩니다. 베이스라인 학생 모델은 약 11k 개의 파라미터를 가집니다. 표현력을 높이기 위해 저자들은 상위 64개 고가 $p_T$ 구성 성분 간의 쌍별 특징(불변 질량, 각도 분리, $k_T$)을 처리하는 **메시지 패싱 레이어(message-passing layer)**를 도입하여 파라미터 수를 약 17k로 증가시켰습니다.
• 증류 전략: 학습 손실은 정답 레이블을 사용하는 교차 엔트로피 항과 학생의 소프트 출력 분포를 교사와 일치시키는 쿨백-라이블러(KL) 발산 항의 조합으로 구성됩니다. 저자들은 다음 세 가지 특정 변수를 비교합니다:
  1. 학생 모델에 메시지 패싱 레이어를 포함하는 여부.
  2. 소프트 레이블(증류)을 사용하여 학습하는 것과 하드 레이블(정답만 사용)을 사용하여 학습하는 것의 비교.
  3. 사전 학습된 파운데이션 모델로부터 증류하는 것과 동일한 아키텍처를 가진 교사가 특정 작업을 위해 처음부터 학습된 경우의 비교.
• 하드웨어 준비성: 배포 가능성을 평가하기 위해, 저자들은 사후 학습 양자화(PTQ)와 양자화 인식 학습(QAT)을 모두 사용하여 8비트 양자화를 평가합니다.

주요 결과
본 연구는 파운데이션 모델의 지식이 효율적인 네트워크로 효과적으로 압축될 수 있음을 입증하며, 특히 낮은 신호 효율에서 작동하는 트리거 시스템에 중요한 지표인 배경 반사(background rejection)에서 특정 이득이 관찰되었습니다.

• 아키텍처 강화: Deep Sets 학생 모델에 메시지 패싱 레이어를 추가함으로써 정확도가 92.85%에서 93.98%로 향상되었으며, 30% 신호 효율에서의 배경 반사(1/FPR)가 748에서 1980으로 크게 증가했습니다.
• 증류 효능: 교사의 하드 레이블(정답)만으로 학습하는 대신 소프트 레이블(증류)을 사용하여 학습했을 때, 30% 신호 효율에서 배경 반사가 1.6배 더 개선되었습니다(1980 vs 1216). 전체 정확도는 미미한 차이(93.98% vs 93.82%)를 보였습니다.
• 사전 학습의 가치: 사전 학습된 OmniLearned 파운데이션 모델로부터 증류하는 것이 동일한 아키텍처를 가졌으나 특정 작업을 위해 처음부터 학습된 교사로부터 증류하는 것보다 우수한 결과를 나타냈습니다. 사전 학습된 교사는 두 교사의 AUC와 정확도가 유사함에도 불구하고, 30% 신호 효율에서 학생의 배경 반사를 1.9배 개선했습니다(1980 vs 1020). 이는 파운데이션 모델의 사전 학습이 ROC 곡선의 "꼬리(tail)" 부분에 유익한 특징을 포착한다는 것을 시사합니다.
• 양자화 및 효율성:
  • PTQ는 상당한 성능 저하를 일으켰습니다(예: 메시지 패싱 학생의 정확도가 59%로 하락).
  • QAT는 성능을 성공적으로 회복하였으며, 메시지 패싱 학생은 93.95%의 정확도(float32 대비 단 0.03 포인트 차이)를 달성하고 높은 배경 반사를 유지했습니다.
  • 계산 비용: 증류된 학생(메시지 패싱 포함)은 제트당 14.3M MACs(곱셈-누산 연산)를 요구하는 반면, 교사는 37.9G MACs를 요구합니다. 이는 3~4 자릿수(orders of magnitude)의 감소를 의미하며, 모델을 FPGA 합성의 유망한 후보로 만듭니다.

의의 및 주장
본 논문은 파운데이션 모델이 압축 가능하고 하드웨어 친화적인 네트워크를 학습시키기 위한 효과적인 교사 역할을 할 수 있다고 주장합니다. 저자들은 이 접근 방식의 주요 이점이 단순히 전체 분류 정확도에 있는 것이 아니라, 하드웨어 트리거와 가장 밀접한 관련이 있는 운용 영역인 낮은 신호 효율에서의 배경 반사에 있다는 점을 강조합니다.

이 연구는 머신러닝 모델과 하드웨어 가속기의 "공동 설계(co-design)"를 위한 경로를 구축합니다. 423M 파라미터의 파운데션 모델을 17k 파라미터 네트워크로 증류하여 8비트 양자화 하에서도 높은 성능을 유지할 수 있음을 보여줌으로써, 저자들은 차세대 하드웨어 배포가 실현 가능하다고 주장합니다. 논문은 결론적으로, FPGA 구현을 위한 고수준 합성(HLS) 작업이 진행 중이지만, 현재의 결과는 효율적인 실시간 입자 물리학 트리거 개발을 위해 사전 학습된 파운데이션 모델을 사용하는 것의 타당성을 검증한다고 밝히고 있습니다.