PQuantML: A Tool for End-to-End Hardware-aware Model Compression

PQuantML 는 제트 태깅과 같은 실시간 LHC 데이터 처리를 포함한 엄격한 지연 시간 제약이 있는 환경에 최적화된 엔드 투 엔드 하드웨어 인식 신경망 모델 압축 라이브러리로, 가지치기와 고정 소수점 양자화를 통합하여 기존 도구들보다 효율적인 모델 경량화를 가능하게 합니다.

핵심

이 논문은 **"PQuantML"**이라는 새로운 도구를 소개합니다. 이 도구를 쉽게 설명하기 위해, 거대한 데이터 처리 공장과 그 안에 설치된 초고속 로봇에 비유해 보겠습니다.

1. 배경: 거대한 데이터 폭풍과 좁은 통로

대형 강입자 충돌기 (LHC) 는 매초 수백 테라바이트의 데이터를 쏟아냅니다. 이는 마치 거대한 폭포수가 쏟아지는 것과 같습니다. 하지만 이 모든 물을 한 번에 저장할 수 있는 통 (저장소) 은 없습니다.

그래서 과학자들은 **초고속 필터 (트리거 시스템)**를 설치했습니다. 이 필터는 폭포수 속에서 가장 중요한 물방울 (중요한 사건) 만 골라내야 합니다. 문제는 이 필터가 마이크로초 (100 만분의 1 초) 단위로 작동해야 한다는 점입니다.

기존의 복잡한 인공지능 (AI) 모델은 이 필터에 넣기엔 너무 무겁고 느립니다. 마치 폭포수 앞에 거대한 코끼리를 세워놓고 "이 코끼리가 물을 걸러줘"라고 하는 것과 비슷합니다. 코끼리는 정확할 수는 있지만, 너무 느리고 공간도 많이 차지합니다.

2. 해결책: PQuantML (코끼리를 미니어처로 만드는 도구)

이때 등장한 것이 PQuantML입니다. 이 도구는 무거운 AI 모델 (코끼리) 을 작고 빠르면서도 똑똑한 미니어처로 변신시키는 마법 지팡이 같은 역할을 합니다.

이 도구는 두 가지 주요 기술을 사용합니다.

A. 가지치기 (Pruning) - "불필요한 나뭇가지 자르기"

AI 모델은 수많은 연결 (가중치) 로 이루어진 거대한 나무입니다. 하지만 모든 가지가 중요한 것은 아닙니다.

• PQuantML은 모델이 학습하는 동안 "이 가지는 쓸모없어"라고 판단하면, 그 가지를 잘라냅니다.
• 비유: 거대한 나무에서 잎이 거의 없는 마른 가지를 잘라내면, 나무는 더 가벼워지고 바람 (데이터) 이 더 빠르게 통과할 수 있습니다.
• 이 도구는 가지를 무작위로 자르는 것뿐만 아니라, 규칙적으로 (예: 2 개 중 1 개) 자르는 등 하드웨어가 이해하기 쉬운 형태로 잘라냅니다.

B. 양자화 (Quantization) - "정밀도 조절"

기존 AI 모델은 숫자를 매우 정밀하게 (예: 32 자리 소수점) 다룹니다. 하지만 필터에서는 그렇게 정밀할 필요가 없습니다.

• PQuantML은 숫자의 정밀도를 낮춥니다. (예: 32 자리 → 8 자리).
• 비유: 고가의 정밀 저울 대신, 일반적인 주방 저울을 사용하는 것과 같습니다. 정확도는 약간 떨어질 수 있지만, 훨씬 가볍고 빠르게 작동합니다.
• 특히 FPGA(현장 프로그래밍 가능한 게이트 어레이)라는 특수한 칩에서 작동할 때, 이 정밀도 조절은 연산 속도를 획기적으로 높여줍니다.

3. PQuantML 의 특별한 점: "한 번에 끝내는 통합 솔루션"

기존의 도구들은 양자화 (정밀도 조절) 는 잘했지만, 가지치기 (불필요한 부분 제거) 는 사용자가 따로 해야 했습니다. 마치 자동차를 개조할 때, 엔진은 A 공장에서, 바퀴는 B 공장에서 따로 주문해야 하는 것과 같았습니다.

PQuantML은 이 두 가지를 하나의 통합된 시스템으로 바꿨습니다.

• 사용자 친화적: 복잡한 코딩 없이 설정 파일 (YAML) 만 작성하면, 도구가 알아서 모델을 다듬고 훈련시킵니다.
• 자동 최적화: "어떤 가지치기 방식이 가장 빠를까?", "어떤 정밀도가 가장 적을까?"를 자동으로 찾아주는 자동 실험실 역할을 합니다.

4. 실제 성과: "코끼리에서 토끼로"

이 도구를 실제 LHC 의 '제트 (입자) 분류' 작업에 적용해 보았습니다.

• 결과: 모델의 크기와 메모리 사용량은 대폭 줄어들었지만, 정확도는 거의 떨어지지 않았습니다.
• 속도: 기존 도구들 (QKeras 등) 보다 더 빠르고, 하드웨어 자원 (전력, 공간) 을 훨씬 적게 사용했습니다.
• 비유: 무거운 코끼리가 가볍고 빠른 토끼로 변신하여, 폭포수 속에서도 중요한 물방울을 놓치지 않고 골라냅니다.

5. 결론: 왜 이것이 중요한가?

PQuantML 은 과학자들이 실시간으로 중요한 데이터를 처리할 수 있게 해주는 핵심 열쇠입니다.

• 간단히 말해: "무겁고 느린 AI 를, 가볍고 빠른 AI 로 만들어서, 초고속 필터에 딱 맞게 설치해주는 도구"입니다.
• 이 도구를 통해 미래의 입자 가속기 실험은 더 많은 데이터를 더 빠르게 분석할 수 있게 되며, 이는 새로운 물리 법칙을 발견하는 데 큰 도움이 될 것입니다.

한 줄 요약:

PQuantML은 거대한 AI 모델을 하드웨어가 소화할 수 있도록 **가볍고 빠르게 다듬어주는 'AI 전용 미용실'**입니다.

기술 요약

1. 문제 정의 (Problem Statement)

대형 강입자 충돌기 (LHC) 와 같은 고에너지 물리 (HEP) 실험에서는 초당 수백 테라바이트의 데이터를 생성하며, 이를 모두 저장하는 것은 불가능합니다. 따라서 ATLAS 와 CMS 실험에서는 하드웨어 기반의 1 단계 트리거 (L1T) 시스템을 통해 마이크로초 (μs) 단위로 이벤트 선별을 수행합니다.

• 도전 과제: 기존 L1T 는 단순한 임계값 알고리즘을 사용하지만, 복잡한 물리 현상을 분석하기 위해 머신러닝 (ML) 모델을 도입하려는 시도가 늘고 있습니다. 그러나 ML 모델은 FPGA(Field-Programmable Gate Array) 와 같은 하드웨어에 배포될 때 엄격한 지연 시간 (latency) 및 리소스 제약 (DSP, LUT, FF 등) 을 충족해야 합니다.
• 기존 방법의 한계: 기존 도구 (QKeras, HGQ 등) 는 양자화 (Quantization) 에 초점을 맞추고 있어, 사용자가 프루닝 (Pruning) 전략을 별도로 구현해야 하는 번거로움이 있었습니다. 또한, 단순한 사후 처리 (Post-training) 방식은 높은 압축 비율에서 성능 저하를 초래하며, 하드웨어 인식 (Hardware-aware) 훈련이 제대로 통합되지 않은 경우가 많았습니다.

2. 방법론 (Methodology)

저자들은 PQuantML이라는 오픈소스 라이브러리를 제안하여 프루닝과 양자화를 통합된 프레임워크에서 처리하고, 하드웨어 제약을 고려한 엔드 - 투 - 엔드 워크플로우를 제공합니다.

핵심 아키텍처 및 기능

1. 통합 압축 프레임워크:

   • 프루닝 (Pruning): 구조화 (Structured), 비구조화 (Unstructured), N:M 반구조화 (Semi-structured) 등 다양한 세분화 (Granularity) 를 지원하는 여러 프루닝 알고리즘을 통합합니다.
   • 양자화 (Quantization): 고정 소수점 (Fixed-point) 양자화를 지원하며, 기존 HGQ(High-Granularity Quantization) 의 기능을 확장하여 가중치와 활성화 값의 비트 폭을 학습할 수 있도록 합니다.
   • 양자화 인식 훈련 (QAT): 양자화 노이즈를 훈련 과정에서 학습하여 성능 저하를 최소화합니다.

2. 유연한 모델 정의:

   • 직접 레이어 정의: PQuantML 의 전용 레이어 (PQConv2d, PQDense 등) 를 사용하여 모델을 구축.
   • 레이어 교체 (Layer Replacement): 기존 Keras/PyTorch 모델을 자동으로 압축 가능한 레이어로 교체하는 기능을 제공.

3. 자동화된 워크플로우:

   • 하이퍼파라미터 최적화: Optuna 를 활용한 베이지안 탐색을 통해 압축 전략 (프루닝 비율, 비트 폭 등) 을 자동 최적화합니다.
   • 실험 추적: MLflow 와 통합되어 설정, 메트릭, 모델 아티팩트를 체계적으로 기록합니다.
   • 다단계 훈련: 프루닝 방법 (예: 소프트 마스크 학습, 하드 마스크 고정) 에 따라 사전 훈련 (Pre-training), 훈련, 파인튜닝 (Fine-tuning) 단계를 자동으로 관리합니다.

4. 하드웨어 배포 연동:

   • 훈련된 모델을 hls4ml 라이브러리와 호환되도록 설계하여, FPGA 펌웨어 (HLS 코드) 로의 변환을 지원합니다.

3. 주요 기여 (Key Contributions)

• 통합 솔루션: 프루닝과 양자화를 별도의 단계가 아닌 하나의 통합된 프레임워크로 제공하여, 물리학자들이 복잡한 압축 워크플로우를 쉽게 구축할 수 있게 함.
• 하드웨어 인식 최적화: FPGA 의 리소스 (DSP, LUT) 와 지연 시간을 고려한 비트 폭 학습 및 프루닝 전략을 지원.
• 다양한 프루닝 알고리즘 지원: MDMM, AutoSparse, PDP, Wanda 등 다양한 프루닝 기법을 단일 인터페이스로 제공.
• 오픈소스 및 재현성: GitHub 에서 오픈소스로 제공되며, YAML 설정 파일을 통해 실험을 재현 가능하게 만듦.

4. 실험 결과 (Results)

저자들은 LHC 의 제트 서브구조 분류 (Jet Substructure Classification, Jet Tagging) 태스크를 벤치마크로 사용하여 PQuantML 의 성능을 검증했습니다.

• 데이터셋: 고수준 특징 (HLF) 과 입자 수준 특징 (PLF) 을 포함한 두 가지 데이터셋 사용.
• 비교 대상: 기존 도구인 QKeras 와 HGQ 와 비교.
• 성능 지표:
  • 정확도 (Accuracy): PQuantML 은 프루닝과 고정 비트 폭 양자화를 결합하여 QKeras 와 유사하거나 더 높은 정확도 (약 76% 대) 를 유지했습니다.
  • 리소스 효율성:
    • LUT/FF 감소: QKeras 대비 LUT 사용량을 약 30~40% 감소시켰습니다 (예: 5,504 → 3,895).
    • 지연 시간 (Latency): 지연 시간을 105ns 에서 약 47ns 로 단축하여 하드웨어 요구 사항을 충족했습니다.
    • DSP 사용량: FITCompress 알고리즘을 적용 시 DSP 사용량을 45 개까지 극도로 낮추는 데 성공했습니다.
  • HGQ 비교: PQuantML 로 훈련된 HGQ 기반 모델은 HGQ 라이브러리로 직접 훈련된 모델과 유사한 정확도와 리소스 사용량을 보여, PQuantML 이 HGQ 의 대체재로서 유효함을 입증했습니다.

5. 의의 및 결론 (Significance & Conclusion)

• 실시간 시스템 적용 가능성: PQuantML 은 LHC 의 1 단계 트리거 시스템과 같이 지연 시간과 리소스가 극도로 제한된 환경에서 ML 모델을 배포하기 위한 필수적인 도구로 자리 잡았습니다.
• 접근성 향상: 물리학자들이 복잡한 하드웨어 최적화 지식을 갖추지 않더라도, 설정 파일과 자동화된 워크플로우를 통해 고성능 압축 모델을 생성할 수 있게 되었습니다.
• 미래 전망: 현재는 TensorFlow/Keras 와 PyTorch 백엔드를 지원하지만, 향후 FITCompress 의 TensorFlow 통합, 더 복잡한 신경망 구조 (잔차 연결, 어텐션 등) 지원, 그리고 지식 증류 (Knowledge Distillation) 등의 추가 압축 기법 통합을 통해 기능을 확장할 계획입니다.

요약하자면, PQuantML 은 하드웨어 제약을 고려한 ML 모델 압축의 진입 장벽을 낮추고, LHC 와 같은 실시간 고에너지 물리 실험에서 효율적인 AI 추론을 가능하게 하는 통합 엔드 - 투 - 엔드 도구입니다.