Parameter Averaging for Feature Ranking

TL;DR

This paper introduces XTab, a robust feature importance estimation method for tabular data that uses parameter averaging of multiple neural network instances to improve consistency over traditional single-model approaches.

Contribution

The paper proposes a novel parameter averaging technique for neural networks to achieve more reliable feature importance rankings in tabular data.

Findings

Parameter averaging improves feature importance stability.

XTab outperforms individual models in identifying ground-truth features.

Method is effective on synthetic and real-world datasets.

Abstract

Neural Networks are known to be sensitive to initialisation. The methods that rely on neural networks for feature ranking are not robust since they can have variations in their ranking when the model is initialized and trained with different random seeds. In this work, we introduce a novel method based on parameter averaging to estimate accurate and robust feature importance in tabular data setting, referred as XTab. We first initialize and train multiple instances of a shallow network (referred as local masks) with "different random seeds" for a downstream task. We then obtain a global mask model by "averaging the parameters" of local masks. We show that although the parameter averaging might result in a global model with higher loss, it still leads to the discovery of the ground-truth feature importance more consistently than an individual model does. We conduct extensive experiments on a variety of synthetic and real-world data, demonstrating that the XTab can be used to obtain the global feature importance that is not sensitive to sub-optimal model initialisation.