Learning tree structures from leaves for particle decay reconstruction

To satisfy the problem requirements outlined above, we employ a modified representation of the lowest common ancestor (LCA) matrix representation of rooted trees [37]. Given two nodes in a graph, a and b, the LCA is defined as the deepest node that is an ancestor of both a and b, deepest in this case meaning farthest graph distance, $\textrm{dist}$, from the root node r, i.e.:

$$\begin{align*} \textrm{LCA}(a,b) = \textrm{argmax}_{p} \{\textrm{dist}(p,r) |\, &p\in\textrm{path}(a,r)\land p\in \textrm{path}(b,r)\}.\end{align*}$$

Figure 2 shows an example tree representing an arbitrary particle decay process (left) and its corresponding LCA matrix (top right), with generations up from the leaves labelled on the far left. Letters indicate arbitrary node labels, and the colours are simply added for clarity and do not hold further meaning ⁹ . In order for this representation to describe a graph unambiguously, it must be a tree where any non-leaf node has two or more children. This way the set of nodes that are elements of the LCA matrix is identical to $\mathcal{V}$. This assumption bounds the maximum depth of the tree to $\mathcal{O}(\log_2 k)$ for k leaves. Note that each entry in the LCA matrix encodes information about an unordered pair of vertices, and as such is symmetric.

Zoom In Zoom Out Reset image size

In order to express the LCA matrix in a form appropriate for machine learning, we propose the LCAG matrix, in which each entry of the LCA matrix is replaced with its corresponding generation in the tree. Figure 2(bottom right) shows the LCAG corresponding to the example tree. We see here a clear distinction between the LCA and LCAG: while the LCA requires unique node labels to identify ancestors, the LCAG is a relaxed form which allows the same identifier (generation) to represent multiple ancestors. In other words, the LCA is surjective to the LCAG. The underlying assumption for making this modification is that if there exist structural rules which allow the inference of the tree structure from the leaf nodes alone, then these rules can also be used on the inferred structure to deduce unseen node labels and hence recover the LCA. Relating this to the example particle decay process in figure 2, if we know that two particles $\textrm{v}$ and $\textrm{w}$ come from the same parent, summing these energies would tell us the mass, and hence the node label of the parent $\textrm{B}$.

For the LCAG matrix, we adopt a pull-down convention: every node in the tree resides at the lowest possible generation, with all leaves fixed at the zeroth generation. Formally:

$$\begin{align*} \textrm{LCAG}_{ij}(\mathcal{G}) & = g(\textrm{LCA}_{ij}(\mathcal {G})), \\ g(u) & = \left\{\begin{array}{l@{\quad}l} 0, & \textrm{if u is a leaf} \\ \max\limits_{v \in \mathcal{C}(u)} \{g(v)\} + 1, & \textrm{otherwise} \end{array}\right. \end{align*}$$

where $\mathcal{C}(u)$ is the set of children of node u. Lemma 1.

lemma The LCAG matrix uniquely describes any rooted tree isomorphism class for trees in which every non-leaf vertex has two or more children. I.e. if two trees $G = (V_G, E_G)$ and $H = (V_H, E_H)$ are described by the LCAG L, then G and H are isomorphic ($G \simeq H$).

We prove this by complete induction. Lemma 1 is trivially true for trees consisting of a single root node, described by the LCAG $\left[ 0\right]$. We will now show that it also holds for any tree of height m + 1 if it holds for trees of height m. Consider two trees G and H of height m + 1, that are described by the LCAG L. Given any subtree G_i of G that includes all descendants of $\textrm{root}(G_i)$ in V(G), G_i is described by L_i containing all entries of L corresponding to the leaves of G_i . The L_i of all subtrees whose roots are siblings in G and are direct descendants of root(G) must be disjoint in L, since their LCA describes the parent of the roots of all G_i in G. The remaining entries of L not in any L_i are all by definition m + 1. The same set of L_i correspond analogously to subtrees of H. Therefore, for each G_i we can find a corresponding subtree H_i that is described the same L_i in such a way that $\dot\bigcup_i V(G_i) \cup \textrm{root}(G) = V(G)$ and $\dot\bigcup_i V(H_i) \cup \textrm{root}(H) = V(H)$. Because of lemma 1, G_i and H_i are isomorphic, as their height is m and there is an isomorphism $f_i: V(G_i) \rightarrow V(H_i)$. We can find an isomorphism $f:V(G) \rightarrow V(H)$ by $f(\textrm{root}(G)) = \textrm{root}(H)$ and $f(u) = f_{i}(u)$ for $u \in V(G_i)$. Therefore, lemma 1 also holds for trees of height m + 1.

In order to generate the adjacency matrix of a tree given its LCAG matrix, we proceed level by level, starting from the leaves. We generate a list of leaves from the dimensions of the input LCAG. We then compute the vertices in the next level by collecting all sets of siblings in the list, i.e. vertices whose LCA is in the next level. For each set of siblings, we remove them from the list and add their parent vertex to the list, keeping track of the LCA of the newly added vertices and all other vertices in the list. Algorithm 1 shows the pseudocode of the conversion algorithm.

Algorithm 1. Pseudocode for the conversion of a LCAG gram matrices L into a tree T, error handling for ill-formatted matrices omitted. Indices and ranges are zero-indexed.
Input: Symmetric LCA(G) gram matrix L (n × n)
Output: Root of tree T described by L
$leaves \leftarrow [~]$
$levels \leftarrow sorted(unique(L))$
$total\_nodes \leftarrow n$
for all n do $leaves.append(Node(level = 0))$
{level by level}
for level in $[0,\ldots,levels.length]$ do
for column in $[0,\ldots, n]$ do
for row in $[column+1,\ldots, n]$ do
if $L[row, column]\neq level$ then continue
$node\_a \leftarrow leaves[row]$
$node\_b \leftarrow leaves[column]$
$root\_a \leftarrow get\_root(node\_a)$
$root\_b \leftarrow get\_root(node\_b)$
{same level, new node}
if $root\_a.level\lt level+1$ and $root\_b.level\lt level+1$
then
$parent \leftarrow Node(level = level+1)$
$parent.children.append([node\_a, node\_b])$
$total\_nodes \leftarrow total\_nodes +1$   root_b is older
else if $root\_a.level \lt level + 1$ and $root\_b = = level + 1$
then
$root\_a.children.append(root\_b)$ vice versa
else
$root\_b.children.append(root\_a)$
end if
end for
end for
end for
$root\gets get\_root(leaves[0])$

To generate the LCAG from an adjacency matrix, we first find the set of leaves, and then fill the values of the matrix with the longest graph distance to the LCA of the corresponding pair of leaf vertices.

Both the LHC and Belle II simulated collision data are restricted to internal member access only, and existing, publicly accessible benchmarks such as the TrackML challenge [38] and the PD4ML [39] focus only on identifying individual particles or whole-graph classification tasks, not any form of decay tree reconstruction. Therefore, we simulate our own experimental dataset using the Phasespace library [40]. This library takes as input a tree structure in the form of parent-daughter relations, as well as the mass of each participating particle (node). Decays are then simulated according to Monte Carlo phase space [41], which obey the usual physical laws (conservation of energy, momentum, etc). The output is the four-momentum (energy and 3D momentum) of each participating particle, of which we use the leaf nodes for input to the network. In a high-energy particle physics experiment, these leaf nodes correspond to particles detected by the experimental hardware, which a physicist would then use to reconstruct the originating decay. We refer to different tree structures as topologies, and we refer to individual simulated trees (of any topology) as samples.

The simulated datasets used by particle physics experiments for building similar types of reconstruction algorithms include only known topologies. For the type of usage of such algorithms, it is intentional that only those topologies that are well understood are reconstructed to avoid discrepancies between simulation (i.e. training) and measured data (i.e. inference). Therefore, it is reasonable to expect all possible topologies intended for reconstruction be present in the training dataset, with no generalisation to new, unseen topologies.

We simulate a synthetic particle decay dataset for our experiments, consisting of topologies with a common root particle of mass 100 (arbitrary units). Intermediate particles are selected at random with replacement from the following masses: $[90, 80, 70, 50, 25, 20, 10]$. Final state particles, which make up the leaf nodes of generated topologies, are drawn with replacement from the following masses: $[1, 2, 3, 5, 12]$. For each intermediate particle (including the root), we limit the minimum number of children to two, and the maximum five. These masses and ranges of child particles are chosen to correspond to the relative mass scales and decay multiplicities observed in nature [42]. For particle physics decay processes seen at the Belle II experiment, for example, the number of leaves is typically around ten or fewer, with existing decay reconstruction algorithms reflecting this ¹⁰ .

Tree topology creation is then as follows: starting from the root particle a set of children are selected from the available intermediate and final state particles such that the sum of their masses totals less than the root, this process is then repeated for each child particle which is not a final state particle and so on until only final state particles remain.

The synthetic dataset consists of 200 topologies in total, with 2000 samples per topology for each of training, validation, and testing. Figure 5(top) shows the distribution of simulated topologies, grouped according to the number of leaves and the total depth of the tree, which we use as a surrogate for topology complexity. All leaf node features are normalized to a normal distribution of mean zero, standard deviation one. We do not enforce any ordering of the nodes and leave them unsorted as created in the dataset.

For the choice of neural network models, we require architectures that operate on unordered sets of vertices. The architectures must be able to learn the inter-dependencies between vertices and model the underlying dynamics of the physical system. We therefore select the encoder components of both the Transformer and NRI models. All model implementations were made in PyTorch v1.8.1 [43] and trained using PyTorch-Ignite v0.4.4 [44].

4.2.1. Transformer

4.2.2. NRI

NRI was originally proposed by [10] as an unsupervised GNN in the form of a variational autoencoder. In our experiments we use only the encoder component of the NRI, as we are performing supervised training. The building blocks of an NRI layer consist of node-to-edge and edge-to-node message passing modules, followed by a fully-connected module each. The latent representation then alternates between node-wise and edge-wise as it propagates through the network. The encoder produces a discrete categorical distribution of the edge label predictions, which in our case are the entries of the LCAG. Each LCAG entry prediction is treated as an individual classification task with a softmax followed by an argmax function after the last layer. The individual losses are averaged to produce the total loss.

Figure 3 shows the complete NRI encoder architecture used in experiments. To update the state embeddings between edge-to-node and node-to-edge transition layers, the NRI employs sequences of multilayer perceptrons (MLPs) containing two linear layers with ELU activations [45] and batch normalisation [46]. We extend the original NRI implementation to include a configurable number of additional MLPs both within the NRI blocks and at the beginning and end of the model to allow for added learning capacity.

Zoom In Zoom Out Reset image size

The neural message-passing mechanism in the NRI, used in the edge-to-node and node-to-edge layers, updates state embeddings by aggregating information across all connected edges in the graph (we adopt the notation of [10]):

$$\begin{align*} \textrm{node-to-edge} &: \textbf{h}^{l}_{(i,j)} = {f\,\,}_{e}^{l} \left( \left[ \textbf{h}_i^l, \textbf{h}_j^l \right] \right) \\ \textrm{edge-to-node} &: \textbf{h}^{l+1}_i = {f\,\,}_{v}^{l} \left( \sum_{j} \textbf{h}_{(i,j)}^l \right), \end{align*}$$

where $\textbf{h}^l_i$ represents the state embedding of node i at layer l, and $\textbf{h}^l_{(i,j)}$ the embedding of the edge connecting nodes i and j. f_e and f_v are feed-forward neural networks used to produce the edge and node embeddings. In contrast to [10], we allow for self-interactions $\left(\textbf{h}_{(i,i)}^l\right)$ in our experiments as we found this empirically to perform better.

Similar to (non-ordinal) classification tasks, a predicted LCAG has no clear concept of an almost correct tree, for example in the case when only one entry prediction is incorrect. To see why this is the case, we can look at an example of what we would intuitively identify as two similar trees, i.e. almost isomorphic, and see that they are not related by valid LCAG representations. Consider the simple operation of exchanging the parent of the nodes $\textrm{x}$ and $\textrm{v}$ in the example from figure 2. Figure 4 demonstrates how this alters two LCAG entries in the process (considering only the upper or lower triangle due to symmetry). Changing only one of the LCAG entries would result in what we consider an invalid LCAG: one that does not represent a coherent tree structure. We therefore introduce the metric $\mathrm{perfectLCAG}$ to evaluate the rate of correct LCAG predictions, which is the ratio of correctly predicted LCAGs to the total number of samples.

Zoom In Zoom Out Reset image size

We optimized and evaluated both the Transformer and the NRI encoder on the simulated dataset described above (cf section 4.1). To find the optimal hyperparameters for both models, we used the Optuna v2.10 [47] library with a Tree-structured Parzen Estimator [48] optimisation approach. The objective value maximized is the $\mathrm{perfectLCAG}$ score on validation samples. All trainings were performed on a single NVIDIA A100 GPU with 40 GB memory.

We performed an initial hyperparameter search on the parameters shown in table 1. Namely, we explored how the network performance changes when varying the building blocks, i.e. the attention layers in the Transformer or the blocks in the NRI, as well as the overall capacity of the models, i.e. the number of neurons in fully-connected layers, and additional initial/final MLPs. As this is a multi-target classification task, i.e. one target per LCAG entry, where getting all LCAG entry predictions correct is more important than getting only a subset correct with a high confidence (as cross-entropy loss does) we also included focal loss [49] in the search.

All hyperparameter search training were performed for a maximum of 100 epochs, a batch size of 128, dropout rate of 0.3, random seed of 100, and with class weights applied, using the Adam optimizer with a learning rate of 0.001.

The results of the searches, shown for each individual hyperparameter, are shown in the supplementary material in figures 1 and 2. Our experiments show the NRI to significantly outperform the Transformer encoder, reaching an average validation $\mathrm{perfectLCAG}$ of around $46\%$, compared to the Transformer's $11\%$. No significant deviations are observed between the training and the validation scores. Two NRI blocks are consistently found to produce optimal results, with larger feed-forward dimensions, favouring 512, and no additional block MLPs within.

Based on these results, we performed a longer training for both the NRI and Transformer encoders using the optimal hyperparameters found and a random seed of 101 until the validation $\mathrm{perfectLCAG}$ score no longer improved (500 and 323 epochs, respectively). The results on the test dataset are shown in figure 5, where the NRI and Transformer achieved an average $\mathrm{perfectLCAG}$ of $46.6\%$ and $5.80\%$, and a corresponding average accuracy of $91.6\%$ and $38.6\%$, respectively. As with the hyperparameter optimisation, we observe no significant difference between the training and validation $\mathrm{perfectLCAG}$ scores, around $54\%$ and $10\%$ for the NRI and Transformer, respectively, but a clear sign of overfitting with respect to the test dataset. Overall, the Transformer struggles to learn more than the most trivial topologies, whereas the NRI achieves reasonable $\mathrm{perfectLCAG}$ scores on most of those below 10 leaves. Looking closer, for trees up to and including 6 leaves, the average $\mathrm{perfectLCAG}$ for the NRI is $92.5\%$, and for cases of trees up to and including 10 leaves $59.7\%$. The scores for all other subdivisions are shown in table 2. This result demonstrates not only that the NRI is an appropriate choice of architecture for learning physical interactions for a variety of tree sizes, but also that it behaves as expected, namely by learning what we intuitively identify as simpler topologies first.

Zoom In Zoom Out Reset image size

Table 2. Average $\mathrm{perfectLCAG}$ scores for subsets of topologies of increasing complexity. Bold numbers indicate the best score at each complexity.

	Average perfectLCAG (%) score for leaves up to and including
Model	2	3	4	5	6	7	8	9	10	11	12	13	14	15	16
Transformer	100.0	63.1	46.3	30.5	24.7	17.1	13.0	9.8	7.7	6.8	6.3	6.1	6.0	5.9	5.8
$\mathrm{NRI}$	100.0	98.8	98.2	96.2	92.5	85.7	79.2	68.5	59.7	54.3	50.3	48.8	48.1	47.6	46.6
$\mathrm{NRI}_{2048}$	100.0	99.3	98.7	97.5	94.2	87.1	79.7	70.1	60.9	55.7	51.5	50.0	49.2	48.7	47.7
$\mathrm{NRI}_{\mathrm{no blocks}}$	100.0	91.2	69.1	46.0	37.2	25.7	19.6	14.7	11.6	10.3	9.4	9.1	9.0	8.9	8.7

Interestingly, we observe no significant advantage to an increased number of topologies at a given complexity. For example the large number of topologies at depth 5 with leaves 9 or 10 in the dataset show no noticeable performance improvement over those at depth 6 with the same number of leaves. This indicates that the model is only able to learn each topology individually, and is unable to generalize to topologies of similar complexity.

Given the significantly better performance of the NRI over the Transformer, we explored the importance of the sequential edge-to-node and node-to-edge encoding layers within the NRI blocks. To do so, we repeated the previous training with all blocks removed, leaving only the initial node-to-edge encoding layer to infer the LCAG from pair-wise combinations of the embeddings alone. The results are shown in figure 6(a), which reached an average test $\mathrm{perfectLCAG}$ of $8.7\%$, only slightly outperforming the Transformer. We therefore conclude that the ability to iteratively combine node pairs is an essential feature of the NRI's ability to predict the LCAG, and similar features should be considered in any other GNN approaches explored in future.

Zoom In Zoom Out Reset image size

The hyperparameter optimisations demonstrated a clear preference for larger feed-forward dimensions and a lower number of initial/final feed-forward layers. We therefore performed an additional hyperparameter search on these parameters up to a maximum feed-forward width of 2048, ¹² with all other hyperparameters set to the optimal found in the previous search. The results are shown in figure 3 in the supplementary material, where no additional initial/final layers were found to be ideal, along with the maximum feed-forward width.

Following this finding, we repeated the training for 177 epochs (until validation $\mathrm{perfectLCAG}$ no longer improved) with these hyperparameters and the same random seed as in the trainings in section 4.4. The results of are shown in figure 6(b). We observe an average $\mathrm{perfectLCAG}$ of $47.7\%$ ($94.2\%$ and $60.9\%$ for the ${\leqslant} 6$ and ${\leqslant}10$ regions, respectively) and accuracy of $92.6\%$, which represents only a small improvement over the previous results. Detailed results are shown in table 2. We see a relatively large discrepancy between this score and the maximum hyperparameter optimisation $\mathrm{perfectLCAG}$ reached in figure 3 in the supplementary material of $59.6\%$ on the validation dataset. Similar to the experiments in section 4.4, we also find a comparable discrepancy in the validation $\mathrm{perfectLCAG}$ score of the 177 epoch training, which achieved $58.4\%$.

To summarise, with a width of 1024, the model achieved $46.6\%$ test and $46\%$ validation $\mathrm{perfectLCAG}$, and for 2048 with all else equal it reached $47.7\%$ test and $59.6\%$ validation $\mathrm{perfectLCAG}$. We therefore conclude that the hyperparameter optimisation has caused overfitting to the validation dataset, and that an increase in the fully-connected widths within the NRI only serves to exacerbate this problem.

While the Belle II experiment typically has ten or fewer leaves, the noisy experimental environments in hadron colliders like the LHC mean that many other particles can be produced, raising the total number of leaves significantly. However, it is often the case that only a subset of leaves that belong to the particle decay of interest are selected. For example, the LHCb triggers [50, 51] have a specific focus on filtering for events with b-hadrons whose decay products are subsequently analysed. As our results demonstrate high performance for scenarios with up to ten leaves, our proposed approach already offers utility in existing experimental scenarios. Nevertheless, there is still significant room for improvement, the key to which we believe will be finding exactly which hyperparameters can be modified to improve performance for more leaves.

We have already explored some of the hyperparameters in the experiments and shown the essential role of the NRI blocks (figure 6). We also demonstrated that a larger model capacity (larger width of feed-forward layers) corresponds to improved performance (table 2), but saw that this quickly leads to overfitting. While a further investigation into this overfitting is outside the scope of this work, as our primary interest was to perform an initial exploration into what appear to be the essential NRI hyperparameters, based on these findings we recommend considering the inclusion of cross-validation, stronger regularization, or ensemble methods to identify and mitigate this.

Looking further, we saw that the choice of architecture has the largest impact on performance, with the Transformer appearing unable to learn a meaningful solution (figure 5). We concluded in section 4.5 that the NRI's ability to iteratively combine node pairs was an essential component for its performance, however graph nodes do not exist only as child pairs, but also triplets and more. Therefore, a broad investigation into various other architectures, especially GNNs, with potentially better inductive biases is likely the most promising avenue for increasing performance to handle more leaves and larger trees.