I/O Lower-Bound Theory and Reinforcement Learning for Efficient Neural Network Inference Optimization
As artificial intelligence models grow in complexity, optimizing neural network inference has become a critical challenge. Existing approaches often rely on manual, expert-driven tuning tailored to specific hardware, which lacks scalability across diverse architectures. In this paper, we propose an automated optimization framework with a hierarchical two-layer tuning mechanism. At the node level (intra-operator), we introduce an I/O lower-bound theory based on the Red-Blue Pebble game and the ( X 1 , X 2 )-Partition theorem to guide tiling and memory-mapping configurations. At the graph level (inter-operator), we employ a reinforcement learning (RL) strategy to adaptively identify optimal operator fusion boundaries across network topologies. By synergizing theoretical I/O constraints with graph-level adaptive fusion while accounting for search overhead, the framework systematically explores high-performance execution patterns. For the TileAttn operator, the ( X 1 , X 2 )-Partition theorem raises DRAM flow estimation accuracy from 77.5%-82.0% under X -Partition to 86.3%-94.8%. Our method reduces shared memory traffic by 10.79% on average, achieves the best performance in 65.3% of cross-platform cases and top-two in 86.1%, and its DQN-based fusion engine outperforms greedy strategies in 91.67% of scenarios. We further analyze the algorithm’s overhead and its amortization break-even points.
Authors
- Huajian Zhang (ORCID: https://orcid.org/0000-0002-2445-6286)
- Qingmao Li (ORCID: https://orcid.org/0009-0004-7397-6802)
- Xiao‐Wei Guo (ORCID: https://orcid.org/0000-0003-4712-1320)
- Rui Xia (ORCID: https://orcid.org/0000-0003-1444-6778)
- Jie Liu (ORCID: https://orcid.org/0000-0003-3745-7541)
- Chuhe Hong
- Xing Peng (ORCID: https://orcid.org/0009-0008-2113-301X)
- Zhenhao Zhao (ORCID: https://orcid.org/0009-0004-5293-2907)
- Genglin Li (ORCID: https://orcid.org/0000-0003-3006-864X)
- Gencheng Liu (ORCID: https://orcid.org/0009-0007-5015-4736)
Institutions
- National University of Defense Technology (CN)
Publication Details
- Journal
- ACM Transactions on Architecture and Code Optimization
- Published
- 2026-09-09
- DOI
- https://doi.org/10.1145/3841638
- Primary Topic
- Advanced Neural Network Applications
- Type
- article
- Field-Weighted Citation Impact
- 0.00