Neural Architecture Codesign for Fast Physics Applications

TL;DR

This paper presents a pipeline combining neural architecture search and model compression to efficiently design hardware-aware neural networks for physics applications, enabling faster deployment on FPGA with improved performance.

Contribution

It introduces a two-stage neural architecture codesign process with hierarchical search space, tailored for physics tasks, and demonstrates its effectiveness on materials science and high energy physics applications.

Findings

Achieved improved accuracy and reduced latency in physics tasks

Enabled FPGA deployment with optimized resource utilization

Demonstrated flexibility of the hierarchical search space

Abstract

We develop a pipeline to streamline neural architecture codesign for physics applications to reduce the need for ML expertise when designing models for novel tasks. Our method employs neural architecture search and network compression in a two-stage approach to discover hardware efficient models. This approach consists of a global search stage that explores a wide range of architectures while considering hardware constraints, followed by a local search stage that fine-tunes and compresses the most promising candidates. We exceed performance on various tasks and show further speedup through model compression techniques such as quantization-aware-training and neural network pruning. We synthesize the optimal models to high level synthesis code for FPGA deployment with the hls4ml library. Additionally, our hierarchical search space provides greater flexibility in optimization, which can easily extend to other tasks and domains. We demonstrate this with two case studies: Bragg peak finding in materials science and jet classification in high energy physics, achieving models with improved accuracy, smaller latencies, or reduced resource utilization relative to the baseline models.