CrossNAS: A Cross-Layer Neural Architecture Search Framework for PIM Systems

2025-05-28

AI Summary

Cross-layer co-optimization of circuits, architectures, and systems remains challenging in analog in-memory computing (PIM) due to abstraction-level fragmentation. Method: This paper proposes the first end-to-end neural architecture search framework tailored for PIM (PIM-NAS), integrating single-path one-shot weight sharing with a multi-objective evolutionary algorithm. It introduces a PIM-aware latency/energy model and a cross-layer coupled search space to jointly optimize neural network topology and hardware mapping strategies. Contribution/Results: Unlike conventional hierarchical optimization, PIM-NAS bridges abstraction gaps to enable holistic circuit–architecture–system exploration. Experiments demonstrate that the method achieves significantly higher accuracy and energy efficiency than existing PIM-adaptation approaches, while incurring comparable or lower search overhead—establishing a new benchmark for PIM-aware NAS.

Technology Category

• Search and Optimization: Sampling/Simulation-based Search
• Machine Learning: Hardware-aware ML
• Cognitive Modeling & Cognitive Systems: Agent Architectures

Application Category

• Search and Retrieval-Augmented AI: Efficiency and scalability of Web search engines
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Economics, Online Markets and Human Computation: Incentives in network design for Web infrastructures and ecosystems

Abstract

In this paper, we propose the CrossNAS framework, an automated approach for exploring a vast, multidimensional search space that spans various design abstraction layers-circuits, architecture, and systems-to optimize the deployment of machine learning workloads on analog processing-in-memory (PIM) systems. CrossNAS leverages the single-path one-shot weight-sharing strategy combined with the evolutionary search for the first time in the context of PIM system mapping and optimization. CrossNAS sets a new benchmark for PIM neural architecture search (NAS), outperforming previous methods in both accuracy and energy efficiency while maintaining comparable or shorter search times.

Problem

• Automated exploration of multidimensional search space for PIM systems
• Optimizing machine learning workloads on analog PIM systems
• Enhancing accuracy and energy efficiency in PIM neural architecture search

Innovation

• Automated multidimensional search space exploration
• Single-path one-shot weight-sharing strategy
• Evolutionary search for PIM optimization