Path Signature Neural Network of Cortical Features for Prediction of Infant Cognitive Scores

A. Analysis of Infant Cognitive Scores

In the literature, there has been intensive researches for infant brain development analysis [23], [24]. However, few studies correlate cognition functions with the longitudinal development of the brain. Some researchers used the statistical analysis framework [6], [25]–[27]. For instance, in [6], the authors intended to determine the association between cortical thickness, surface area and measures of cognitive ability through mediation analysis and longitudinal analysis. Factor analysis has also been used to qualitatively evaluate the impact of different brain features on cognition function and illustrate the high correlation between the anatomical differentiation of white matter fibers and cognitive development [25]. Nevertheless, for the lack of quantitative formulation, they are incapable to predict the cognitive scales for individuals directly.

To enable individual prediction, recent advances of cognitive scores prediction are mainly based on machine learning methods [12], [13] with challenges of the missing data and the trade-off between large dimensional features and the SSS problem. Specifically, in [13], a latent representation is generated for each participant to leverage the complementary information among different time-point. A set of indicators is also introduced to eliminate the loss brought by the incomplete data. Nonetheless, in their formulation, the cortical morphological features extracted from 70 brain regions of each MRI scan are flattened into a 490 dimensional vector, which may cause the loss of structural information. Meanwhile, in [12], researchers proposed a Bag-of-Words [28] based model to decrease the number of features at an early stage with the aim of balancing the overlarge dimensionality of neuroimaging features and limited sample size. They achieved encouraging performance, but the linear transformation they used to gather features of different time points may be incapable to sufficiently explore the dynamic patterns in growth trajectories of brain regions.

B. Path Signature Method

Path signature is an infinite graded sequence of statistics known to characterize data streams which can also be regarded as paths. As a type of noncommutative formal power series, it originated from Chen’s study [29] who works on piece-wise smooth path. Lyons first used it to make sense of controlled differential equations driven by very rough paths [16]. The uniqueness of signature was initially proved by [30] for paths of bounded variation and it was extended to paths of finite p-variation for any p > 1 in [31]. Recently, the applications of path signature in machine learning is an emerging areas; in particular, the path signature can be used as an effective feature extractor in various tasks [17]–[21]. Specifically, in [32], Diehl applied path signature transformation onto handwritten curves for recognition. Yang pioneered in constructing paths in spatial domain by connecting skeleton nodes in a single frame for skeleton-based human action recognition and obtained great results [20]. However, the choice of constructing the paths from high dimensional data stream to be best accompanied with the the path signature feature is not yet clear and needs further investigation.

C. Comparison with the Previous Work

In our previous work [22], which is the first of building paths on the cerebral cortex structure, we designed the TPS layer to extract the geometrically and analytically meaningful path signature features for describing the growth patterns of brain regions. Together with an attention mask generator, our network achieves a competitive result on an in-house dataset. It provides a new solution on the longitudinal brain analysis and validates its feasibility. However, it is still challenging for the missing data issue and the trade-off between the excessive feature dimension and limited data. Therefore, in this paper, we extend the previous version by introducing the hierarchical structure, existence embedding and the top K selection module to fully leverage the existing data and decrease model size. More technical details and extended analyses are also added. Our contributions can be concluded as follows,

• We demonstrate the potential of path signature method to have a deeper integration with the machine learning structure and work as a typical sequential descriptor in longitudinal brain analysis.

• We demonstrate the further performance gain enabled by introducing the hierarchical structure, existence embedding and top K selection module to construct the CF-PSNet, compared to the conference version, thanks to its thoroughly exploration and paying more attention on the existing visits and the influential brain regions.

• We perform comprehensive experiments to evaluate the impact of missing data and the number of brain regions used in cognition prediction, and validate the effectiveness our method and its superiority performance compared to the other advanced approaches.

III. Preliminaries of Path Signature

A path X in ${ℝ}^d$ is a continuous mapping of finite length from interval [a, b] to a d-dimensional vector space, i.e. $X:[a,b]\to {ℝ}^d$. We use the subscript notation $X_t=\left\{X_t^1,X_t^2,\cdots ,X_t^d\right\}$ to denote the d-dimensional vector for any $t\in [a,b]$, where $X_t^i$ denotes the i^th coordinate of X_t and $i\in \{1,2,\cdots ,d\}$. Before introducing the signature, let us introduce $S_l{(X)}_{a,b}$, the l^th fold iterated integral of the path X, for the convenience of understanding.

The 1^st iterated integral of Xⁱ equals to the increment of X along the i^th coordinate, i.e., $X_b^i-X_a^i$, which can be expressed by $Sig{(X)}_{a,b}^i=\underset{a<t_1<b}{\int }d{X}_{t_1}^i$. Meanwhile, the 1^st fold iterated integral of X is defined as the collection of $Sig{(X)}_{a,b}^i$ for $i\in \{1,2,\cdots ,d\}$, i.e.,

which is a d-dimensional vector. The 0^th fold iterated integral $S_0{(X)}_{a,b}$ is set to 1 conventionally.

Similarly, the 2^nd fold iterated integral of X is defined as the integral of S₁(X)_a,. against X, which is the collection of all 2^nd iterated integrals of X with possible index (i₁, i₂), i.e., ${\left(Sig{(X)}_{a,b}^{i_1,i_2}\right)}_{i_1,i_2\in \{1,\cdots ,d\}}$. The collection of 2^nd iterated integrals can be written in the tensor form as follows:

The l-th fold iterated integral of X, $S_l{(X)}_{a,b}$, is the integral of $S_{l-1}{(X)}_a$, against X:

whose dimensionality is d^l.

In general, the signature of a path is a graded infinite series containing all the l fold iterated integrals. In practice, for dimensionality reduction, we may truncate the signature up to a finite degree l, which is denoted as $Si{g}_l{(X)}_{a,b}$:

The dimension of $Si{g}_l{(X)}_{a,b}$ defined in Equation (4) is:

In practice, the discrete time series are more common, which can be embedded into the path space by the linear interpolation. The corresponding signature of the embedded path can be computed explicitly by Chen’s identity [29] as follows:

which is a non-linear feature set of the time series data. Specifically, in this application, we interpret X as the growth trajectories of brain regions extracted from the longitudinal MRI scans. It is noteworthy that the 1^st fold iterated integral of growth trajectories are the variations of biological measurements, while the linear combination of the 2^nd iterated integrals, $\frac{1}{2}\left(Sig{(X)}_{a,b}^{i_1,i_2}-Sig{(X)}_{a,b}^{i_2,i_1}\right)$, is equal to the area enclosed by the curve $\left(X^{i_1},X^{i_2}\right)$ and a chord connecting the starting and the ending points of the path, which can be used to evaluate the curvature of the growth trajectories.

The signature of a path is an effective feature set of the streamed data containing many algebraic and analytic properties. It remains invariant under time re-parameterization of path X [33]. It means that the signatures capture the information on the path trajectory while removing infinite-dimensional noise caused by speed variation. Intuitively, the signature of a path plays a role as the non-commutative polynomial on the path space. Further, in practice, the signature feature set can be used to handle time series of variable length, and variation caused by the time parameterization. Therefore, it can be used as a a fixed dimensional and principled descriptor to summarize the path information over intervals effectively and offer dimension benefits especially for long time series. For interested readers, [34] and [35] are suggested to find more details.

A. Top K Selection Module

First, we briefly review the top K selection module, as illustrated in Fig. 3, which is proposed to select the top K influential brain regions for each task based on features over all 9 time points. As shown in Fig. 4, features of the same brain region are flattened into a 36-dimensional feature vector and then fed into a fully connected layer to produce a scalar as its influence coefficient under a certain cognitive task. In terms of these influence coefficients, top K influential ROIs are selected to be fed into following layers. Further, we multiply influence coefficients with the corresponding feature vectors of different ROIs as an initial weight. By selecting the top K influential brain regions, the excessive dimension of input features is substantially decreased. Then, parameters needed in the next following layers are accordingly reduced.

B. Temporal Path Signature Layer

Stacked temporal path signature layers are then proposed to extract multi-scale dynamic information and generate discriminative representation for each participant. In each TPS layer, for the first step, K paths corresponding to K most influential ROIs are defined along the time axis as shown in Fig. 5 and then, like what CNN does, split by overlapping sliding windows with the size of W and a sliding stride of 1 as illustrated in Fig. 6. Consequently, for each path, $T^i=\left(T^{i-1}+1-W\right)$ sub-paths are generated to further explore local temporal properties, each covers a temporal receptive field of a $\widetilde{T^i}=9-i\times (W+1)$ time points period. In this work, we denote i as the index of TPS layers, while T⁰ = 9 for input. Then, for every sub-path, Equations (4) and (6) together with Chen’s identity are employed to compute its corresponding path signature features and obtain an output with the dimension of $n_{\text{PS}}=\left(4^{l+1}-1\right)/(4-1)$ according to Equation (5) with d being set as 4. The outputs of different sliding windows are gathered to construct a new path, shown in Fig. 6 for details. Afterwards, an 1 × 1 convolutional layer is again introduced to conduct a feature transformation as well as dimensionality reduction from n_PS to 4 for SSS problem. Notably, this is the only set of parameters needed in the TPS layer. Given the short time series we have in this study, we only apply two TPS layers in CF-PSNet and generate a global temporal representation which covers all existing visits. We hope that it can be tested with longer sequence in future works.

Additionally, as the existence of missing visits, it would be better to highlight the position of the interpolated visits in paths explicitly. To this end, inspired by positional encoding in [47], we design an existence embedding in path generation to lift the path into a higher dimension. The existence embedding is a 1-dimension sequence which has the same length with the path generated to indicate the more reliable existing visits, so that it can be concatenated with the path. There are many choices of existence embedding. In this work, we use the cumulative sequence:

In the example above, Path denotes for the path we generated, while ∘ denotes for the missing visits. We use Binary to suggest the existence of each time point and compute the existence embedding as the cumulative sum of the Binary sequence.

C. Multi-stream Neural Network with Group Fully Connected Layers

With stacked TPS layers, hierarchical temporal information is extracted. A multi-stream neural network is then included in our CF-PSNet to process the raw features and multi-scale temporal PS features separately, believing that each stream transforms a set of features aggregated to a certain level. Considering that MR scans are collected sparsely along time, the influence of each brain region may vary in different time points. Therefore, we introduce group fully connected layers in all streams in Fig. 3, regard features sharing the same receptive field in time domain as a group and process them separately by applying a group-specific fully connected layer. For instance, the numbers of groups are 9, 5, 1 respectively for Stream I, II and III when the sliding window size W is set to 5. The output of each group-specific fully connected layer is a set of features which encodes the cortical structure at corresponding developmental periods. At the end of the multi-stream neural network, the informative vectors produced by all streams are concatenated and fused with a fully connected layer and output the predicted cognitive score y (depicted in purple in Fig. 3).

D. Attention Mask Generator

For emphasizing the most influential brain region in each developmental stage, an attention mask generator is introduced as a parallel branch of the multi-stream neural network in Fig. 7. In this branch, a set of intermediate cognitive score $y_i,i\in \{1,2,\cdots ,9\}$ is obtained for each time point of the input data by applying the group fully connected layers sequentially with the group number of 9. The importance of each brain region at a certain time point is implicitly learned by the weights of the corresponding group-specific layer. Thus, we sum over parameters of group fully connected layers, which green dash lines surrounded in Fig. 7 along input channels and generate a 9 × K attention mask corresponding to 9 time points and K ROIs. For features in Stream II and III whose receptive fields are more than a single time point, averaging are conducted over the corresponding receptive field of attention mask. Then, element-wise multiplications are conducted between groups of features and attention masks to weight brain regions differently for different time stages and time scales. Notably, in this work, the intermediate output, $\widehat{y}=\left(y_1,y_2,\cdots ,y_9\right)$ is used to assist generating attention masks, but not the final output.

E. Loss Function

All modules mentioned above are learned simultaneously with the loss function defined as:

which consists of two set of L1 loss to constrain the predicted cognitive scores y and the intermediate output y separately. Their ground truth are cognitive scales estimated at 48 months after birth which are denoted as Y. By constraining the intermediate output $\widehat{y}$ with Y, we assuming that the cognitive function development at 48 months can be predicted by a single visit before. λ A is a hyper-parameter introduced to balance the two L1 loss.

A. Model Configurations

As described in Section IV, we conduct experiments on an in-house longitudinal dataset. Following the experimental settings in [13], we perform leave-one-out validation and use the Root Mean Squared Error (RMSE) between the predict values and ground truth of five cognitive scores as the evaluation metric. Learning rate is tuned in {10⁻³,10⁻⁴,10⁻⁵} with Adam as our optimizer. ReLU is employed as the nonlinear activation for hidden neurons. After comparison, we fix λ = 0.1. In all our experiments, the number of epochs is at a maximum of 400.

Notably, we normalize five cognitive scores with their maximum and minimum values separately within a [0,1] range to have a unified comparison setting. The mean ± standard deviation for 5 tasks are 0.6615 ± 0.2467, 0.5842 ± 0.2571, 5565 ± 0.2913, 0.6910 ± 0.2014, 0.6087 ± 0.2628 respectively. An open-source package, signatory², is used to provide efficient path signature computation with backpropagation.

B. Comparison on the TPS Layer

In this work, for the first time, we introduce path signature method to extract rich dynamic information hidden in longitudinal data of structural MRI. In this subsection, we will evaluate its capability by selecting different hyper-parameters and comparing with other sequential models. First, we run our method with different sliding window sizes W and truncated levels l. It is noteworthy that two TPS layers in Fig. 3 are sharing the same choices of W and l in practice. Nevertheless, our method achieves great performance. To illustrate the impact of these two hyper-parameters on model size, some examples are shown in Table. II. It can be seen that, the number of parameters decreases with a smaller truncated level and a larger sliding window size. Then, we compare the performance of CF-PSNet with different hyper-parameter choices in terms of averaging RMSE of all five prediction tasks. In Table III, the best performance of each row improves when we increase the sliding window size, since a larger W may lead to a larger receptive field, more existing visits and lighter fully group connected layers. In most cases, the PS of a higher level typically characterizes more trivial details and does not leads to further improvements. Based on the discussion above, we selected W = 5 and l = 2 in following experiments.

To validate the effectiveness and efficiency of our TPS layer, we select some well-known sequence encoders (i.e., Bi-LSTM [48], GRU [49] and Transformer [47]) as our substitutes. In practice, we replaced each TPS layer in Fig. 3 with single layer of other substitutes to have a fair comparison. In Table I, we found out that our path-signature-based method shows superior performances in various metrics without introducing extra parameters under the SSS problem.

C. Comparison on the Top K Selection Module

On the top of CF-PSNet, a pooling layer is proposed to select the most influential brain regions for each cognitive functions and eliminate redundancy. DiffPool [50] are selected to be compared with the top K selection module as both of them do not require the adjacent information which is lacked in structural MRI data. In Table IV, we find that the top K selection module and DiffPool both effectively decrease the number of parameters. Nonetheless, the top K selection module outperforms DiffPool by a large margin. The reason might be that the former one preserves the cortical structural information by retaining the features of top K influential ROIs, whereas the latter generates K clusters using features from all brain regions, some of which may be useless or containing noises. Further, we test the top K selection module by varying the selection of K. In Fig. 8, the average RMSE for all five cognitive scores is illustrated. It can be observed that with more brain regions selected, the model needs more parameters to be optimized, while the performance is deteriorating. Of note, the top K selection module only contains 36 parameters to optimize, while the model size decreased for more than 90% after applying this module.

D. Model Analysis

In Fig. 9 and Fig. 10, several experiments are conducted to investigate the impact of feeding different number of visits and adjusting the missing rate on the model performance. Specifically, in Fig. 9, we use the CF-PSNet with single TPS layer as the model configuration, as it requires at least 9 time points to apply the stacked TPS layers. Notably, model achieves the best performance when we feed in time points from OM to 18M after birth. One possible explanation is the excessive missing rates in the 24M and 36M as illustrated in Fig. 2. Further, we try to adjust missing rate by randomly drop existing visits. The initial missing rate is around 0.306, while we add 2% missing visits each time. A sharp decrease on the model performance can be observed as the missing rate grows, while the performance with existence embedding is more steady suggesting its effectiveness. Three different runs are conducted for the reliability.

E. Comparison with the State-of-the-art Methods

Several recent algorithms are selected as baselines to validate the performance of our method, including: 1) NN (nearest neighbour) [53]; 2) MtJFS (Multi-Task Learning with Joint Feature Selection) [52]; 3) TrMTL (Trace-Norm Regularized Multi-Task Learning) [51]; 4)RMTL (Robust Multi-Task Feature Learning) [54], 5) MTMLR (Multi-Task Multi-Linear Regression) [12] and 6) LPMvRL (Latent Partial Multi-view Representation Learning) [13]. It is noteworthy that in these methods, five cognitive scores are learned simultaneously with the belief that these scores are essentially inter-related and can benefit each other for the prediction tasks. Hence, we proposed CF–PSNet(multi–task) to predict these five scores at once and have a fair comparison. From Table VI, we find that CF–PSNet(multi–task) outperforms the other algorithms in three tasks and Average under the same settings which validates the performance of our method. Additionally, since MTMLR [12] used different morphological features in its experiments, we adjusted the data configuration accordingly, which is denoted as Multi–taskt in Table VI. The huge margin between CF–PSNet(multi–task) and CF–PS Net may indicate that the important features for these five tasks are not the same and some task-specific brain regions exist. We conducted the paired t-test between the result using model in our conference version [22] and CF-PSNet. The p-values are small (< 0.05) in most tasks, i.e., 0.029 for VRS, 0.001 for FMS, 0.674 for RLS, 0.001 for ELS, 0.001 for ELC proving the improvement is significant. The scatter plots of 10 different runs for predicting the cognitive scores are depicted in Fig. 11 demonstrating that the scores are predicted fairly good.

In general, we conclude our experimental results as follows. 1) Features of multiple temporal resolution may contribute to the predicting accuracy of cognitive scores. TPS layers can effectively and efficiently extract the dynamic information hidden in cortical morphological growth trajectories. 2) Every cognition function may be associated with few brain regions. Most of ROIs are redundant. 3) Further, there are some general and task-specific brain regions for these five tasks.