Multiple Instance Learning with Trainable Decision Tree Ensembles

TL;DR

This paper introduces STE-MIL, a novel random forest-based model for Multiple Instance Learning on small tabular datasets, utilizing soft decision trees, neural network approximations, and attention mechanisms for end-to-end training.

Contribution

It proposes a new soft decision tree model converted into neural networks and integrated with attention for MIL, trained end-to-end, with demonstrated effectiveness on tabular data.

Findings

Effective on small tabular datasets

End-to-end trainable model combining trees and neural networks

Code is publicly available for reproducibility

Abstract

A new random forest based model for solving the Multiple Instance Learning (MIL) problem under small tabular data, called Soft Tree Ensemble MIL (STE-MIL), is proposed. A new type of soft decision trees is considered, which is similar to the well-known soft oblique trees, but with a smaller number of trainable parameters. In order to train the trees, it is proposed to convert them into neural networks of a specific form, which approximate the tree functions. It is also proposed to aggregate the instance and bag embeddings (output vectors) by using the attention mechanism. The whole STE-MIL model, including soft decision trees, neural networks, the attention mechanism and a classifier, is trained in an end-to-end manner. Numerical experiments with tabular datasets illustrate STE-MIL. The corresponding code implementing the model is publicly available.