JEDI-linear: Fast and Efficient Graph Neural Networks for Jet Tagging on FPGAs

본 논문은 정밀한 양자화와 승산기 없는 연산을 활용하여 FPGA 에서 기록적인 지연 시간과 자원 효율성을 달성함으로써 HL-LHC CMS 레벨 1 트리거 시스템의 실시간 제트 태깅 요구 사항을 성공적으로 충족시키는 새로운 선형 복잡도 그래프 신경망 아키텍처인 JEDI-linear 를 소개합니다.

핵심

당신은 거대한 공항 (대형 강입자 충돌기) 의 초고속 보안 검색대를 운영한다고 상상해 보세요. 매 25 나노초마다 새로운 입자 '비행'이 지면과 충돌하며 파편의 혼란스러운 분출을 만들어냅니다. 당신의 임무는 이 분출을 즉시 살펴보고 결정하는 것입니다: "이건 지루한 쓰레기 더미인가, 아니면 드물고 귀한 보물인가?"

모든 파편 조각을 저장하려 한다면, 순식간에 저장 공간이 부족해질 것입니다. 따라서 트리거 시스템—오직 흥미로운 사건만 선별하기 위해 순간적인 결정을 내리는 초고속 필터—이 필요합니다.

이제 이 논문이 등장합니다. 저자들은 보안 요원들이 더 나은 결정을 내릴 수 있도록 돕는 새로운 초고속 '뇌' (JEDI-linear라고 명명됨) 를 구축했으며, 이 뇌를 극도로 빠르게 작동해야 하는 작고 특수한 컴퓨터 칩 (FPGA) 에 탑재하는 데 성공했습니다.

다음은 그들의 발명을 간단한 비유로 설명한 내용입니다:

1. 문제: "악수" 병목 현상

이러한 입자 분출 ( '제트'라고 함) 을 분류하는 기존 방법들은 거대한 라운드 로빈 악수와 유사한 기법을 사용했습니다.

• 기존 방식: 64 명이 있는 방을 상상해 보세요. 그룹을 이해하기 위해, 기존 방법은 모든 단일 개인이 서로 다른 모든 사람과 개별적으로 돌아서 악수해야 했습니다.
• 결과: 64 명이 있다면, 악수는 4,000 회 이상 발생합니다. 이는 너무 오래 걸리며, 동시에 이야기하려는 사람들로 방이 너무 붐빕니다. 입자 물리학 세계에서 이 '악수' 과정은 너무 느리고 하드웨어 공간을 너무 많이 차지하여 실시간 보안 검색에 유용하지 않습니다.

2. 해결책: "그룹 허들" (JEDI-linear)

저자들은 모든 사람이 개별적으로 악수할 필요가 없다는 것을 깨달았습니다. 대신, 그들은 선형 복잡도 접근법을 고안했습니다.

• 새로운 방식: 개별 악수 대신, 방 안의 모든 사람이 단순히 현재 기분을 공유하기 위해 손을 든다고 상상해 보세요. 그리고 한 명의 '캡틴'이 모든 기분을 하나의 큰 요약으로 수집합니다. 그런 다음 캡틴은 모두에게 말합니다, "이것이 전체 그룹의 분위기입니다."
• 마법: 이제 4,000 회의 악수 대신, 64 명의 사람이 한 번씩 말하기만 하면 됩니다. 작업은 선형적으로 확장됩니다 (사람을 두 배로 늘리면 작업도 두 배가 되지, 네 배가 되지 않습니다). 이것이 바로 'JEDI-linear' 부분입니다: 이는 지저분하고 느린 쌍별 상호작용 없이 그룹의 맥락을 유지합니다.

3. 하드웨어 해킹: 작은 칩에 맞추기

새로운 '허들' 방법으로도 여전히 이 뇌는 보안 시스템에 사용되는 특정 유형의 칩에 들어갈 만큼 작고 빨라야 했습니다. 저자들은 두 가지 영리한 트릭을 사용했습니다:

• 맞춤형 유니폼 트릭 (양자화):
  일반적으로 컴퓨터는 모든 숫자를 같은 방식으로 처리합니다 (모든 병사에게 똑같은 무거운 코트를 주는 것과 같습니다). 저자들은 수학의 일부 부분은 매우 민감하여 고정밀도 (무거운 코트) 가 필요하지만, 다른 부분들은 크게 신경 쓰지 않는다는 (가벼운 티셔츠) 사실을 깨달았습니다. 그들은 시스템이 '맞춤형 유니폼'을 착용하도록 훈련시켜, 정밀도가 많이 필요하지 않은 숫자들에게 작고 효율적인 비트 너비를 할당했습니다. 이로 인해 메모리 사용량이 크게 줄었습니다.

• 승산기 제거 트릭 (분산 산술):
  표준 칩은 수학 계산을 위해 특수하고 비싼 '승산기' 블록을 사용하는데, 이는 무겁고 전력을 많이 소모하는 엔진과 같습니다. 저자들은 이러한 엔진을 가산기와 시프터의 영리한 시스템 (슬라이드 룰이나 블록 쌓기를 사용하는 것과 같습니다) 으로 대체했습니다.

  • 결과: 그들은 무거운 '승산기 엔진' (DSP 블록) 을 완전히 제거했습니다. 이로 인해 막대한 공간과 전력을 절약하여, 이전에는 부하를 처리할 수 없었던 칩에서 시스템이 작동할 수 있게 되었습니다.

4. 결과: 속도와 효율성

이 새로운 시스템을 기존 최우수 방법들과 비교하여 테스트했을 때:

• 속도: 기존 방법보다 3.7 배에서 11.5 배 더 빠릅니다. 60 나노초 미만 (눈을 깜빡이는 것보다 빠름) 에 결정을 내릴 수 있습니다.
• 효율성: 결정 사이의 시동 시간을 최대 150 배 줄이고 칩 상의 공간을 6.2 배 더 적게 사용합니다.
• 정확도: 더 작고 빠르지만, 이전의 무거운 모델들보다 희귀한 입자 제트를 식별하는 데 실제로 더 정확합니다.

왜 이것이 중요한가

저자들은 이것이 상호작용 기반 AI 모델이 CERN 의 고광도 대형 강입자 충돌기 (High-Luminosity Large Hadron Collider) 의 레벨 -1 트리거 시스템에서 사용될 만큼 빠르고 작아진 첫 번째 사례라고 주장합니다.

이는 느린 수동 검색에서 희귀한 물건을 절대 놓치지 않으면서도 줄을 절대 늦추지 않는 초고속 자동 스캐너로 공항 보안 검색을 업그레이드하는 것과 같습니다. 이를 통해 과학자들은 이전에는 너무 빨라 볼 수 없었던 희귀한 물리 현상을 포착할 수 있게 되었으며, 이는 표준 계산기보다 적은 하드웨어를 사용하면서 가능합니다.

간단히 말해: 그들은 복잡하고 느린 AI 를 가져와, 끊임없이 '자기 자신과 대화'할 필요가 없도록 수학을 단순화하고, 공간을 절약하기 위해 맞춤형 옷을 입혔으며, 무거운 엔진을 경량 기어로 교체했습니다. 그 결과, 칩에 들어가는 초고속 소형 뇌가 만들어졌으며, 이는 실시간으로 희귀 입자를 포착할 수 있습니다.

기술 요약

"JEDI-linear: FPGAs 에서 제트 태깅을 위한 빠르고 효율적인 그래프 신경망" 논문에 대한 상세한 기술 요약입니다.

1. 문제 제기

CERN 의 대형 강입자 충돌기 (LHC) 에서 수행되는 고에너지 물리 실험은 매초 수백 테라바이트에 달하는 방대한 데이터를 생성합니다. 이를 관리하기 위해 Level-1 트리거 (L1T) 시스템은 FPGAs 를 사용하여 실시간 (수 마이크로초 이내) 으로 이벤트를 필터링해야 합니다.

• 과제: 제트 태깅 (입자 분무의 기원을 식별하는 작업) 은 이러한 필터링에 필수적입니다. 그래프 신경망 (GNN), 특히 JEDI-net 과 같은 상호작용 네트워크 (INs) 는 입자 간 상호작용을 모델링하여 뛰어난 정확도를 제공하지만, 다음과 같은 이유로 L1T 를 위한 FPGA 배포가 어렵습니다:
  • 계산 복잡도: 표준 GNN 은 명시적인 쌍별 엣지 계산 ($O(N^2)$) 을 필요로 하여 입자가 많은 제트에서 병목 현상을 초래합니다.
  • 하드웨어 제약: 엄격한 지연 시간 요구사항 (<100 ns), 제한된 FPGA 리소스 (일반적으로 <1 Super Logic Region), 그리고 낮은 시작 간격 (II) 필요성.
  • 리소스 소비: 기존 FPGA 기반 GNN 들은 종종 수천 개의 디지털 신호 처리 (DSP) 블록과 큰 룩업 테이블 (LUT) 수를 요구하여, 다른 알고리즘과 함께 실제 환경에 배포하는 것을 비현실적으로 만듭니다.

2. 방법론

저자들은 하드웨어 효율성을 위해 특별히 설계된 새로운 GNN 아키텍처인 JEDI-linear와 고급 최적화 기법을 결합하여 제안합니다.

A. 알고리즘 혁신: 선형 복잡도

• 핵심 개념: 원래 JEDI-net 은 모든 입자 쌍 간의 상호작용을 계산합니다 ($O(N^2)$). JEDI-linear 는 엣지 상호작용 함수 $f_R$을 단순한 아핀 변환 (단일 밀집 계층) 으로 재구성합니다.
• 수학적 유도: $f_R(I_i \| I_j) = W_1 I_i + W_2 I_j + C$라고 가정하면, 명시적인 쌍별 합계를 전역 집계로 다시 쓸 수 있습니다. 입자 $i$에 대한 상호작용 임베딩은 모든 입자 특성의 전역 평균의 함수와 개별 입자 특성의 변환의 합으로 표현됩니다.
• 결과: 이는 계산 복잡도를 2 차 $O(N^2)$에서 선형 $O(N)$으로 줄여, 전역 컨텍스트를 유지하면서 명시적인 엣지 수준 계산을 제거합니다.

B. 하드웨어 최적화 전략

1. 세분화된 양자화 인식 학습 (QAT):

   • 균일 양자화와 달리, 저자들은 파라미터별 비트 너비 최적화 접근법을 사용합니다.
   • 미분 가능한 서로게이트 기울기를 사용하여, 훈련 과정에서 정확도와 하드웨어 비용 (Effective Bit Operations, EBOPs 로 측정) 에 미치는 영향에 따라 각 가중치에 특정 비트 너비를 자동으로 할당합니다.
   • 이를 통해 많은 가중치가 가지치기되거나 (비트 너비가 0 으로 설정됨) 1~2 비트로 축소되는 혼합 정밀도 모델을 가능하게 하여, 정확도 손실 없이 모델 크기를 크게 줄입니다.

2. 승산기 없는 MAC 를 위한 분산 산술 (DA):

   • 리소스 사용을 더욱 줄이기 위해, 구현은 기존 승산기를 **분산 산술 (Distributed Arithmetic)**로 대체합니다.
   • DA 는 행렬 - 벡터 곱셈을 LUT 를 통해 구현된 시프트 - 덧셈 연산으로 분해합니다.
   • 결과: 이 설계는 DSP 블록의 필요성을 완전히 제거하여, FPGA 에서 더 풍부하고 유연한 LUT 와 레지스터에만 의존합니다.

3. 완전 언롤된 데이터 흐름 아키텍처:

   • 이 설계는 모든 연산이 전용 하드웨어에 매핑되는 완전히 정적이며 언롤된 데이터 흐름을 사용합니다.
   • 이는 리소스 공유 및 제어 오버헤드를 방지하여 1 사이클 시작 간격과 결정론적이며 초저 지연 시간을 가능하게 합니다.

3. 주요 기여

• JEDI-linear 아키텍처: 명시적인 쌍별 상호작용을 제거하여 선형 복잡도를 달성한 제트 태깅용 최초의 상호작용 기반 GNN 으로, 대규모 입자 수에 대해 확장 가능합니다.
• 하드웨어 인식 공동 설계: 세분화된 혼합 정밀도 양자화와 분산 산술을 통합하여 승산기 및 DSP 가 없는 구현을 생성합니다.
• 자동화 프레임워크: 이러한 복잡하고 언롤된 아키텍처에 대한 심볼릭 계산 그래프를 자동으로 추적하고 합성 가능한 Verilog 를 생성하는 확장된 da4ml 프레임워크.
• 오픈 소스: 재현성을 지원하기 위한 JEDI-linear 템플릿 및 코드 공개.

4. 실험 결과

모델은 CMS Level-1 트리거 시스템 (Correlator Layer 2) 을 대상으로 한 AMD VU13P FPGA에서 평가되었습니다.

• 지연 시간 및 처리량:
  • <60 ns 지연 시간 달성 (예: 16 개 특성으로 16 개 입자의 경우 52 ns) 및 1 클록 사이클 시작 간격.
  • 최신 (SOTA) GNN 설계 (예: LL-GNN, JEDI-net 변형) 에 비해 3.7 배에서 11.5 배 낮은 지연 시간 및 최대 150 배 낮은 시작 간격.
• 리소스 효율성:
  • 모든 구성에서 0 개의 DSP 블록 사용 (SOTA 설계는 종종 5,000~9,000 개 이상의 DSP 를 사용함).
  • SOTA 모델 대비 최대 6.2 배 낮은 LUT 사용량.
  • 예: 32 개 입자 JEDI-linear 모델은 GNN J5 모델보다 LUT 를 6.2 배 적게 사용하며 11.5 배 낮은 지연 시간을 달성하면서도 더 높은 정확도 (81.4% 대 79.9%) 를 제공합니다.
• 정확도:
  • 최대 82.4% 분류 정확도 달성 (64 개 입자로 16 개 특성 입력 기준).
  • 다양한 입자 수 (8~128) 에서 DeepSets (DS) 및 이전 GNN 구현체보다 우수한 성능을 보임.
  • 다른 모델들이 성능이 저하되거나 실현 불가능해지는 반면, 입자 수가 증가함에 따라 높은 정확도를 유지하는 뛰어난 확장성을 입증함.

5. 의의

• 최초의 현실적 실현 가능성: 이는 HL-LHC CMS Level-1 트리거의 엄격한 지연 시간 (<60 ns) 과 리소스 제약을 충족하는 제트 태깅용 최초의 GNN 입니다. 이는 이전에 더 단순하고 정확도가 낮은 모델이 지배하던 실시간 하드웨어 트리거에서 강력한 GNN 의 사용을 가능하게 합니다.
• 확장성: 선형 복잡도는 시스템이 지수적인 리소스 증가 없이 많은 수의 입자 (최대 128 개) 를 가진 제트를 처리할 수 있게 하여, 더 높은 광도 실행을 위한 시스템을 미래 지향적으로 만듭니다.
• 광범위한 영향: 선형화된 처리, 혼합 정밀도 양자화, 분산 산술, 그리고 완전 언롤된 설계와 같은 기법은 입자 물리학을 넘어 신뢰할 수 있는 DNN, VAE, Transformer 와 같은 다른 저지연 도메인에도 적용 가능합니다.

결론적으로, JEDI-linear 는 고정밀 딥러닝 알고리즘과 실시간 하드웨어 트리거의 극한 제약 사이의 간극을 성공적으로 연결하여, LHC 차세대 지능형 트리거 시스템의 길을 마련했습니다.