Automating Feature Subspace Exploration via Multi-Agent Reinforcement Learning


Feature selection is the preprocessing step in machine learning which tries to select the most relevant features for the subsequent prediction task. Effective feature selection could help reduce dimensionality, improve prediction accuracy and increase result comprehensibility. It is very challenging to find the optimal feature subset from the subset space as the space could be very large. While much effort has been made by existing studies, reinforcement learning can provide a new perspective for the searching strategy in a more global way. In this paper, we propose a multi-agent reinforcement learning framework for the feature selection problem. Specifically, we first reformulate feature selection with a reinforcement learning framework by regarding each feature as an agent. Then, we obtain the state of environment in three ways, i.e., statistic description, autoencoder and graph convolutional network (GCN), in order to make the algorithm better understand the learning progress. We show how to learn the state representation in a graph-based way, which could tackle the case when not only the edges, but also the nodes are changing step by step. In addition, we study how the coordination between different features would be improved by more reasonable reward scheme. The proposed method could search the feature subset space globally and could be easily adapted to the real-time case (real-time feature selection) due to the nature of reinforcement learning. Also, we provide an efficient strategy to accelerate the convergence of multi-agent reinforcement learning. Finally, extensive experimental results show the significant improvement of the proposed method over conventional approaches.

Meeting Name

25th ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, KDD 2019 (2019: Aug. 4-8, Anchorage, AK)


Computer Science

Second Department

Electrical and Computer Engineering

Research Center/Lab(s)

Center for Research in Energy and Environment (CREE)

Keywords and Phrases

Automated exploration; Feature selection; Multi-agent reinforcement learning

International Standard Book Number (ISBN)


Document Type

Article - Conference proceedings

Document Version


File Type





© 2019 Association for Computing Machinery (ACM), All rights reserved.

Publication Date

01 Aug 2019