X-RLflow: Graph Reinforcement Learning for Neural Network Subgraphs Transformation
Tensor graph superoptimisation systems perform a sequence of subgraph substitution to neural networks, to find the optimal computation graph structure. Such a graph transformation process naturally falls into the framework of sequential decision-making, and existing systems typically employ a greedy search approach, which cannot explore the whole search space as it cannot tolerate a temporary loss of performance. In this paper, we address the tensor graph superoptimisation problem by exploring an alternative search approach, reinforcement learning (RL). Our proposed approach, X-RLflow, can learn to perform neural network dataflow graph rewriting, which substitutes a subgraph one at a time. X-RLflow is based on a model-free RL agent that uses a graph neural network (GNN) to encode the target computation graph and outputs a transformed computation graph iteratively. We show that our approach can outperform state-of-the-art superoptimisation systems over a range of deep learning models and achieve by up to 40% on those that are based on transformer-style architectures.
Code (1)
Tasks
Decision MakingGraph Neural Networkreinforcement-learningReinforcement LearningReinforcement Learning (RL)Sequential Decision MakingMethods 이 논문이 사용한 방법론
Similar Papers 제목 키워드 기반
RLFlow: Optimising Neural Network Subgraph Transformation with World Models
Training deep learning models takes an extremely long execution time and consumes large amounts of computing resources. At the same time, recent research proposed systems and compilers that are expected to decrease deep …
Deep LearningModel-based Reinforcement Learningreinforcement-learningReinforcement Learning (RL)Controllable Flow Matching for Online Reinforcement Learning
Model-based reinforcement learning (MBRL) typically relies on modeling environment dynamics for data efficiency. However, due to the accumulation of model errors over long-horizon rollouts, such methods often face challe…
Reinforcement Learning5* Knowledge Graph Embeddings with Projective Transformations
Performing link prediction using knowledge graph embedding models has become a popular approach for knowledge graph completion. Such models employ a transformation function that maps nodes via edges into a vector space i…
Graph EmbeddingKnowledge Graph CompletionKnowledge Graph EmbeddingKnowledge Graph Embeddings+3ReGAE: Graph autoencoder based on recursive neural networks
Invertible transformation of large graphs into fixed dimensional vectors (embeddings) remains a challenge. Its overcoming would reduce any operation on graphs to an operation in a vector space. However, most existing met…
DecoderMAG-GNN: Reinforcement Learning Boosted Graph Neural Network
While Graph Neural Networks (GNNs) recently became powerful tools in graph learning tasks, considerable efforts have been spent on improving GNNs' structural encoding ability. A particular line of work proposed subgraph …
Combinatorial OptimizationGraph LearningGraph Neural Networkreinforcement-learning+2