Multi-Stage Prediction Networks for Data Harmonization
Stefano B. Blumberg Marco Palombo Can Son Khoo Chantal M. W. Tax Ryutaro Tanno Daniel C. Alexander
cs.LG cs.CV stat.ML





Multi-Stage Prediction Networks for Data Harmonization








Stefano B. Blumberg et al.

Department of Computer Science and Centre for Medical Image Computing, 
 University College London (UCL)
Cardiff University Brain Research Imaging Centre (CUBRIC), Cardiff University

 stefano.blumberg.17@ucl.ac.uk
Multi-Stage Prediction Networks for
 Data Harmonization
    Stefano B. Blumberg1 Marco Palombo1 Can Son Khoo1 
 Chantal M. W. Tax2 Ryutaro Tanno1 Daniel C. Alexander1
    [10pt]0pt0pt December 30, 2023
==============================================================================================================




In this paper, we introduce multi-task learning (MTL) to data harmonization (DH); where we aim to harmonize images across different acquisition platforms and sites.  This allows us to integrate information from multiple acquisitions and improve the predictive performance and learning efficiency of the harmonization model. Specifically, we introduce the Multi Stage Prediction (MSP) Network, a MTL framework that incorporates neural networks of potentially disparate architectures, trained for different individual acquisition platforms, into a larger architecture that is refined in unison.  The MSP utilizes high-level features of single networks for individual tasks, as inputs of additional neural networks to inform the final prediction, therefore exploiting redundancy across tasks to make the most of limited training data.  We validate our methods on a dMRI harmonization challenge dataset,  where we predict three modern platform types, from one obtained from an old scanner.  We show how MTL architectures, such as the MSP, produce around 20% improvement of patch-based mean-squared error over current state-of-the-art methods and that our MSP outperforms off-the-shelf MTL networks.  
Our code is available <cit.>.


Data Harmonization, Deep Learning, Diffusion Magnetic Resonance Imaging, Multi-Task Learning, Transfer Learning.







§ INTRODUCTION


Lack of standardization amongst imaging acquisitions is a long-standing problem, that confounds large scale multi-centre imaging studies.  DH aims to remove differences arising from specifics of scanner, centre or acquisition protocol and increase the power of group studies.  The problem aligns with the wider issue of estimating, from an image obtained from a particular subject and scanner, the image that would have been obtained from the same subject on another scanner.  Various recent studies, such as Image Quality Transfer <cit.> and Modality Transfer <cit.>, present solutions to this problem that may be repurposed for DH.

 Current approaches used in practice, such as <cit.>, generally register all images to a common template and align the mean and variances of image intensities from each platform (i.e. centre, scanner, and/or acquisition protocol), either voxelwise, regionally, or of whole images.     Recently, deep learning has shown much promise, e.g. in the Multi-Shell Data Harmonization Challenge (MUSHAC) <cit.>, outperforming statistical approaches. Methods to date include the Convolutional Neural Network with Rotationally Invariant Spherical Harmonics (CNN-RISH) <cit.> – a 5-layer CNN that harmonizes dMRI from 3T to 7T; the Deeper Image Quality Transfer Network (DIQT) <cit.> – an extension of the FCSNet <cit.>, which holds state-of-the-art results in dMRI super-resolution and the Spherical Harmonic Residual Network (SHResNet) <cit.> – A 7-layer CNN that processes different spherical harmonic (SH) coefficients separately.  A common feature of all these methods is that they use a separate single CNN to estimate the target images from those of each individual platforms. 


 Multi-scanner or multi-centre studies often acquire traveling heads (TH) – a small number of subjects specifically scanned at each site, to support DH learning mappings from scanner to scanner. These learned mappings then transform other subjects to a common representation.  However, acquiring TH data is expensive and typically the number of subjects is small (rarely more than 10).  Therefore, single-network deep learning approaches easily overfit <cit.> and fail to exploit the synergy between multiple, often strongly related, prediction tasks.    

 MTL potentially offers a powerful solution to the paucity of training data in DH, by allowing the model to integrate information from acquisitions from multiple platforms. In contrast to single-network approaches, it exploits commonalities and differences across multiple tasks (in our case, predicting different acquisitions) to improve the predictive power of neural networks.  Furthermore by incorporating predictions and loss functions into a single network, MTL approaches acquire additional regularisation <cit.>.  There are a variety of existing approaches that cascade over multiple predictions of varying resolutions, that might be adapted to a MTL approach in DH.  Examples include the Convolutional Pose Machine Networks (CPM) <cit.> comprising a sequential architecture of CNNs, to provide increasing refined estimates and the Holistically-Nested Edge Detection Network (HNED) <cit.> – a single-stream deep network with multiple side outputs to perform multi-scale and multi-level feature learning.  However, these approaches process tasks sequentially, whilst the distinct tasks in DH differ in terms of resolution, contrast and noise patterns, with no simple notion of ordering.  Furthermore it is unclear as to what might be the optimal subnetworks for such a larger network.


 In this paper, we introduce a new MSP architecture designed to exploit, for the first time, MTL in DH.  The MSP draws on sub-network structures from <cit.>, but cascades across various estimation tasks, rather than assuming they are sequentially structured.  This avoids both duplicating features and allows us to use redundancy across tasks to make the most of limited training data.
Furthermore the MSP allows us to incorporate pre-trained state-of-the art networks into a MTL framework, letting us explicitly use high-level features of multiple predictions, as network inputs to inform the final prediction.

 We utilize the recent MUSHAC <cit.> data to evaluate various strategies.  The specific task is to predict acquisitions of both low and high quality from different scanners, from data from the same subjects acquired on an ageing scanner.  We compare performance of the state of the art, i.e. various single network approaches, with MTL frameworks, including the MSP, where the best performing single networks inform construction of a particular MSP architecture.   The MSP outperforms both single networks and simple off-the-shelf MTL approaches.  We expect this initial demonstration to motivate wider development and usage of MTL strategies for DH, therefore we release our implementation in <cit.>.




§ METHODS




In this section we outline the general DH problem with THs and setting from MUSHAC <cit.> with which we validate our methods.  Then we construct the MSP.



 In the typical DH setting, THs provide image data sets I_ij i=0,...,P-1  j=1,...,N for each of N subjects imaged on each of P platforms (a scanner and imaging protocol combination).  The challenge is to construct mappings M_iT i=0,...,P-1 from images acquired on platform i to the corresponding image acquired on some target platform T ∈{0,...,P-1}. 


 MUSHAC <cit.> tests the ability to predict multiple high quality targets (modern platforms) from a single low-quality input (an old out-of-date platform) thus concentrating on the most challenging aspect of DH.  With the same notation as before, we predict images I_ij i=1,...,P-1  j=1,...,N from I_0j j=1,..,N. 

 When predicting acquisitions from platform T ∈{1,...,P-1}, current state-of-the-art deep learning approaches extract patch pairs { (x(k),y_T(k)),... } from respective I_0j,I_Tj j=1,...,N.  These approaches train a neural network 𝒩_T on these patches, where for input patch x we denote the neural network's estimate (also denoted as prediction) of y_T as y_T i.e. y_T = 𝒩_T(x) and the loss is calculated as L(y_T,y_T).  This process is repeated for different T.  In this paper we use CNN-RISH <cit.>, DIQT <cit.>, SHResNet <cit.> to represent this class of technique.

 In figure-<ref> we contrast the MSP architecture that we propose, with the single network approach.  The MSP integrates information from multiple
platforms and utilizes high-level features from individual networks, to inform the final prediction with additional neural networks.  In training we use the set of corresponding patches
{ (x(k),y_1(k),...,y_T(k),...),... } from all images.  We first take the pre-trained networks 𝒩_i i=1,...,P-1 as the state-of-the-art single-networks, which may have different architectures, that have already been optimized.  We denote for input patch x, the predictions of these networks as the first-stage prediction of the MSP: y^1_i = 𝒩_i(x)  i = 1,...T.  For networks 𝒩_i i ≠ T, we denote its last feature map as z_i.  We then create the new networks 𝒩_iT i ≠ T, each taking input z_i, where 𝒩_iT informs the final prediction of platform T with information from platform i ≠ T.  The second-stage prediction of the MSP, is a linear combination of network outputs of 𝒩_iT(z_i)  i ≠ T, with the fist-stage prediction y^1_T:

    y^2_T =  (1-α)y^1_T + α/P-1 (∑_i ≠ T𝒩_iT(z_i) + y^1_T)    α∈ [0,1].

We train the whole MSP in unison, using supervised loss functions calculated at both stage predictions: L(y^1_i,y_i),L(y^2_T,y_T)  i=1,...,P-1, which are backpropagated through the network.  For α∈ [0,1], we manually increase α from 0 to 1 during training, to combine pre-trained 𝒩_T with the untrained 𝒩_iT as in <cit.>.

 We provide an example of the MSP in figure-<ref> for P=4, T=1 and code <cit.>.




§ EXPERIMENTS AND RESULTS


In this section we present an overview of our MUSHAC dMRI hamonization dataset <cit.>.  We describe how to process the raw data to produce patches for our supervised neural network approach.  We then illustrate how MTL approaches, such as the MSP, produce improved results over state-of-the-art methods.



 Harmonization Data  Data are obtained from the MUSHAC <cit.>. 10 healthy volunteers were scanned on three different scanners, an ageing 3T General Electric Excite-HDx scanner (max. gradient 40 mT/m), a modern but standard 3T Siemens Prisma scanner (max. gradient 80 mT/m) and a state-of-the-art bespoke 3T Siemens Connectom scanner (max. gradient 300 mT/m). dMRI images were acquired for b = 1200 s/mm2 with two different manufacturer quality protocols: a standard (st) protocol – with voxel side 2.4 mm, 30 directions per b-value, TE = 89 ms from all scanners; a state-of-the-art (sa) protocol on the Connectom – with voxel size 1.2 mm, 60 directions per b-value, TE = 68 ms.  Note we excluded the sa protocol of the Prisma scanner, due to severe mis-alignments.  For sa protocol, multiband-acquisition and stronger gradients shortened TE and improved both spatial and angular resolution per unit time. Additional b=0 s/mm2 images were acquired with TE and/or TR matching between protocols, as well as structural MPRAGEs for each scanner.  The data were corrected for EPI and eddy-current distortions, subject motion and gradient non-linearity, as well as co-registered together as in <cit.>. In addition, brain masks excluding the skull and background were provided for each subject and each acquisition.  The dataset size is representative of typical TH data-sets, where it is difficult to transport a large number of subjects from different centres and obtain multiple scans.  See figure-<ref> for a visualisation.




 Processing and Training Raw data were pre-processed to extract signal features by SH deconvolution – variants of this technique are used as input-output common to the challenge and our comparison networks <cit.>.  We employed  <cit.> where the 28 coefficients of the 6th order real-and-symmetric SH deconvolution were estimated from the normalized raw signal considered as separate channels of our data.  We first normalized the data, on each acquisition, per channel, to be mean 0 and standard deviation 1. 

 Neural-network data-harmonisation transformations work patch-by-patch on an input image. Each voxel within the brain mask of the input acquisition defines a patch with that voxel at the centre. The corresponding target patch comes from the corresponding location in the target image.  For this paper we utilised input patch size as 11^3, with target patch size 11^3,19^3, depending on the target resolution.  We then separated the patches into 90%, 10% training, test set.

 Our implementation used PyTorch with Python, with a minibatch size of 12, ADAM optimizer, learning rate starting at 1E-4 and decaying by √(2) per 15 or 25 epochs, for a minimum of 50 total epochs <cit.>.

 

 Comparison We evaluate our approach via the tasks in MUSHAC <cit.>: to predict images from more modern platforms prisma_st,  connectom_st, connectom_sa, from an older platform ge_st (see figure-<ref>).  Our evaluation was performed on a local scale, as in <cit.>, we calculate the mean-squared-error (MSE) on the SH space.  This metric is not based on the raw dMRI and encompasses all key aspects of the diffusion signal (anistropy, mean diffusivity, principle orientation).

 We select three baselines, representing state-of-the-art approaches <cit.>, that predict each acquisition using separate neural networks: CNN-RISH <cit.>, DIQT <cit.>, SHResNet <cit.>.  We also consider simple MTL sub-networks from CPM <cit.>, HNED <cit.>, illustrated in the supplementary material, which were formulated for different tasks and inspired the MSP.  The super-resolution networks replaced a standard 3D convolution with a deconvolution / strided convolution layer, the connection networks are a stack of two 3D convolutions.  We then constructed the MSP by combining the best single-network approaches (bold in table-<ref>).


Table-<ref> shows that MTL in general improves with respect to the single-network approaches, for all three target platforms. We use the Wilcoxon signed-rank test, a non-parametric statistical test, confirming that for all predicted platforms 
the differences in error scores are statistically significant, as p values obtained are all less than 10^-5, after Bonferroni correction.

 Figure-<ref> shows qualitative results revealing differences in results from DIQT and MSP for connectom_sa, on a subject excluded from the training set.  Improvements from MSP over  DIQT are subtle but visible on both the SH and direction-encoded colour maps, e.g. MSP reduces errors directly underneath the left ventricle.  Quantitatively in that figure, MSP improves the global SH MSE by 5% and the global direction-encoded colour MSE by 11%.





§ CONCLUSION AND FUTURE WORK


In this paper we demonstrated how MTL approaches such as the MSP, increase the predictive power of deep learning in DH and improve over state-of-the-art single network approaches.

 There are many potential directions of future work.  Here we test just one MSP architecture, but many variations of deeper structures with more sophisticated connections between tasks are possible.  For example, the MSP chooses z_i to be the
the deepest feature of 𝒩_iT, in future work we could experiment with connecting different layers.


 One key problem in evaluating DH performance is residual misalignment of input and output images. This can mask strong performance when using simple global comparison metrics.  It may be useful to devise alternative evaluation metrics that are robust to misalignment, or by moving away from direct image evaluation and analyzing downstream processes.


 In this paper, we demonstrated the MSP via a diffusion MRI data set.  However the same architecture extends naturally to harmonizing data between multiple different modalities, within or beyond MRI.



§ ACKNOWLEDGEMENTS

We thank: Tristan Clark and the HPC team (James O'Connor, Edward Martin) and the organizers of the 2017 MUSHAC challenge (Francesco Grussu, Enrico Kaden,  Lipeng Ning, Jelle Veraart).  This work was supported by an EPRSC and Microsoft scholarship and EPSRC grants M020533 R006032 R014019 and the NIHR UCLH Biomedical Research Centre.  The data were acquired at the UK National Facility for In Vivo MR Imaging of Human Tissue Microstructure located in CUBRIC funded by the EPSRC (grant EP/M029778/1), and The Wolfson Foundation. Acquisition and processing of the data was supported by a Rubicon grant from the NWO (680-50-1527), a Wellcome Trust Investigator Award (096646/Z/11/Z), and a Wellcome Trust Strategic Award (104943/Z/14/Z).

splncs




19




CODE
Code: <https://github.com/sbb-gh/>   <http://mig.cs.ucl.ac.uk/>

IQTVRFR
Alexander, D.C., et al.: Image quality transfer via random forest regression: Applications in diffusion MRI. In: MICCAI. (2014)

BIQT
Tanno, R., et al.: Bayesian image quality transfer with CNNs: exploring uncertainty in dMRI super-resolution. In: MICCAI. (2017)

DIQT
Blumberg, S.B., et al.: Deeper image quality transfer: Training low-memory neural networks for 3D images. In: MICCAI. (2018)

MPCSSSS
Ye, D.H., et al.: Modality propagation: Coherent synthesis of subject-specific scans with data-driven regularization. In: MICCAI. (2013)

HCTM
Fortin, J.P., et al.: Harmonization of cortical thickness measurements across scanners and sites. NeuroImage 167 (2018) 104-120

MSHDMRI
Mirzaalian, H., et al.: Multi-site harmonization of diffusion MRI data in a registration framework. Brain Imaging and Behavior 12 (02 2017)

CSCPDMRIDH
Tax, C.M., et al.: Cross-scanner and cross-protocol diffusion MRI data harmonisation: A benchmark database and evaluation of algorithms. NeuroImage 195 (2019) 285-299

CSCPHMSDMRIOCEV
Ning, L., et al.: Cross-scanner and cross-protocol harmonisation of multi-shell diffusion MRI data: open challenge and evaluation results. In: ISMRM. (2018)

MSDMRIEC
Ning, L., et al.: Muti-shell diffusion MRI harmonisation and enhancement challenge (MUSHAC): Progress and results. In: MICCAI CDMRI Workshop. (2018)

HDMRID
Cetin Karayumak, S., et al.: Harmonizing diffusion MRI data across magnetic field strengths.  In: MICCAI. (2018)

SHRN
Koppers, S., et al.: Spherical harmonic residual network for diffusion signal harmonization.  In: MICCAI CDMRI Workshop. (2018)




DSN
Lee, C.Y., et al.: Deeply-supervised nets. In: AISTATS. (2015)

CPM
Wei, S.E., et al.: Convolutional pose machines. In: CVPR. (2016)

HNED
Xie, S., Tu, Z.: Holistically-nested edge detection. In: ICCV. (2015)

PGG
Karras, T., et al.: Progressive growing of GANs for improved quality, stability, and variation. In: ICLR. (2018)

DEFODF
Tournier, J.D., et al.: Direct estimation of the fiber orientation density function from diffusion-weighted MRI data using spherical deconvolution.  NeuroImage 23(3) (2004) 1176-1185

DIPY
Garyfallidis, E., et al.: DIPY, a library for the analysis of diffusion MRI data. Frontiers in neuroinformatics 8 (02 2014)

FSL
Jenkinson, M., et al.: Multi-site harmonization of diffusion MRI data in a registration framework.  NeuroImage 62 (2012) 782-790





§ SUPPLEMENTARY MATERIAL





