Neural Architecture Codesign for Fast Physics Applications

2025-01-09

AI Summary

To address the high dependency on domain experts and low hardware-adaptation efficiency in deploying neural networks for physics applications, this paper proposes an end-to-end neural architecture co-design methodology. We introduce a novel two-stage search framework comprising global hardware-aware neural architecture search (NAS) followed by local compression optimization, supporting a hierarchical and scalable search space tailored to diverse physics tasks. Our approach jointly integrates NAS, quantization-aware training, structured pruning, and HLS4ML-based high-level synthesis to enable fully automated FPGA code generation. Evaluated on Bragg peak localization and high-energy physics jet classification, the resulting models achieve +1.2–2.8% accuracy gains, −37–51% latency reduction, and −29–44% LUT/BRAM resource savings over baselines. This work significantly advances automation and efficiency in physics-driven AI deployment.

Technology Category

• Machine Learning: Hardware-aware ML
• Computer Vision: Low Level & Physics-based Vision
• Search and Optimization: Learning to Search

Application Category

• Search and Retrieval-Augmented AI: Vertical and domain-specific search
• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications
• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs

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.

Problem

• Neural Network Design
• Physical Applications Optimization
• Machine Learning Accessibility

Innovation

• Neural Architecture Search
• Model Compression
• Hardware-friendly Design