WARGM-PDPG: A dual-phase policy gradient graph mamba neural network for device placement algorithms

Area of Science:

• Artificial Intelligence
• Machine Learning
• Deep Learning

Background:

• Increasing neural network complexity strains device computational limits.
• Efficient parallel training and optimal resource allocation are crucial for large-scale AI.
• Device placement is a critical bottleneck in distributed AI training.

Purpose of the Study:

• To propose an efficient strategy for optimizing device placement in large-scale neural network training.
• To address the computational constraints of complex neural networks.
• To improve the efficiency and performance of AI model training.

Main Methods:

• Integration of Weight Adaptive Residual Graph Mamba Network (WARGMamba) for superior graph embedding.
• Utilization of Perturbation-Based Dual-Phase Policy Gradient Algorithm (PDPG) for device placement optimization.
• WARGMamba mitigates over-smoothing and balances node features; PDPG employs dual-phase updates for short-term and long-term optimization.

Main Results:

• The WARGM-PDPG strategy significantly outperforms existing methods.
• Achieved a 95.89% reduction in device placement computation time.
• Demonstrated a 10.24% to 28.66% improvement in execution time.

Conclusions:

• The proposed WARGM-PDPG strategy offers a highly effective solution for device placement optimization.
• This method enhances the efficiency and scalability of training large neural networks.
• WARGM-PDPG represents a significant advancement in distributed AI training methodologies.