Economical Jet Taggers -- Equivariant, Slim, and Quantized

이 논문은 성능 저하를 적게 유지하면서도 에너지 비용을 10분의 1 수준으로 줄인, L-GATr 제트 태거의 슬림하고 양자화되었으며 파라미터가 축소된 버전을 제시하며, 이는 LHC의 효율적인 트리거 레벨 제트 태깅을 위한 길을 열어준다.

핵심

대형 강입자 충돌기(LHC)를 거대하고 빠른 입자 공장이라고 상상해 보십시오. 매초 수많은 양성자를 충돌시켜 혼란스러운 파편의 분사를 만들어냅니다. 물리학자들은 이 파편 속에서 수십억 개의 평범한 입자들 사이에 숨겨진 특정 희귀 입자(예: '톱 쿼크')를 찾아내기 위해 분류 작업을 수행해야 합니다. 이 분류 과정을 **제트 태깅(jet tagging)**이라고 부릅니다.

수년 동안 과학자들은 이 분류를 위해 복잡한 컴퓨터 프로그램(머신러닝)을 사용해 왔습니다. 현재의 챔피언은 '트랜스포머(Transformers)'입니다. 이는 매우 정확하지만 거대하고, 느리며, 에너지를 많이 소비하는 강력한 AI 모델입니다. 이들은 마치 편지 한 통을 배달하기 위해 엄청난 연료를 소모하는 대형 트럭 함대를 운용하는 것과 같습니다. 그들은 일을 완벽히 수행하지만, 데이터를 수집하는 바로 그 순간(트리거 단계)에 사용하기에는 너무 크고 비쌉니다.

이 논문은 다음과 같은 단순한 질문을 던집니다: 우리는 이 거대한 트럭들을 편지를 배달하는 능력을 잃지 않으면서도 작고 연료 효율적인 스쿠터로 줄일 수 있을까?

저자들은 세 가지 주요 전략을 사용하여 이 문제를 해결했습니다.

1. "슬림한" 버전 (L-GATr-slim)

원래의 "L-GATr" 모델은 스칼라, 벡터, 텐서 등 가능한 모든 도구를 지닌 스위스 아미 나이프와 같습니다. 하지만 저자들은 대부분의 입자 물리학 작업에는 두 가지 도구, 즉 스칼라(숫자)와 벡터(방향이 있는 화살표)만 있으면 된다는 것을 깨달았습니다.

• 비유: 간단한 샌드위치를 만들기 위해 오븐, 블렌더, 믹서기까지 갖춘 산업용 주방 전체를 사용하겠다고 고집하는 요리사를 상상해 보십시오. 저자들은 "칼과 도마만 사용하자"라고 말했습니다.
• 결과: 저자들은 불필요한 도구들을 제거한 "Slim" 버전의 AI를 구축했습니다. 이 모델은 거대한 버전만큼 성능이 뛰어나면서도 학습 속도가 훨씬 빠르고 메모리를 적게 사용합니다. 이는 육중한 트럭을 민첩한 스포츠카로 교체하여 동일한 작업을 수행하게 만든 것과 같습니다.

2. "초소형" 버전 (Ultra-mini Taggers)

그다음 저자들은 "얼마나 더 작게 만들 수 있을까?"라는 질문을 던졌습니다. 그들은 이 AI 모델들을 원래 모델의 수백만 개 파라미터와 비교했을 때 약 1,000개 수준인 아주 작은 장난감 자동차 크기로 줄이려 시도했습니다.

• 비유: 도서관 한 권 분량의 지식을 단 한 장의 엽서에 담으려는 상황을 생각해 보십시오. 보통은 이야기를 놓치게 됩니다. 하지만 저자들은 정보를 올바르게 구성한다면(물리학 법칙을 준수하는 특정 '로런츠 불변(Lorentz-equivariant)' 규칙을 사용하여), 아주 좁은 공간에도 핵심적인 지식을 담을 수 있다는 것을 발견했습니다.
• 결과: 저자들은 매우 작은 모델의 경우, "LLoa" 구조는 레이어의 수를 줄이는 것이 가장 효과적이며, "L-GATr-slim"은 레이어의 *폭(width)*을 줄이는 것이 가장 효과적이라는 것을 발견했습니다. 이 미세한 크기에서도 이 모델들은 물리학적 특성을 반영하지 않는 기존의 AI 모델들보다 뛰어난 성능을 보였습니다.

3. "양자화된" 버전 (저정밀도 수학)

이것은 가장 극적인 에너지 절약 방법입니다. 표준 AI는 매우 정밀한 수학(예: 거리를 10억 분의 1 밀리미터 단위까지 측정하는 것)을 사용합니다. 저자들은 제트 태깅을 위해 그 정도의 정밀도가 필요하지 않다는 것을 깨달았습니다. 숫자를 상당히 반올림하더라도 충분히 괜찮습니다.

• 비유: 창고에 있는 사과 개수를 세고 있다고 상상해 보십시오.
  • 표준 AI: 모든 사과의 무게를 마이크로그램 단위까지 측정합니다. (정확하지만 시간이 오래 걸리고 많은 에너지 소식을 소모합니다.)
  • 양자화된 AI: 그냥 정수로 개수를 셉니다. (빠르고 에너지를 거의 쓰지 않으며, "사과가 몇 개인가"를 알기 위한 목적에는 충분합니다.)
• 방법: 저자들은 PARQ(Piecewise-Affine Regularized Quantization)라고 불리는 기술을 사용했습니다. 이것은 학습 과정 중에 숫자를 (0, 1, 또는 -1처럼) 단순하게 부드럽게 유도하는 스마트한 반올림 규칙이라고 생각하면 됩니다.
• 결과: 이 "거친" 숫자로 전환함으로써, AI를 실행하는 데 드는 에너지 비용을 10배(한 자릿수) 줄였습니다. AI는 매우 빠르고 에너지 효율적으로 변했으며, 정확도는 아주 약간만 감소했습니다.

종합적인 그림

저자들은 이 세 가지 전략인 슬리밍(Slimming), 소형화(Miniaturating), 그리고 **양자화(Quantizing)**를 결 결합하여 "경제적인 제트 태거(Economical Jet Taggers)"를 만들어냈습니다.

• 이것이 왜 중요한가? 현재 이러한 강력한 AI 모델들은 어떤 충돌을 보존하고 어떤 것을 버릴지 결정하는 하드웨어(트리거)에서 직접 실행하기에는 너무 큽니다.
• 목표: 이 모델들을 작고, 빠르고, 에너지 효율적으로 만듦으로써, 저자들은 궁극적으로 이들을 트리거 하드웨어에서 직접 실행하기를 희망합니다. 이를 통해 LHC는 데이터가 너무 빨리 폐기되어 놓쳤을지도 모를 새로운 물리학을 발견할 수 있도록, 충돌 순간에 AI를 사용하여 어떤 충돌을 저장할지 결정할 수 있게 될 것입니다.

요약하자면, 그들은 거대하고 에너지를 많이 소비하는 AI에게 다이어트를 시키고, 크기를 줄였으며, 더 적은 소수점으로 수학을 할 수 있도록 가르쳤습니다. 그 결과, 가장 중요한 입자들을 여전히 식별할 수 있는 작고 초효율적인 엔진을 만들어냈습니다.

기술 요약

기술 요약: 경제적인 제트 태거(Jet Tagger) – 등변성, 슬림화 및 양자화

문제 정의
현대 머신러닝(ML)은 거대 강입자 충돌기(LHC)에서의 제트 태깅(jet tagging)을 변화시켰으며, 로런츠 등변성 트랜스포머(Lorentz-equivariant transformers)는 최첨단 아키텍처로 부상했습니다. 그러나 L-GATr와 같은 주요 모델들은 계산 비용이 많이 들어 상당한 메모리와 훈련 시간을 요구합니다. 산업계의 추세는 네트워크와 데이터셋을 확장하는 방향을 선호하지만, LHC 물리학은 특히 이벤트 트리거 하드웨어의 메모리 및 지연 시간 요구사항과 관련된 구체적인 제약에 직면해 있습니다. 현재 제트 분류는 아직 트리거 단계에서 역할을 수행하지 못하고 있으나, 저자들은 이것이 수행되어야 한다고 주장합니다. 핵심 과제는 성능 저하를 최소|화하면서 현대적인 등변성 제트 태거의 크기와 계산 비용을 어떻게 줄일 것인가이며, 이는 잠재적으로 트리거 레벨에서의 배포를 가능하게 합니다.

방법론
본 논문은 자원 효율성을 최적화하기 위한 두 갈래의 전략을 제안합니다: 아키텍처 슬림화와 수치 양자화입니다.

1. L-GATr-slim 아키텍처:
   저자들은 간소화된 로런츠 등변성 트랜스포머(L-GATr) 버전을 도입합니다. 표준 L-GATr는 스칼라, 의사 스칼라(pseudo-scalar), 벡터, 축 벡터(axial-vector), 그리고 반대칭 랭크-2 텐서를 포함하는 기하 대수 표현을 사용합니다. 저자들은 대부분의 LHC 응용 분야에서 의사 스칼라, 축 벡터, 텐서가 불필요하다는 점에 주목했습니다. 따라서 L-GATr-slim은 잠재 표현을 스칼라와 벡터로만 제한합니다.

• 선형 계층(Linear Layers): 결합된 스칼라 및 벡터 표현을 작동하도록 확장되었으며, 로런츠 등변성을 유지하기 위해 벡터 성분이 단일한 학습 가능한 스칼라 계수를 공유하도록 보장합니다.
• 비선형성(Nonlinearity): 두 벡터의 내적에 비선형성을 적용하고 이를 벡터 출력에 곱하는 방식으로 Gated Linear Unit(GLU)을 적응시킵니다.
• 정규화(Normalization): 벡터 채널에 대해 민코프스키 내적의 절댓값을 사용하도록 RMSNorm을 수정합니다.
• 어텐션(Attention): 전체 L-GATr에서 사용하는 계산 비용이 높은 외적(outer product)을 피하는 특정 공식을 사용하여 스칼라 어텐션 행렬을 구축합니다.
• 구현: 이 아키텍처는 효율성을 위해 torch.compile로 컴파일되도록 설계되었습니다.

2. 양자화 전략:
   저자들은 비용을 더욱 줄이기 위해 저정밀도 데이터 타입과 가중치 양자화를 적용합니다.

• 데이터 타입 양자화: 선형 계층의 입력은 제로 포인트 양자화(zero-point quantization)를 사용하여 int8로 양자화하되, 정밀도가 중요한 연산과 역전파(backward pass)를 위해서는 bfloat16을 유지합니다. 이는 Transformer, ParT, L-GATr-slim, LLoCa-Transformer의 은닉층에 적용됩니다.
• 가중치 양자화: 선형 가중치는 **근사 경사 양자화(Proximal Gradient Quantization, PARQ)**를 사용하여 이진 또는 삼진 값으로 양자화됩니다. 이 방법은 양자화를 정규화 제약 조건으로 취급하며, 가중치를 업데이트하기 위해 근사 연산자(proximal operator)를 사용합니다. 저자들은 PARQ를 Straight-Through Estimation(STE)과 비교하여 PARQ가 더 나은 안정성과 성능을 제공함을 확인했습니다.
• 등변성 보존: 양자화가 로런츠 등변성을 위반하지 않도록 특별한 주의를 기울였습니다. LLoCa의 경우, 직교 정규화(orthonormalization)와 프레임 투영은 풀 정밀도(float32)로 유지되어, 저정밀 연산을 로런츠 불변량(Lorentz invariants)으로 제한합니다. L-GATr-slim의 경우, 전체 벡터가 양자화된 가중치와 곱해지는데, 이는 추가적인 대칭성 위반을 유발하지 않습니다.

3. 울트라 미니 스케일링(Ultra-Mini Scaling):
   저자들은 네트워크의 블록 수나 폭(채널)을 줄임으로써 아키텍처를 1,000개 파라미터 수준까지 축소하여 성능을 조사합니다.

주요 결과
본 연구는 세 가지 작업(탑 태깅, 진폭 회귀, 이벤트 생성)에 대해 제안된 방법들을 벤치마킹합니다.

• 성능 대비 효율성 (L-GATr-slim):

  • JetClass 데이터셋(다중 클래스 제트 태깅)에서 L-GATr-slim은 전체 L-GATr 및 LLoCa-Transformer과 동일한 성능(AUC ~0.9885)을 보이면서도, 훈련 시간을 6배(H100 GPU 기준 166시간에서 27시간으로) 단축하고 메모리 소비량을 2배 줄였습니다.
  • 진폭 회귀($Z + 4g$)에서 L-GATr-slim은 전체 L-GATr와 동일한 평균 제곱 오차(MSE)를 달성하면서도, 20배 적은 훈련 연산과 절반의 훈련 시간만을 필요로 했습니다.
  • 이벤트 생성($t\bar{t} + nj$)에서 슬림 아키텍처는 전체 모델과 동일한 음의 로그 가능도(negative log-likelihood) 성능을 보여주었습니다.

• 울트라 미니 태거:

  • 블록 수(깊이)를 줄일 때, LLoCa-Transformer은 매우 작은 크기(예: 1,000개 파라미터)에서 L-GATr-slim보다 우수한 성능을 보였습니다.
  • 블록 수(10개)를 고정하고 채널(폭)을 줄일 때, L-GATr-slim은 단 2개의 벡터와 4개의 스칼라 채널만으로도 1,000 이상의 배경 거부율(background rejection rate)을 유지하며 다른 1,000-파라미터 아키텍처들을 능가했습니다.

• 양자화 이득:

  • 입력을 int8로, 가중치를 삼진 값으로 양자화하면 성능 저하를 최소화하면서 에너지 소비를 약 한 자릿수(10배) 줄일 수 있습니다.
  • LLoCa-Transformer와 L-GATr-slim은 양자화에 강건하며, 일반적인 트랜스포머가 성능이 크게 저하되는 상황에서도 높은 성능을 유지합니다.
  • 가장 자원이 제한된 시나리오(1개 블록, 16차원 잠재 공간, int8)에서, 양자화된 LLoCa-Transformer(글로벌 정준화)은 전체 크기 모델에 비해 배경 거부율이 2배 감소했음에도 불구하고, 프리-그래프(pre-graph) 태거보다 우수한 성능을 유지했습니다.

의의 및 주장
본 논문은 이러한 "경제적인" 등변성 트랜스포머 버전이 고휘도 LHC(HL-LHC)를 위한 트리거 레벨 제트 태깅을 향한 실행 가능한 경로를 나타낸다고 주장합니다. 아키텍처 슬림화(불필요한 기하 대수 성분 제거)와 공격적인 양자화(PARQ 및 int8)를 결합함으로써, 저자들은 로런츠 등변성의 물리적 이점을 유지하면서도 약 1,000개의 파라미터를 가진 태거를 만들 수 있음을 입증했습니다.

저자들은 산업계의 표준은 확장이지만, LHC 물리학에는 "물리 기반 다운스케일링(physics-aware downscaling)" 접근 방식이 필요하다고 강조합니다. 결과는 작고 양자화되었으며 등변성을 가진 네트워크가 물리적 대칭성을 희생하지 않고도 FPGA와 같은 자원 제한적인 하드웨어에 배치될 수 있음을 시사하며, 이는 제트 하부 구조(jet substructure)의 실시간 분석을 위한 새로운 길을 열어줄 수 있습니다.