Integer linear programming for unsupervised training set selection in molecular machine learning

TL;DR

The paper tackles the problem of selecting small, highly informative training sets for predicting size-extensive molecular properties. It introduces an ILP-based method that maps target atomic environments to a subset of fragments by minimizing a distance-based objective while enforcing one-to-one and fragment-diversity constraints, and it uses lazy constraints to ensure new molecules are added per solution; a penalty term discourages unnecessary extra atoms. Compared to baselines like SML, FPS, CUR, and Random, the ILP approach generally delivers superior extrapolation performance, especially when a penalty (p=1) is applied, enabling reliable predictions for molecules larger than those in the training set. The method demonstrates practical accuracy improvements for large-molecule property predictions and provides a flexible framework for incorporating additional constraints or representations in training-set design, with publicly available software for replication.

Abstract

Integer linear programming (ILP) is an elegant approach to solve linear optimization problems, naturally described using integer decision variables. Within the context of physics-inspired machine learning applied to chemistry, we demonstrate the relevance of an ILP formulation to select molecular training sets for predictions of size-extensive properties. We show that our algorithm outperforms existing unsupervised training set selection approaches, especially when predicting properties of molecules larger than those present in the training set. We argue that the reason for the improved performance is due to the selection that is based on the notion of local similarity (i.e., per-atom) and a unique ILP approach that finds optimal solutions efficiently. Altogether, this work provides a practical algorithm to improve the performance of physics-inspired machine learning models and offers insights into the conceptual differences with existing training set selection approaches.

Figures (10)

• Figure 1: Outline of our approach to identify an optimal training set from an existing database (here, QM7rupp2012fast), for a dedicated target molecule (here, sildenafil). Target and fragment molecules are represented with atom-centered feature vectors, $\mathbf{T}_i$ and $\mathbf{M}_j$ respectively, where $i$ and $j$ are the corresponding atom indices. With ILP, we find solutions as combinations of atoms in fragment molecules that best match the target atoms. Here, the top three solutions (smallest objective value) are shown. After removing duplicates, and labelling molecules with properties if necessary, an optimal training set is obtained. puck_thesis
• Figure 2: Weighted bipartite graph corresponding to a target $T$ and database with two elements $D = \{M, M'\}$. For each fragment $M \in D$ and for each atom $i \in T$ and $j \in M$ with same atom types, the edge between $i$ and $j$ has cost $w_{ij}^{(M)} = \Vert \mathbf{T}_i - \mathbf{M}_j \Vert_2^2$, where $\mathbf{T}_i$ and $\mathbf{M}_j$ are the representations of each atom.
• Figure 3: Example of linear programs in two dimensions. Left: The point $\mathbf{x}^*_\mathrm{frac}$ maximizes $\mathbf{c}^\top \mathbf{x}$ for points $\mathbf{x}\equiv(x,y)$ restricted to the grey area. The point $\mathbf{x}^*_\mathrm{int}$ is a maximum within the grey area intersected with the $\mathbb{Z}^2$ grid. Right: An additional constraint $(C)$ is added to the ILP algorithm from the left, making the old maximum $\mathbf{x}^*_\mathrm{old}$ infeasible, leaving $\mathbf{x}^*_\mathrm{new}$ as the only integer maximizer of $\mathbf{c}^\top \mathbf{x}$.
• Figure 4: 2D structures of the ten drug molecules used as targets. From left to right; 1st row: oxycodone, penicillin, troglitazone, imabinib; 2nd row: pregabalin, oseltamivir, sildenafil, pemetrexed; 3rd: salbutamol, apixaban.
• Figure 5: Density plot (i.e., $d = \frac{c}{\Delta b \sum{c}}$, where $d$ is density, $c$ is counts, and $\Delta b$ is the discrete difference between the number of bins, as implemented in matplotlib.pyplot.hist with density=True) of energy correction $\hat{E}$ of QM7 molecules, ten random QM9 molecules (QM9*), and ten drug molecules, computed using PBE0-D3/def2-SVP. The smooth curves are continuous fits to the histograms using kernel density estimates, as implemented in sns.kdeplot.