Non-Gaussian normative modelling with hierarchical Bayesian regression

Brain disorders affect millions of people worldwide, and have a large impact on the quality of life of the patient and the people surrounding them. Correct treatment in the form of medication or therapy can help reduce symptoms. However, studies to understand mechanisms or develop treatments for these disorders are usually performed in a case-control framework, which provides group-level inferences in which homogeneity within clinical cohorts and clinical groups is presumed. These studies aim to detect differences between the clinical and the control (i.e., healthy) groups and are therefore limited to providing inferences about group averages (e.g., the “average patient”). However, it has been argued that the individual deviations from the groups’ means contain the most important information for stratification (Marquand, Rezek, et al., 2016; Marquand, Wolfers, et al., 2016). An improved understanding of the heterogeneity within clinical cohorts may ultimately lead to better diagnosis and treatment.

Normative modelling provides a framework for understanding heterogeneity within clinical cohorts (Marquand, Rezek, et al., 2016). A normative model aims to provide estimates of centiles of variation of a specific set of phenotypes over the full range of the clinical covariates, that is, across the lifespan, much like growth charts in pediatric medicine. Centiles of variation of imaging derived phenotypes (IDPs) are estimated, and subjects are assigned a point somewhere along the estimated centiles. The assumption is that healthy subjects typically land in areas of higher density than non-healthy subjects. The mean or “average” subject is thus also a meaningful concept in normative models, but contrary to case-control studies, the deviations from that mean are interpreted as anomaly in the brain structure or function, and are seen as valuable resources for further downstream processing like analysis of underlying mechanism of brain disorders.

In probabilistic modelling, a distinction is made between epistemic and aleatoric uncertainty (Kendall & Gal, 2017). The former is due to uncertainty about a parameter or phenomenon, and generally becomes smaller as more data are collected; the latter is uncertainty due to the variation that is naturally present in the measured phenomenon. In normative modelling, we aim to capture and separate these forms of uncertainty to the highest possible degree. Separating all sources of aleatoric and epistemic uncertainty allows adequately controlling for nuisance variation, and guiding further decisions by relevant (biological) variation only.

As in many other fields, data availability in neuroscience has increased dramatically in recent years, and that trend will likely continue. Medical data, however, is sometimes subject to restrictions involving sharing, due to certain permissions not being granted. There is no guarantee that all data on which a model is to be trained can be centrally located. The question of how to do normative modelling on large, distributed datasets is thus a natural one, and it comes with a number of challenges. Firstly, the method must be able to handle large amounts of data. We have seen that Gaussian processes would be viable candidates for normative modelling, if it were not for their time complexity that scales cubically with the number of input data points (Marquand, Rezek, et al., 2016). Second, distributed datasets inevitably come with sources of nuisance variance; since every scan site has its own protocols, every scanner has its own unique noise characteristics, and there may be unknown sources of noise present in specific sites as well. In addition, it is well known that there are sex-related differences in neurobiological phenotypes. Hirearchical Bayesian Regression (HBR) can deal with these so-called batch effects in data by assuming a shared prior over batch-specific model parameters, as we describe in Appendix B. Third, neuroimaging data are complex and provide rich information about the individual from which they were derived. These data therefore need to be treated with care. Permission to share personal data is not always available, and transportation of the data is therefore not always possible. This requires federated learning; a distributed learning approach that does not require transportation of the data, only of the inferred parameters of model (Kia et al., 2021; Li et al., 2020). In this study, we employ and extend the hierarchical Bayesian regression (HBR) approach for normative modelling (Gelman & Rubin, 1992; Kia et al., 2020, 2021) on non-Gaussian IDPs. The method involves defining a generative model, and inferring the posterior distribution over the model parameters using Markov chain Monte Carlo (MCMC) sampling. By absorbing and summarizing information contained in training data in learned posteriors, HBR supports federated learning; all information required to extend or adapt the model is contained in the model, eliminating the need of data transportation when using or even extending the model.

The current implementation of HBR assumes that the residuals follow a Gaussian distribution. However, the variability in many imaging phenotypes cannot be properly described by a Gaussian distribution, as they are skewed, kurtotic, or both. The issue of non-Gaussianity in neurological phenotypes has partially been addressed in Fraza et al. (2021). This method accounts for non-Gaussianity by applying a monotonic warping function to the data (Snelson et al., 2003), which yields a closed form for the marginal likelihood that can be maximised using standard optimisation techniques (Dinga et al., 2021). Here, we utilise the same warping function—the sinh-arcsinh transformation, Appendix A—but in a fully Bayesian approach, that is, in the HBR framework, yielding full posterior distributions for all parameters. Applying the sinh-arcsinh transformation to a standard Gaussian distribution yields the SHASH distribution (Jones & Pewsey, 2009), which is detailed in Section 2.1, a flexible distribution with separate parameters roughly corresponding to skew and kurtosis. Here, we apply the HBR framework to the a canonical variant of SHASH distribution and a novel reparameterisation that substantially reduces dependency between the parameters, which is crucial for more efficient sampling. We assess its performance on some known problematic imaging phenotypes, and show that the HBR framework with a SHASH likelihood is able to model phenotypes in a way that was previously impossible.

The contributions of this work are: (i) a reparameterisation of SHASH distribution and proposal of an MCMC sampling approach to infern this model; (ii) a thorough analysis of the HBR method with a SHASH likelihood in comparison with a warped-Bayesian linear regression as a baseline method (Fraza et al., 2021); (iii) an extensive evaluation of these methods on a large multi-site neuroimaging dataset containing estimates of cortical thickness and subcortical volume; and (iv) an extension to the existing HBR implementation in the pcntoolkit.

The paper is structured as follows: we first provide a theoretical background of HBR, MCMC sampling, and the SHASH distribution in Section 2, where the problems of using the SHASH distribution in a sampling context will also become clear. We introduce a reparameterisation that addresses these problems in Section 2.1.1. A framework for concisely describing common variations of HBR models is introduced in Appendix B. Then, we evaluate our approach using several experiments with a large neuroimaging dataset in Section 3, where we show the proposed method performs on par or better than existing methods. Then, we apply the new HBR models to an even more non-linear dataset (Zabihi et al., 2022), where we show a clear advantage over existing methods. A discussion of the advantages and limitations of our approach follows in Section 4, and we summarise all our findings in Section 5.

At its heart, normative modelling involves fitting a probabilistic regression model to predict centiles of variation in a biological or psychometric response variable as a function of a set of clinically relevant covariates. The covariates are often taken to be demographic variables such as age and sex, but this is not necessarily the case and many other mappings are possible (e.g., using cognitive scores to predict brain activity). Early reports used classical Gaussian process regression (Marquand, Rezek, et al., 2016; Ziegler et al., 2014) which is appealing due to its Bayesian non-parametric nature, exact inference and flexibility in modelling non-linearity with a relatively small number of parameters. However, the exact Gaussian process approach used in these works has two main limitations, namely its extremely poor computational scaling to large numbers of data points (owing largely to the cubic computational scaling required to solve the GPR predictive equations) and for the exact formulation (i.e., having an analytical solution), it is restricted to modelling Gaussian centiles. While generalisations are possible within the Gaussian process literature to address these problems, these typically have other shortcomings. For example, modelling non-Gaussian noise distributions usually requires approximations which further increases computational complexity. The result is that other approaches with better computational scaling have been proposed in the neuroimaging literature. These are reviewed in Marquand et al. (2019) and include quantile regression and hierarchical linear regression techniques. More recently, the field has increasingly recognised the importance of modelling non-Gaussianity in the response variables, while using a distributional form to increase the precision with which outer centiles can be estimated (Borghi et al., 2006). There are two main approaches that have been proposed to solve this problem in the neuroimaging literature, namely likelihood warping (Fraza et al., 2021) and generalised additive models for location, scale, and shape (GAMLSS) (Dinga et al., 2021). While these approaches are both capable of flexibly modelling non-Gaussianity and site effects, both have potential shortcomings: GAMLSS is not probabilistic due to the heavy regularisation penalties that are applied to constrain the flexible smoothers underlying the approach. Likelihood warping is probabilistic in nature, but still does not model the uncertainty associated with shape parameters and does not offer completely flexible control over the nature of the modelled distribution. Neither approach supports fully decentralised federated learning at training time although ad-hoc solutions have been proposed to allow model parameters to be transferred to a new site at test time. The approach proposed in this paper attempts to address these shortcomings.

The HBR framework is a Bayesian modelling framework that assumes a generative model over the response variables $Y$ in the form of a likelihood $ℒ$ and a prior distribution $p$ for each parameter $θi$ of the likelihood. A general model can be expressed as such:

where $P$ is the number of parameters of the likelihood. One of the main objectives in Bayesian modeling is to retrieve the posterior, which is the distribution over the model parameters given the data; $p(θ|Y)$⁠. Bayes’ rule gives us an expression of the posterior: $p(θ|Y)=p(Y|θ)$$p(θ)/p(Y)$⁠, but this expression can only be evaluated analytically in special cases where the likelihood is conjugate to the prior, and this is not the case in HBR. As of such, the posterior is approximated by MCMC sampling, which we discuss in Appendix D. Assuming the response variable follows a Gaussian distribution, we substitute the Gaussian $N$ for $ℒ$⁠, and our parameters $θ$ become the mean and variance $μ$ and $σ2$⁠. The general model described in Equation 1 is fully capable to support non-Gaussian likelihoods, like the family of SHASH distributions, which is detailed in Section 2.1. The way we model $pθi$ depends on our further assumptions about the data. Typically $μ$⁠, and optionally also $σ$ are taken to be linear functions of a set of covariates, which in this example are clinical or demographic variables. This is further detailed in Appendix B, where we also introduce a framework for model specification.

To accurately fit normative models on non-Gaussian-distributed response variables, one could substitute a family of flexible distributions for $ℒ$ in Equation 1. This family must then contain positively as well as negatively skewed members, leptokurtic as well as platykurtic members, and preferably with smooth transitions between them. All those requirements are fulfilled by the family of SHASH distributions (⁠$S$⁠). Figure 1a and 1b show the flexibility of the $S$ in modelling various distributional forms.

To generate samples from a SHASH distribution, one applies an inverse sinh-arcsinh transformation to samples from a standard Gaussian (Jones & Pewsey, 2009). The sinh-arcsinh transformation $ξ$ and its inverse $ξ−1$ are defined as:

where $∈∈ℝ$ and $δ∈ℝ+$ govern the shape of the resulting distribution. Now assuming a random variable $Z$ follows a standard Gaussian, we can construct samples $X$ that follow a SHASH distribution by applying $ξ∈,δ−1$ to $Z$¹:

For the remainder of Section 2, $Z$ and $z$ will indicate Gaussian distributed samples, and $X$ and $x$ will indicate SHASH distributed samples.

To derive the density of this SHASH distribution, we need to apply a change of variables, which states that the density of transformed samples will be the density of the original samples multiplied with the Jacobian determinant of the inverse transformation (Bishop, 2006). The constraint being that this Jacobian determinant exists and that it is not 0. This is true for any transformation that is a diffeomorphism². Because the transformation $ξ$ is a diffeomorphism, we can get the density of $S$ by a change of variables. We multiply the density of the samples in the Gaussian domain, $φ$⁠,—which we find by applying $ξ$—with the derivative of $ξ∈,δ$⁠.

where $C∈,δ(x)=cosh(δsinh−1(x)−∈)$⁠. Figures 1 and 2 illustrate the effect of the parameters $∈$ and $δ$ on the shape of the SHASH density. We clearly see that $∈$ and $δ$ modulate the skew and kurtosis, respectively.

A constraint on the parameter $δ$ can be derived from this result. A well-known property of probability distributions is that they are positive everywhere. The only term in Equation 3 that can possibly be negative is the left fraction, because the middle fraction is a positive constant and the right exponential is strictly positive for all real inputs³. The left fraction is negative if exactly one of $δ$ or $C∈,δ$ is negative, and we have $cosh(α)>0$ for all $α∈ℝ$⁠. Thus, the only way in which we can get a negative density is to have $δ<0$⁠, from which we derive the constraint that $δ>0$⁠.

The SHASH distribution we have seen so far has two parameters, $∈$ and $δ$⁠, but additional parameters of location and scale are necessary to achieve the desired flexibility. Jones and Pewsey (2009) suggest adding scale and location parameters by muliplication with $σ$ and addition of $μ$⁠. The resulting distribution is known as the SHASH_o or $So$ distribution.

Assuming a random variable $Z$ follows a standard Gaussian, we can construct samples $Ω$ that follow a $So$ distribution by applying $ξ∈,δ−1$ to $Z$⁠, multiplying with $σ$⁠, and adding $μ$⁠:

Equivalently, because of 2, we can write

And because we already know the density of $X$ to be $S$⁠, we can use the chain rule in conjunction with the change of variables to arrive at the density of $So$⁠:

where $∈$ and $δ$ are parameters that govern the higher order moments of the distribution (explained below). While we believe this family is flexible enough for the purposes of modeling a wide variety of neuroimaging phenotypes, inferring the parameters of this distribution is highly challenging because of their strong correlations⁴. Figures 1a and 2a illustrate the correlation between $μ$ and $∈$⁠. Clearly, $∈$ controls both the skew and the mean, but $μ$ also controls the mean. Similarly, in Figures 1b and 2b, we see that $δ$ controls the kurtosis as well as the variance, while $σ$ also controls the variance. Writing down the first two moments of the $So$ distribution will make this correlation explicit. Because the samples $Ω$ from SHASH_o are simply scaled and shifted samples $X$ from SHASH (see Equation 5), we can express the moments of $Ω$ simply in terms of the moments of $X$ as:

To break the correlation between $μ$ and $∈$ and between $σ$ and $δ$⁠, we propose to use the expressions for $E[X]$ and $E[X2]$ that are provided by Jones et al. to derive a new density. Jones et al. give us an analytical expression for the $r$’th non-central moment, given here in Equation 8. For derivations of these quantities, please refer to Jones and Pewsey (2009).

for $r∈ℕ$⁠, where

and $K$ is the modified Bessel function of the second kind (Bowman, 1939).

The solution proposed here is to first apply a standardising shift and scale to the SHASH samples, such that the mean and variance of those samples are 0 and 1, respectively. Then, we apply the location and scale parameters $μ$ and $σ$ like before to attain samples $B$⁠. Because the moments of $S$ are all obtainable by Equation 8, this is a simple operation. We define the newly designed transformation as $λ$ and the corresponding density as SHASH_b or $Sb$⁠:

where $η∈,δ2$ is the central variance given by $η∈,δ2=m∈,δ(2)−$$(m∈,δ(1))2$⁠. We can again use the chain rule in conjunction with the change of variables to arrive at the density of $Sb$⁠:

Writing down the first and second moments of $B$⁠, the expectations that appeared in Equations 6 and 7 now cancel with the moments that we added in Equation 10, and the correlation breaks completely:

The new parameterisation has nicely interpretable parameters for location and scale, which the $So$ does not. Compare Figure 1a and 1b with 1c and 1d. The densities in the two rightmost figures all have zero mean and unit variance. This may not seem obvious for the density with $δ=0.4$ in Figure 1d, but the tail behaviour responsible for this becomes clearer from the inset. In Appendix G, H, the empirical moments of the $So$ and $Sb$ distributions are plotted over a range of parameters. It is important to note that the $Sb$ family is isomorphic to $So$⁠, which means that they model exactly the same set of distributions, just under a different parameterisation.

The $Sb$ reparameterisation greatly reduces the problem of correlation in parameter space, but it comes at a cost. Some of the reparameterisation terms get extremely large quickly, and may therefore lead to numerical instability. This may be addressed by appropriately setting the priors for the $Sb$ likelihood. Comparing Figure 2b and 2d, the warp in the low domain of $δ$ is apparent. The stretching effect of the change of variables is enormous in those regions, because it has to compensate for the flattening effect of $δ$ visible in Figure 1b in order to standardise the SHASH distribution. The curve in the line associated with $δ=0.1$ is barely visible in this range. The variance of $X∈,δ$ as a function of $δ$ is shown on a double log scale in Figure 3. Notice that small differences in the lower domain of $δ$ amount to substantial differences in the variance. The enormous magnitude of $ddδm∈,δ(2)$ in those regions is mainly due to the behaviour of the Bessel function $Kn(z)$ in Equation 9, and to a lesser degree due to the fraction in the exponent in Equation 8. In practical applications, values of $δ<0.3$ are thus associated with numerical instability, and sampling in those regions should be avoided at the cost of losing some expressivity. In our models, we enforce this constraint by first enforcing positivity of the samples by applying the softplus function, and then adding a small constant of 0.3. The softplus function enforces positivity and is defined as:

To allow a meaningful comparison with the current methods, we adapt and use the data from (Rutherford et al., 2022). This is a large neuroimaging dataset collected from 82 different scanners, containing 58834 control subjects and 2925 patient subjects⁵. Measures of cortical thickness and subcortical volume were extracted with Freesurfer version 6.0 using image derived phenotypes derived either from the Destrieux cortical parcellation or from the subcortical segmentation. Furthermore, the different UK-Biobank datasets (Sudlow et al., 2015) were merged, because based on prior work (Fraza et al., 2021; Kia et al., 2022), we have observed that site effects in the UKB cohort are minimal for these freesurfer measure derived from the different scanners used (due to a careful harmonisation of acquisition parameters and scanning hardware across this cohort). We therefore consider that we can safely consider them to be drawn from a single site. Counts per site and sex can be found in Appendix F. For full details about the data, we refer to Rutherford et al. (2022). Most data in this dataset are publicly available, but for several of the clinical samples, we do not have explicit permission for data sharing. Such permissions for sharing were not obtained from the data from the University of Michigan, and as such these data were not included in this study (n = 1394 (1149 controls, 245 patients)).

The family of SHASH distributions has the flexibility to model skewed and kurtotic distributions, but also has the standard Gaussian as a central member. We pick some “hard” (i.e., difficult) phenotypes, the age-conditioned distributions of which are not described well by a Gaussian, and some “easy” phenotypes, the age-conditioned distributions of which are described well by a Gaussian. We aim to model both easy and hard phenotypes with the same model, using the same hyperparameters. The selected hard phenotypes are: right cerebellum white matter, because it is slightly skewed; right lateral ventricle, because it is moderately skewed and heteroskedastic; and white matter hypointensities, because it is severely skewed and heteroskedastic. For easy phenotypes, we selected the estimated total intracranial volume, the thickness of the Jensen sulcus, and the volume of the brain stem. These “easy” phenotypes do not show a lot of skew or kurtosis, but the mean of the thickness of the Jensen sulcus shows an interesting trajectory (see Fig. I.26). The distributions of the different phenotypes are shown in Figure 6 and in Appendix I below. In some instances, we used the following abbreviations for the “hard” phenotypes: Right-Cerebellum-White-Matter: RCWM, Right-Lateral-Ventricle: RLV, WM-hypointensities: WMH, and similarly for the “easy” phenotypes: EstimatedTotalIntraCranialVol: ETICV, rh_S_interm_prim-Jensen_thickness: RIPJ, Brain-Stem: BS,

A stratified 10-fold split was made using the Sklearn package (Pedregosa et al., 2011), to which end we filtered out data from sites with less than 10 subjects of one or either sex. This ensures that the test folds contain at least one unique sample for each batch effect. As preprocessing, all data were feature-wise-normalised by subtracting the mean and dividing by the standard deviation. Normalisation is not a necessity, but it allowed the same hyperparameters to be used for different response variables. This would otherwise not be possible, because the scale of the response variables differ by several orders of magnitude. Sampling performance could be hampered if the prior distribution differs exceedingly much from the likelihood.

Age and sex were used as covariates, and site was used as a batch effect. The HBR framework supports multidimensional batch effects, allowing to use both sex and site as a batch effect and allowing site effects to be modelled in a single-step regression together with effects of interest. However, we consider that there is little reason to assume that the differences related to sex vary much between sites, and having this multidimensional batch effect results in a larger number of model parameters. In practice, we saw no difference in performance between models that used sex as a batch effect, and those that used sex as a covariate. For further discussion, see Section 4.

As a second, more challenging, test of the flexibility of our model, we fit a normative model on a highly non-linear representation of brain data derived from the latent representation learned by a convolutional autoencoder. Full details are provided in Zabihi et al. (2022), but briefly; data derived from the UK Biobank sample (n = 20,781, 47% female) were compressed to a 2-dimensional latent representation using a pipeline consisting of a conditional convolutional autoencoder with a 100-dimensional latent space (Kingma & Welling, 2013), followed by a UMAP dimensionality reduction (McInnes et al., 2018). The two-step pipeline was constructed because the intermediate representation produced by the autoencoder would otherwise not be reflective of age and sex, it would only reconstruct the data. In this case, it was desirable to have a latent representation conditioned on these variables. The resulting representation followed a distribution with a nonlinear trajectory with respect to age n the mean, variance, and skew. A stratified train/test (75%/25%) split was made on the raw data level, specifically not at the 100-dimensional latent representation (see Fig. 9). This means that the same train data were used for training the convolutional autoencoder as well as the UMAP embedding, and all of the test data were held out from training. Following the rationale described above, no batch effects were modelled. Again, we used age and sex as covariates.

We employed a No U-Turn sampler for inference, implemented within the PyMC version 5.4.1 software using two chains per sample. Chains were run for 1500 iterations using 500 samples as burn-in. Convergence checks were performed using the $R^$ statistic. Further details about the theory behind MCMC inference are provided in Appendix D. All scripts necessary for running the analyses conducted in this paper are freely available via GitHub.

All participants provided written informed consent for their data to be used, and all studies were approved by the institutional review boards at the respective institutions.

In Section 3.1, the convergence of the MCMC samplers utilised by the HBR models is assessed.

We evaluate four variants of the HBR method: The first uses a simple Gaussian likelihood, where $μ$ and $σ$ are modelled exactly like $μ$ and $σ$ in Appendix B.3. This model is indicated with an $N$⁠. Model $N$ is the model that was previously the limit of what could be done within HBR in the pcntoolkit. We also evaluate a model with an $So$ likelihood and two variants of the $Sb$ likelihood, which differ only in the way $∈$ and $δ$ are modelled. In the first model $Sb1$⁠, both $∈$ and $δ$ assume constant values throughout the full range of covariates. In model $Sb2$⁠, both $∈$ and $δ$ are linear regressions on the design matrix $Φ$⁠, which in our case contains 5-knot b-spline basis expansion (Boor, 1978) of the age covariate, concatenated with the age and sex covariates. These models are formalised in Table 1, using the notation detailed in Appendix B. Details about the priors can be found in Table 2. Full specifications of the generative models can be found in Appendix C.

In Section 3.2, we compare different types of HBR normative models, with one another and with warped Bayesian Linear Regression (Fraza et al., 2021), which is used as a baseline, abbreviated with W-BLR. The models used in Section 3.3 are the same as in Section 3.2, just applied to different data.

Note that the results in Section 3.2 are based on MCMC estimates of z-scores (see Appendix E). The results in Appendix C are based on the MAP estimates, which were chosen because they are less computationally demanding, and because in practice, the MCMC results were only marginally better.

A main objective in normative modelling is to map observations to z-scores in such a way that healthy variation within those z-sores follows a Gaussian distribution, and that patients fall in the outer centiles of variation. Gaussianity in the outer centiles is compared by the third and fourth moment (skew and kurtosis) of the predicted z-scores. The qq-plots, and a visual inspection of the regression curves further help assess the centiles.

Here a simple analysis of the samplers is performed, by analysing the trajectory of the Gelman-Rubin, or $R^$ statistic (Brooks & Gelman, 1998), which can be thought of as the ratio between the interchain variance and the intrachain variance. The $R^$ statistic is 1 (one) when those are equal, and this is a good indicator that the sampler has converged. In practice, a rule of thumb is that an $R^$ lower than 1.1 is no reason to worry, and an $R^$ less than 1.05 indicates good convergence. To compute the $R^$ statistic, at least two chains need to be sampled. We sampled two chains of 1500 samples using the pcntoolkit, the first 500 of which were used for tuning and were discarded. The last 1000 samples of all chains were stored and used for further analysis.

In Figure 4, the $R^$ statistics of $∈$ (unbroken line) and $δ$ (broken line) are displayed as a function of chain length. The results are averaged over the 10 folds. Overall, the convergence of the $So1$ model, the $Sb1$ model, and the $Sb2$ models are similar on these phenotypes. Except for the WM-hypointensities phenotype—which is very hard to model—all samplers converge with $R^$ very close to 1 within at most 300 iterations.

Since it is a strictly positive measure, we also evaluate the convergence of the samplers on a log-transformed version of the WM-hypointensities, which shows much better convergence.

Note also that the $So$ sampler did not converge for the most difficult phenotype reported in Section 3.3, so we use the $Sb$ parameterisation in that case.

First, we discuss how the HBR with the $Sb$ likelihoods compare to the W-BLR method, and the HBR with the Gaussian likelihood. We assess (i) how well the predicted z-scores follow a Gaussian distribution, by looking at the higher-order moments of the z-scores and at qq-plots; (ii) how good the fit of the percentile lines is visually, by superimposing them on the data; and (iii) how well the models have captured the batch effects, by looking at classification performance. Specifically: for any combination of two sites, we computed the area-under-the-curve (AUC) metric (Hastie et al., 2009) by using the z-scores as classification thresholds, and the batch-indices as labels.

Table 3 forms the first basis of our discussion. For convenience, it is useful to juxtapose Figure 6 or Appendix I, which visualise the phenotypes discussed here. The skew and kurtosis are the third and fourth standardised moment, respectively (Casella & Berger, 2021). As the name suggests, skew describes how much the data is skewed about the mean. The kurtosis measures the tailed-ness of the distribution. A spike distribution has a kurtosis of 0, and a standard Gaussian has a kurtosis of 3. For this reason, the kurtosis values reported here are reduced by 3, which is sometimes named the excess kurtosis. For both measures, the ideal result is thus 0, which suggests a symmetric distribution, with tails as slim as those of a Gaussian.

The best results—results that are closest to 0—are printed in bold font but note that the most direct comparison is between W-BLR, $So$ and $Sb1$⁠. The bold figures on the left of the vertical line that separates W-BLR, $So$ and $Sb1$ from $Sb2$ show the best performing method amongst these methods. The $Sb2$ models are more flexible by construction and have an additional random effect. Phenotypes (rows) that are bold on the right side of that vertical line are those that benefit from this additional flexibilty. For most phenotypes, $So$ or $Sb1$ have the best result, with a few benefiting from the additional flexibility introduced by $Sb2$⁠, but the difference with other methods is sometimes small. The two moments need to be weighted differently as well. For instance, for the Brain-Stem phenotype, the W-BLR has only a marginally smaller skew than $So$⁠, but $So$⁠, $Sb1$ and $Sb2$ have a significantly smaller kurtosis. It is important to remember that the kurtosis is a sum of fourth powers, and the skew is a sum of third powers. Deviations thus generally have a larger effect on the kurtosis than on the skew. For the WM-hypointensities phenotype, the W-BLR method performs on par with the $S$ models regarding the skew, and better when regarding the kurtosis, indicating a better fit on those phenotypes that are also most difficult to model. The W-BLR has inherent batch effects in all aspects of the predicted percentiles, possibly explaining why it fits better to these hard phenotypes. This is further discussed in Section 3.2.1. We emphasise that the comparison is not completely straightforward with the HBR models being more conservative, as described in detail in Section 3.2.1. Despite that, the numbers indicate a fit that is overall better than W-BLR.

Figure 5 further illustrates the good performance of the $So$ and $Sb$ models relative to the $N$ model. The three hard phenotypes are not modelled well at all by the $N$ model. Overall, the $So$ models perform comparable to W-BLR, except for Right-Lateral-Ventricle, where the W-BLR line looks slightly better. Note that this figure shows only a random selected cross-validation fold, whereas the statistics in Table 3 are averaged across all folds. Note also that the log transform does not by itself adequately model non-Gaussianiy in the WM-hypointensities phenotype, indicating that while it addresses non-negativity, it is still necessary to combine this transformation with a principled approach for modelling non-Gaussianity, such as afforded by the SHASH distribution.

Figure 6 visualises the centiles of the fitted HBR models on two hard phenotypes; Right-Lateral-Ventricle, WM-hypointensities, and the log-transformed WM-hypointensities. For brevity, we show only the $Sb2$ model, but the fitted centiles for all models and phenotypes are shown in the appendix (i.e., including $So$ and $Sb1$⁠). Overall, this figure shows that the non-Gaussian models fit the data better. Focusing on area A, the W-BLR the $Sb$ model fits the data well, but the fitted centiles of the Gaussian model (b) cannot accurately model the centiles in the lower range, and many extend implausibly into the negative range (also for e). Focusing on area B, and realising that this phenotype is a volume, and can therefore never be negative, the problem is obvious. Focusing on area C, we see that the W-BLR model must severely warp the input space to be able to accurately model the data at this point. For models (d) and (f), the 0.1’th fitted centile is still estimated to run into the negative domain, towards the range of the data points. This means that the data are not modelled accurately in those regions. As noted above, this motivates the application of the log transformation to the data before fitting which enforces positivity (g-i) and avoids the fitted centiles becoming negative toward the limits of range of the input variables. However, it is still necessary to model non-Gaussianity. Focusing on area D, we see that the Gaussian model applied to the log-transformed data has a poor fit in these regions, while the other two models perform well. These visualisations indicate that the SHASH transformation fits the data better and better accommodates site effects relative to a Gaussian model.

Another important aspect of the HBR method is the way it can deal with batch effects. Unlike the W-BLR method, where one-hot encoded⁶ batch vectors are appended to the design matrix (effectively treating site as a fixed effect), causing the batch effects to directly or indirectly affect the mean, variance, skew, and kurtosis, in HBR the batch effects only affect aspects of the likelihood by explicitly giving the parameter that controls that aspect a random effect. By having a batch-specific offset in $μ$⁠, sampled as a deviation from a group mean, the HBR models’ predictions are ultimately indistinguishable between batches. To assess to which degree the predictions are indistinguishable between batches, we compare the Area Under the Reciever-Operator-Curve (AUC) scores (Hastie et al., 2009), which vary between 0 and 1, in a 1v1 setting. Only site was used as a batch effect in our models, so the batch effects correspond exactly to the site-ids. For each combination of two sites, we compute the AUC score using the sklearn.metrics.roc_auc_score method, where we use the binary class labels as y_true, and the z-score predictions for y_pred. Because an AUC score of 0.5 would indicate perfect indistinguishibility, we plot the absolute deviation from the 0.5 mark in Figure 7.

The comparison between HBR and W-BLR seems to indicate the W-BLR should be preferred in some cases where the modelled response variable takes on an extreme shape. However, the W-BLR method and the HBR differ in the order in which they apply certain transformations, leading to a subtle difference.

The W-BLR method applies the sinh-arcsinh transformation after applying the affine transformation for mean and variance, and the HBR method applies the sinh-arcsinh transformation before applying the affine transformation for $μ$ and $σ$⁠. Because $μ$ and $σ$ potentially contain batch effects, and because in W-BLR, $μ$ influences the shape of the distribution (as illustrated in Fig. 8), the random effects in W-BLR indirectly influence the skew and the kurtosis. The better fit of W-BLR on the Right-Lateral-Ventricle and the WM-hypointensities data in Table 3 may partly be due to this flexibility. We assume that in most cases, batch effects are not expected in the variance, skew, and kurtosis as much as they are in the mean. Especially for features with a variance that is small relative to the mean, a small batch effect in the mean would not constitute a large scaling in the variance. Since the skew and kurtosis are central measures—meaning that they are computed on mean-centred data that are divided by the variance—they are not influenced by scaling or shifting transformations at all. For phenotypes derived from images from a single specific scanner may be overestimated by some constant amount, but it is less likely that the shape of the trajectory of the feature depends much on the scan location itself (although see Section 4 for further discussion about this assumption). Nonetheless, a latent correlation between the site and the measurement could be present. In contrast, the HBR method provides better control because it allows modelling batch effects optionally, but not necessarily (see Appendix B). With W-BLR, the random effects are encoded as covariates, and may therefore influence all parameters of the likelihood. The results in Table 3 should be seen in light of this subtle difference. The W-BLR model may outperform HBR on some phenotypes, but the results listed under HBR are made by more conservative models, which can be controlled better in the sense that the shape parameters in HBR are implicitly regularised under the prior, whereas the shape parameters estimated by W-BLR are effectively unconstrained and can potentially more easily overfit to the data.

To demonstrate the importance of the flexibility of model $S$ likelihood, we take as an example a recent study in which problematic data emerged from a dimensionality reduction (see Fig. 9).

In Zabihi et al. (2022), fMRI data were compressed using an autoencoder (Kingma & Welling, 2013), and then further compressed using UMAP dimensionality reduction (McInnes et al., 2018). The result was a 2d latent representation of brain activation patterns, which followed a somewhat problematic, highly nonlinear, distribution. Positively skewed on the lower end, and negatively skewed on the higher end of the age range, this representation cannot be modelled by HBR with a Gaussian likelihood, and W-BLR is also not designed to deal with these kinds of effects. If we take $Sb1$⁠, we find that we run into the limits of that model as well, because it assumes a fixed skew and kurtosis throughout the age range. With model $Sb2$⁠, we have the flexibility to model heterogeneous skew and kurtosis, and Figure 9 shows that this flexibility is absolutely essential for modelling this phenotype.

We provide an indication here about the computational complexity for the time required to estimate the different variants of our models. Specifically, for the lifespan normative models presented in Section 2.2 above the mean [stdev] time averaged over 10 folds and is specified as hours:minutes. This measures the time required to draw 1000 samples after 500 samples burn in was: $N$ = 2:19 [0:36], $So$ = 6:29 [1:00], $Sb1$ = 6:54 [0:31] and $Sb2$ = 27:26 [1:27]. Another important question in the use of MCMC samplers is determining the number of samples to acquire (i.e., the chain length). We provide a simple evaluation to determine whether 1000 samples after burn-in is sufficient. To achieve this, we compare the empirical mean and variance of different parameters across chains having 500, 1000, and 1500 samples, post burn-in. We also show Monte Carlo estimates for a representative set of parameters in Figure J.36. This shows that the parameters are effectively identical, with perhaps a slight reduction in the Monte Carlo estimates of the variance for some parameters for the short chains. Overall, this provides confidence that the chain length we employ is sufficient for reliable inferences.

In this work, we proposed a flexible method for modelling non-Gaussianity in normative models based on the SHASH distribution. We proposed a novel reparameterisation and developed an efficient MCMC sampling approach for these models. We first applied this method to a large neuroimaging dataset and showed competitive performance with respect to competing methods, while providing full Bayesian inference and full distributions over all model parameters. Then, we applied this method to a highly challenging problem where we aim to model a set of learned parameters from a highly non-linear model (derived from an autoencoder). In this case the classical methods are not able to satisfactorily model the distribution of the brain phenotype, while the proposed method can.

Normative modelling aims to find centiles of healthy variation in neuroimaging phenotypes, much like growth charts in pediatric medicine. Many of those phenotypes follow complex nonlinear, asymmetric distributions; so for a normative model to be generally applicable, it has to be able to model features like skewness and kurtosis. Our extension builds upon the HBR framework for Normative modelling introduced by (Kia et al., 2021), which retrieves a distribution over model parameters given a prior distribution over the parameters and a likelihood over the data. The original paper reported only on results where a Gaussian likelihood was used, but the framework is flexible enough to support any parameterised likelihood. Here, we adapted the HBR framework to make use of a different one, namely the SHASH likelihood, which has additional parameters roughly modelling for skew and kurtosis.

The basic SHASH distribution in Equation 3 has no parameters for location and scale, but those can be added, resulting in the SHASH_o (⁠$So$⁠) distribution, as in Equation 5. We also introduce a reparameterisation called the SHASH_b (⁠$Sb$⁠) distribution, which aims to remove the severe correlations by applying an affine tranformation for location and scale to a standardised SHASH distribution.

Our results on a large neuroimaging dataset show that the $Sb$ reparameterisation in some cases converges better and more reliably than the $So$ distribution. For example, in the analysis of the highly challenging data, shown in Figure 9, the $So$ model did not converge. However, for most phenotypes, the convergence is similar to the original variant. A more important benefit of this reparameterisation is that the parameters of the distribution are more interpretable, as shown schematically in Figure 1 and empirically in Figure 2. We note that the $Sb$ distribution is isomorphic to $So$—the two can model exactly the same set of distributions— which might not give a preference for one distribution over the other in the context of HBR. We also observed a similar sampling performance (at least in terms of convergence) for both. However, the $So$ parameterisation could in some circumstances be preferred over $Sb$⁠, (e.g., for “easy” phenotypes with near-Gaussian shapes) if parameter interpretation is not of primary importance, although the latter achieved slightly better performance in our experiments. However, as noted above, this $So$ model did not converge for the most difficult nonlinear phenotype we considered, in which case $Sb$ is preferred.

We showed that HBR with an $Sb$ likelihood performs equivalently or in some cases slightly better than Warped Bayesian Linear Regression (W-BLR) (Fraza et al., 2021) on most datasets, in terms of Gaussianity of the predicted deviation scores, although W-BLR predictions are better on specific pathological phenotypes. We argue that the W-BLR method fits better to that data due to the fact that batch effects are effectively encoded as a fixed effect and can therefore potentially affect all parameters of the likelihood, which is not always desired (see Fig. 8). The W-BLR method has no option to disable this flexibility, as it is inherent to the method. This is why we argue that the HBR method should be preferred as it supports enabling or disabling random effects in each parameter of the likelihood as a modelling choice. Moreover, the W-BLR method is unable to model the most complex phenotype we evaluated (i.e., derived from the latent representation of an autoencoder) because it does not provide the ability to sufficiently control the skewness and kurtosis across the range of the covariates. In other words, HBR provides more flexible parametric control over the shape of the distributions that the approach is able to model.

We have shown that HBR with $Sb$ likelihood is able to model highly non-Gaussian distributions with intricate trajectories in several parameters in a way that is impossible to replicate with existing techniques. This indicates that HBR will possibly play an important role in pushing the field of Normative modelling forward.

The computational complexity of the likelihood may be the most important limitation, especially for the $Sb2$ variant. It must be noted that the computation time is significantly reduced if fewer batch effects are present. We hope that future releases of software packages like pymc will help reduce the required computational time.

Lastly, all our work is freely available to be used and extended, in the open-source pcntoolkit package.

In the current work, we have not applied the method in a transfer learning setting, but the methods for doing so are described in Kia et al. (2021) and the updates we provide to the pcntoolkit package support this. For transferring a learned model to a new site, one would have to create a factorised approximation of the posterior by using the MCMC samples retrieved earlier. Restricting ourselves to a known distributional form, this can easily be done by standard optimisation techniques on the log likelihood. The factorised posterior is then fixed, and can be used as an informed prior for MCMC sampling with data from a new site. The samples retrieved in this manner can then be used for approximation of z-scores or any other downstream analysis.

We have shown the value of the $So$ and $Sb$ likelihoods in the context of MCMC sampling. The reparameterisation removes much of the correlation in the posterior that is problematic for sampling methods, leading to more stable results. Nonetheless, the MCMC sampling procedure is still very computationally demanding, especially on large datasets with many batch effects. A variational approach (Blei et al., 2017) may now be a viable solution to this problem. In the HBR context, the $Sb$ reparameterisation may be particularly beneficial in this regard because it removes dependencies between the parameters, such that the posterior can be approximated by a factorised distribution. Variational inference is a fast approximation, and MCMC is a slow but (asymptotically) exact method. We believe that in some cases a fast approximation suffices, so variational approximation could be a suitable extension to the present work.

The extension presented here is just the tip of the iceberg of Hierarchical Bayesian modelling. This extension was developed for modelling skewness and kurtosis, and while we showed that this can fit a wide range of distributions, the flexibility of the SHASH distribution is still limited. Specifically, the limitations in tail thickness are easily read from plot 1d. Mixture models could further help increase the expressive power of the HBR method. One could design a mixture of a single SHASH_b (or simply a Gaussian, which might be easier to fit) and two Gamma distributions, one left-tailed and one right-tailed, conveniently located as to act as surrogate tails. The central $S$ component would be dedicated to learning the shape of the bulk of the data, while the Gamma distributions absorb the largest outliers. With the current version of the pcntoolkit, this extension is now within reach, but we leave this for future work.

Multidimensional batch effects are currently implemented as a $D$ dimensional matrix in the pcntoolkit, containing one value per combination of batch effects. The number of values grows exponentially with the number of batch effects. An alternative approach is to model a separate vector for each batch effect, and to learn parameters linking the vectors together. So instead of having $θi,j∼N(μi,j,σ)$⁠, one could have $θi,j∼N(αθi1+βθj2+γθi1θj2)$⁠. The number of required linking parameters grows as $n(n+1)/2∈O(n2)$⁠, much better than the current $O(dn)$⁠, where $d$ is the number of unique values for a batch effect. We expect that the correlation between batch effects is not sufficient to justify the approach, and we believe the most important effects in the batches can be modelled this way. In addition, throughout this work we have assumed that the batch effects are principally evident in the mean, and observed empirically that this provided reasonable performance. It is conceivable, however, that batch effects could also be evident in the variance. While the original Gaussian HBR method can easily handle batch effects in the variance, the use of the SHASH distribution complicates this somewhat; while we could add batch effects to the variance as well, this might interact with the shape and/or heteroskedasticity of the resulting model. We therefore leave this for future work.

We have extended a method for federated normative modelling of neuroimaging data to support non-Gaussian attributes such as skew and kurtosis. To this end, we have introduced a 4-parameter reparameterisation of the SHASH distribution, along with an accompanying software implementation. We have demonstrated that our method performs equivalently or in some cases slightly better than a baseline method on simple and complex data, thereby showing that this method can push the field of normative modelling forward.

Most data used in this manuscript are publicly available and can be downloaded from the repositories providing the data. See Rutherford et al. (2022) for details. Data from the “Thematically Organised Psychosis” sample is available under reasonable request by contacting the responsible PIs (L.T.W. and O.A.A.). A software implementation of the tools presented here is provided in the pcntoolkit package and the scripts for running the experiments are also provided online via GitHub.

A.A.A.d.B., J.M.M.B.: primary analysis, drafting and revising the manuscript. S.M.K.: conceptualisation, supervision, analysis, and drafting and revising the manuscript. S.R., M.Z., C.F., and P.B.: analysis, drafting and revising the manuscript. L.T.W., O.A.A.: provision of data, drafting and revising the manuscript. M.H., C.F.B., and A.M.: conceptualisation, supervision, analysis, and drafting and revising the manuscript.

C.F.B. is director and shareholder of SBGNeuro Ltd. The other authors report no conflicts of interest.

Every phenotype has its own defining characteristics. Some may be severely skewed, while others appear to be nicely Gaussian. We may see strong correlations between age and variance (heteroskedasticity) in one, while another seems to retain a consistent variance throughout the entire age spectrum. Batch effects like sex-specific differences, or site-specific noise characteristics may or may not be present in the data. Characteristics like these can be learned by a sufficiently flexible model. However, Occam dictates that we should prefer a simple model over a complex model, if they perform equally well. A model is therefore optimally defined if it has enough flexibility to capture the parameters of interest, whilst being constrained in every other way. An ideal modelling framework would allow specifying constraints on all aspects of the model, be flexible enough to support a wide range of models, while still remaining concise. Here, a model definition framework is proposed that allows us to concisely write down a wide range of different models. This framework will simplify the later sections of this report. In addition, the framework introduced here aligns with the implementation of the source code in the pcntoolkit software package. The graphical notation used here is slightly adapted from Bishop (2006) chapter 8, where here we indicate the range and name of the iterator variables more explicitly in the plates, i.e., $N$ becomes $n∈{1,...,N}$⁠. For the remainder of this section, we will uphold the nomenclature in Table B.5. Without loss of generality, let us now focus on modelling a single parameter, $θ$⁠. For more complex distributions that take more parameters, like almost any distribution, we just apply what follows to every parameter separately.