JetFormer: A Scalable and Efficient Transformer for Jet Tagging from Offline Analysis to FPGA Triggers

JetFormer는 고정밀 오프라인 분석부터 초저지연 FPGA 트리거까지 다양한 환경에서 효율적으로 작동하도록 설계된 확장 가능한 트랜스포머 기반의 제트 태깅(jet tagging) 모델입니다.

핵심

1. 배경: "모래사장에서 바늘 찾기"

거대한 입자 가속기(LHC)는 초당 수십억 번의 충돌을 일으킵니다. 이 충돌이 일어나면 입자들이 사방으로 튀어나가는데, 이를 '제트(Jet)'라고 부릅니다. 문제는 이 '제트'가 너무나도 많이 발생한다는 것입니다.

우리가 진짜 궁금한 건 '힉스 입자' 같은 아주 특별한 보물인데, 이 보물은 평범한 돌멩이(배경 데이터)들 사이에 아주 가끔 섞여 있습니다. 이 엄청난 양의 데이터를 전부 다 저장할 수는 없기 때문에, 데이터가 저장소로 들어가기 직전(트리거 단계)에 "이건 보물이다!" 혹은 "이건 그냥 돌멩이다!"를 0.000001초 만에 판단해서 걸러내는 초고속 필터가 필요합니다.

2. 기존의 문제: "똑똑하지만 너무 느린 천재"

최근에는 '트랜스포머(Transformer)'라는 아주 똑똑한 AI 모델이 등장했습니다. 이 모델은 입자들 사이의 복잡한 관계를 아주 잘 파악해서 보물을 기가 막히게 찾아냅니다.

하지만 문제가 있습니다. 이 AI는 '너무 똑똑해서 몸집이 너무 큽니다.' 마치 엄청난 지식을 가진 천재 교수님 같은데, 이 교수님이 판단을 내리는 데 시간이 너무 오래 걸리는 거죠. 입자 가속기처럼 눈 깜짝할 새 데이터가 지나가는 곳에서는, 아무리 똑똑해도 판단이 늦으면 이미 데이터는 지나가 버리고 맙니다.

3. JetFormer의 등장: "천재의 지식을 가진 초고속 요원"

연구진이 만든 **'JetFormer'**는 바로 이 문제를 해결하기 위해 설계되었습니다.

• 범용성 (Versatility): JetFormer는 아주 정밀한 분석이 필요한 '연구실용(Offline)'부터, 찰나의 순간에 결정을 내려야 하는 '현장용(Online/FPGA)'까지 모두 대응할 수 있는 '변신 로봇' 같은 모델입니다.
• 효율성 (Efficiency): 기존의 최고 모델(ParT)과 실력은 거의 비슷하면서도, 계산량(FLOPs)은 훨씬 적습니다. 즉, **"똑같은 실력을 갖췄는데 훨씬 가볍고 빠르다"**는 뜻입니다.

4. 핵심 기술: "다이어트와 압축의 마법"

연구진은 JetFormer를 아주 작은 하드웨어 칩(FPGA)에 넣기 위해 세 가지 마법을 부렸습니다.

1. 최적의 몸매 찾기 (Hyperparameter Optimization): 'Optuna'라는 도구를 써서, 정확도는 최대한 높이면서 몸집(계산량)은 최대한 줄이는 가장 완벽한 모델의 설계를 찾아냈습니다. (마치 근육은 유지하면서 체지방만 쏙 뺀 상태를 만드는 것과 같습니다.)
2. 가지치기 (Pruning): 중요하지 않은 신경망 연결들을 과감하게 잘라냈습니다. 성능은 거의 떨어뜨리지 않으면서 몸집을 절반으로 줄였습니다.
3. 극단적인 단순화 (1-bit Quantization): 원래 AI는 아주 복잡한 소수점 숫자로 계산을 하는데, 이를 **'1 또는 -1'**이라는 아주 단순한 숫자로 바꿔버렸습니다. 복잡한 수학 계산을 단순한 더하기/빼기로 바꿔서, 하드웨어가 훨씬 빠르게 처리할 수 있게 만든 것입니다.

5. 결론: "미래의 눈"

결과적으로 JetFormer는 매우 작은 칩에서도 작동할 수 있을 만큼 가벼워졌으면서도, 보물을 찾아내는 눈은 여전히 날카롭습니다.

이 기술이 완성되면, 인류는 거대한 입자 가속기에서 발생하는 엄청난 데이터의 홍수 속에서도, 우리가 찾고자 하는 우주의 비밀(새로운 입자)을 놓치지 않고 실시간으로 낚아챌 수 있게 됩니다.

---

요약하자면:

"JetFormer는 **천재적인 두뇌(Transformer)**를 가졌지만, **초고속 요원(FPGA용 최적화)**처럼 가볍고 빠르게 움직일 수 있도록 **다이어트(Pruning & Quantization)**를 완벽하게 마친, 입자 가속기용 차세대 AI 필터입니다."

기술 요약

[기술 요약] JetFormer: 오프라인 분석부터 FPGA 트리거까지 아우르는 확장 가능하고 효율적인 제트 태깅용 트랜스포머

1. 문제 배경 (Problem Statement)

고에너지 물리(HEP) 실험, 특히 LHC(대형 강입자 충돌기)에서는 초당 수십억 건의 충돌이 발생합니다. 이 중 물리적으로 의미 있는 신호(예: 힉스 보존, 톱 쿼크)를 식별하기 위해서는 '제트 태깅(Jet Tagging)' 기술이 필수적입니다.

• 기존 기술의 한계: 최근 트랜스포머(Transformer) 기반 모델(예: ParT)은 높은 정확도를 보여주지만, 모델이 너무 크고 복잡하여 실시간 데이터 선별을 담당하는 L1 트리거(Level-1 Trigger) 시스템의 하드웨어(FPGA)에 구현하기에는 지연 시간(Latency)과 자원 제약이 매우 큽니다.
• 딜레마: 고성능을 위한 '오프라인 분석용 모델'과 초저지연이 필요한 '온라인 트리거용 모델' 사이의 간극을 메울 수 있는 통합된 아키텍처가 부족했습니다.

2. 연구 방법론 (Methodology)

본 논문은 오프라인의 고성능과 온라인의 저지연성을 모두 충족하기 위해 JetFormer라는 범용 인코더 전용(Encoder-only) 트랜스포머 아키텍처를 제안합니다.

A. 아키텍처 설계 (JetFormer Architecture)

• BERT 기반 구조: BERT와 유사한 인코더 전용 구조를 채택하고, 제트 전체의 특징을 요약하기 위해 학습 가능한 [CLS] 토큰을 사용합니다.
• 하드웨어 친화적 최적화:
  • Batch Normalization: 하드웨어 구현이 까다로운 Layer Norm 대신, 추론 시 파라미터가 고정되는 Batch Norm으로 대체했습니다.
  • ReLU 활성화 함수: 계산 비용이 높은 SiLU 대신 ReLU를 사용하여 연산 효율을 높였습니다.
  • 순열 불변성(Permutation Invariance): 입력을 정렬되지 않은 집합으로 처리하여 입자 순서에 상관없는 물리적 특성을 반영했습니다.

B. 최적화 및 압축 파이프라인 (Optimization & Compression)

하드웨어 배포를 위해 세 가지 핵심 단계를 거칩니다.

1. 다목적 하이퍼파라미터 최적화 (HPO): Optuna를 사용하여 정확도는 극대화하고 연산량(FLOPs)은 최소화하는 최적의 모델 크기(JetFormer-tiny)를 찾아냅니다.
2. 구조적 가지치기 (Structured Pruning): Taylor 중요도 기반의 가지치기를 통해 채널/필터를 제거하여 정확도 손실을 최소화하면서 연산량을 약 50% 줄입니다.
3. 1비트 양자화 (1-bit Quantization): BitNet 프레임워크를 활용하여 가중치를 +1 또는 -1로 양자화합니다. 이를 통해 모델 크기를 82~92%까지 획기적으로 줄입니다.

C. 하드웨어 구현 (FPGA Implementation)

• Allo 프레임워크 확장: 기존의 Allo 프레임워크가 지원하지 않던 트랜스포머 핵심 연산(Log Softmax, ReLU 3D, Class Token 조작 등)을 직접 구현하여 FPGA 합성이 가능하도록 확장했습니다.

3. 주요 기여 (Key Contributions)

1. 통합 프레임워크 제공: 하나의 아키텍처로 고정밀 오프라인 분석부터 초저지연 FPGA 트리거까지 대응 가능한 확장성을 증명했습니다.
2. 효율적인 트랜스포머 설계: 입자 간의 명시적인 상호작용 항(Pairwise interaction) 없이도 기존 SOTA 모델과 대등한 성능을 내면서 연산량은 크게 줄였습니다.
3. 하드웨어-알고리즘 공동 설계(Co-design): 모델 설계 단계부터 하드웨어 제약 조건을 고려한 최적화 파이프라인을 구축했습니다.

4. 연구 결과 (Results)

A. 모델 성능 (Accuracy & Efficiency)

• JETCLASS 데이터셋: 최신 모델인 ParT와 비교했을 때, 정확도 차이는 단 **0.7%**에 불과하지만, 연산량(FLOPs)은 37.4% 더 적게 사용합니다.
• HLS4ML 150P 데이터셋: 기존 MLP, Deep Sets, Interaction Networks 모델보다 정확도가 3~4% 향상되었습니다.

B. 압축 효과 (Compression)

• Pruning: 정확도 손실을 0.5% 미만으로 유지하며 연산량을 50% 절감했습니다.
• Quantization: 모델 크기를 82~92% 감소시켰으며, 정확도 하락은 1.5~3.5% 수준으로 억제했습니다.

C. 하드웨어 검증 (Hardware Evaluation)

• JetFormer-tiny: FPGA 상에서 동작 가능한 수준의 자원 사용량을 확인했으며, 배치 사이즈를 줄임으로써 sub-microsecond(마이크로초 미만) 수준의 지연 시간을 달성할 수 있는 가능성을 보여주었습니다.

5. 의의 (Significance)

이 연구는 딥러닝의 최첨단 기술인 트랜스포머를 실제 입자 물리학 실험의 가장 가혹한 환경인 **실시간 하드웨어 트리거 시스템에 적용할 수 있는 실질적인 경로(Practical Pathway)**를 제시했다는 점에서 매우 큰 의의가 있습니다. 이는 향후 LHC의 업그레이드(Phase-2) 과정에서 더 정교한 물리 신호를 실시간으로 포착하는 데 핵심적인 역할을 할 것으로 기대됩니다.