Learning shared neural manifolds from multi-subject FMRI data

Functional magnetic resonance imaging ( fMRI ) is a notoriously noisy measurement of brain activity because of the large variations between individuals , signals marred by environmental differences during collection , and spatiotemporal averaging required by the measurement resolution . In addition , the data is extremely high dimensional , with the space of the activity typically having much lower intrinsic dimension . In order to understand the connection between stimuli of interest and brain activity , and analyze differences and commonalities between subjects , it becomes important to learn a meaningful embedding of the data that denoises , and reveals its intrinsic structure . Specifically , we assume that while noise varies significantly between individuals , true responses to stimuli will share common , low-dimensional features between subjects which are jointly discoverable . Similar approaches have been exploited previously but they have mainly used linear methods such as PCA and shared response modeling ( SRM ) . In contrast , we propose a neural network called MRMD-AE ( manifold-regularized multiple decoder , autoencoder ) , that learns a common embedding from multiple subjects in an experiment while retaining the ability to decode to individual raw fMRI signals . We show that our learned common space represents an extensible manifold ( where new points not seen during training can be mapped ) , improves the classification accuracy of stimulus features of unseen timepoints , as well as improves cross-subject translation of fMRI signals . We believe this framework can be used for many downstream applications such as guided brain-computer interface ( BCI ) training in the future . 1 INTRODUCTION . Brain activation dynamics , as measured by fMRI , exist in an extremely high-dimensional collection space and contain high levels of noise . Noise sources include measurement imprecisions and physiological factors such as a subject ’ s blood pressure variability , head motion , and respiration ( Brooks et al. , 2013 ; Mumford & Poldrack , 2007 ) . Moreover , even after accounting for the idiosyncrasies in noise , neural responses evoked by the same stimulus can vary greatly between different subjects . This makes using fMRI to extract meaningful signals or general trends that are task-relevant quite difficult . Experiments involving fMRI are expensive and usually involve pilot studies conducted in a smaller scale with simpler stimuli and on a smaller group of subjects . Thus , it is desirable to have a method that can learn common trends from pilot data and can generalize to a larger group of subjects or more complex tasks . Like for other high dimensional biomedical data , it is broadly acknowledged that neural activity actually lies in a lower-dimensional latent space . Moreover we hypothesize that several features of these latent spaces are shared across subjects given that measurements are usually task-based and are taken when subjects experience the same stimuli . Dimensionality reduction methods have been used to discover such low-dimensional spaces and facilitate understanding the underlying brain activities . Methods such as principal component analysis ( PCA ) and factor analysis have been the common choice ( Smith et al. , 2014 ) . However , these linear methods are sensitive to noise and restrict interactions in the latent space . Nonlinear dimensionality reduction algorithms are shown to better capture the geometry of high-dimensional biological data ( Moon et al. , 2019 ; Van der Maaten & Hinton , 2008 ) and account for dynamics in neural activities ( Gao et al. , 2020 ) . PHATE ( Moon et al. , 2019 ) , a diffusion based manifold learning method , was shown to capture both global and local geometry of complex biological data , which appears to be a good candidate to model fMRI data . Using a nonlinear manifold such as PHATE to reduce the dimensionality of data is appealing as it can extract meaningful signals , denoise the data , and accelerate downstream analysis . However , unlike a method like PCA which learns a projection operator that can be applied to new data , manifold learning methods are fixed to the input data and do not naturally extend to new data from new tasks or from new subjects . To extend an existing manifold to out-of-sample data , landmark interpolation or Nyström extensions are commonly used . However it is usually unsatisfactory as shown by decremented performance ( Appendix Fig . 1 ) . To tackle these shortcomings , we turn to neural networks such as autoencoders that provide nonlinear dimensionality reduction via a learned parametric function that is readily applicable to new data . While our key goal is learning an informative low-dimensional data manifold shared across multiple subjects , this is not something automatically done by autoencoders . First , there is nothing enforcing autoencoder latent embeddings to respect data manifold geometry , and indeed they often just spread out points for ease of decoding . Second , individual variation often dominates embeddings which prevents a common space to be automatically learned . To address these issues , we propose a manifold-regularized multi-decoder autoencoder ( MRMD-AE ) that can process fMRI data from a group of subjects and extract common latent space representations that respect individual data geometry . Key features of the MRMD-AE include : ( a ) A common encoder that projects fMRI data from any subject on to a shared latent space , and subject-specific decoders that learn to reconstruct data for each subject faithfully . ( b ) A manifold-geometry regularization of the latent space based on a precomputed PHATE embedding . ( c ) Penalties for distances between individual patient encodings of common stimuli to ensure that the latent space is not split into individual embeddings . We show results on three types of tasks in two different datasets . First , we show that we are able to extend the manifold embedding to new timepoints of data not used in training . This shows the extensibility of our manifold in contrast to non-neural network based manifold learning methods . Second , we show that our latent space improves the ability to classify or infer stimulus features based on subject fMRI measurements . Surprisingly , the multiple decoder improves classification accuracy over a common decoder , showing that this framework allows the network to separate common from individual variations . Finally , we show that even untrained , translation between subjects has increased accuracy on withheld timepoints . Further , we show highly improved cross subject translation after training for this task . 2 RELATED WORK . Manifold Learning Manifold learning assumes that high-dimensional ambient data X ∈ Rd lies on a low dimensional manifold . Given a set of datapoints measured in the ambient space , assumed to be sampled from a manifoldM , manifold learning optimizes a low dimensional Euclidean space that encodes the intrinsic geometry ofM and the original data . Popular manifold learning methods include diffusion maps ( Coifman & Lafon , 2006 ) , Laplacian Eigenmaps ( Belkin & Niyogi , 2003 ) and PHATE ( Moon et al. , 2019 ) . We note that multiscale manifold learning methods ( Wolf et al. , 2012 ; Brugnone et al. , 2019 ; Kuchroo et al. , 2020 ) can be used to consider multiple manifolds that capture relations and structure in data at different resolutions . Other methods , such as tSNE ( Van der Maaten & Hinton , 2008 ) and UMAP ( McInnes et al. , 2018 ) , relax the manifold assumption to focus on neighborhood preservation in low dimensions . While these are often considered as variations of manifold learning , they are mostly suitable for visualization of clustered data that does not necessarily have global structure , and typically do not preserve overarching trends or relations between data regions ( see , e.g. , Moon et al. , 2019 ; Gigante et al. , 2019 ) . SAUCIE In ( Amodio et al. , 2019 ) , SAUCIE ( Sparse Autoencoder for Clustering , Imputation , and Embedding ) is an autoencoder-based generative model that manipulated internal representations to force the network to produce specific desired transformations in the data . By imposing these structure on the latent representations of data , it can quickly extract information about datasets more massive than alternative approaches can easily handle . Further , this framework establishes that multiple regularizations , in conjunction , can be used to encourage interpretable latent representations that accentuate biologically reliable relations in data coming from multi-patient cohorts , as demonstrated , e.g. , in immune response profiling ( Zhao et al. , 2020 ) . Inspired by the success of this approach in immunology , we apply a similar multi-objective neural network approach here to process multisubject fMRI data and alleviate some of the main challenges associated with this data type . SRM Shared response modeling ( SRM ) is a linear method to account for differences in functional topographies in fMRI data in a low-dimensional space . SRM uses a singular value decomposition approach to learn a set of low-dimensional signals ( shared response ) common to multiple subject ’ s fMRI data during the same experiment . SRM learns subject-specific transformation to align subject data with the shared model using an orthonormally-constrained weight matrix , similar to that of PCA ( Chen et al. , 2015 ) . 3 PRELIMINARIES . Manifold learning and dimensionality reduction with PHATE To capture the complex interaction in the latent space and balanced global and local structures of highly noisy fMRI data , we use Potential of Heat-diffusion for Affinity-based Transition Embedding ( PHATE ) ( Moon et al. , 2019 ) , a manifold learning method based on diffusion geometry ( Coifman & Lafon , 2006 ) . PHATE computes the diffusion operator which is the probability transition matrix P between data point pairs via an α−decay kernel . Then by raising P to the power of t , the PHATE optimal diffusion time scale , it simulates a t−step random walk over the data graph . Diffusion potential distances D is extracted from P t , which accounts for both global and local data geometry . Finally , metric multi-dimensional scaling ( metric MDS ) ( Abdi , 2007 ) , is applied to D to get the PHATE embedding of input data . PHATE has been shown to effectively capture geometry of high-dimensional biological data and was shown to preserve manifold structure in fMRI data in ( Rieck et al. , 2020 ) . We utilize the extensible framework of PHATE and use PHATE embedding of the training data as regularization applied to the manifold embedding layer as shown in Fig . 1 . Note that we use PHATE over other popular dimensionality reduction methods such as tSNE or UMAP because of its established ability to retain manifold and trajectory structure in biomedical data ( Moon et al. , 2019 ; Kuchroo et al. , 2020 ; Chung et al. , 2020 ; Pappalardo et al. , 2020 ) . Geometry-regularized Autoencoders ( GRAE ) The latent space of the vanilla autoencoders are usually difficult to interpret and not necessarily reflect data geometry ( see , e.g. , Mishne et al. , 2019 ; Moor et al. , 2020 ; Duque et al. , 2020 ) . In order to encourage meaningful data geometry to emerge , or be retained , in latent representations , the training process can be enhanced with appropriate regularizations , often derived from manifold learning principles . For example , Jia et al . ( 2015 ) and Yu et al . ( 2013 ) propose regularization terms that penalize inaccurate preservation of neighborhood relationships in the ambient space . Recent work by Duque et al . ( 2020 ) has shown that the geometry extracted by PHATE can be leveraged to formulate Geometry-Regularized Autoencoders ( GRAE ) that faithfully capture the instrinsic data geometry , while demonstrating advantages in extendibility and invertibility compared to previous approaches . In this work we adapt this approach to fMRI data in conjunction with adjusted architecture and additional regularizations that address further challenges encountered in these settings . Cross-subject batch effects in fMRI Relative to other neuroimaging methods , fMRI is noninvasive , minimal risk , and has broad coverage with high spatial resolution . fMRI experiments commonly use naturalistic stimuli such as movies to probe an array of cognitive mechanisms in conditions that closely mimic the human experience of the world . However , the data generated by such experiments are notoriously high-dimensional and complex , as fMRI is rife with noise from sources like the scanner , subject movement , and the measured blood-oxygen-level-dependent ( BOLD ) signal , which is a secondary effect of neural activity . This makes it difficult to aggregate data across sessions , either within or between subjects in a study . Beyond the noise sources of fMRI data , another problem with aggregating data across subjects to uncover something common about brain functioning pertains to the idiosyncratic nature of functional topographies in the brain . In many brain regions , shared stimuli will reliably evoke similar patterns of neural activity across subjects ( Nastase et al. , 2019 ) . However , that signal of interest could be located in slightly different anatomical locations across subjects or otherwise obscured by noise , causing a mismatch between functional and anatomical locations for a signal . A large body of neuroimaging research centers around rectifying this problem with a variety of linear functional alignment methods by learning either high-dimensional ( Haxby et al. , 2011 ; Busch et al. , 2021 ) or reduced-dimensional ( Chen et al. , 2015 ) common spaces . A solution to this issue must consider both the complexity of the data at the individual subject level , where there is biological and environmental noise , and the broader cross-subject representational idiosyncrasies .

Summary: The paper proposes a model to learn a low-dimensional representation of fMRI data over multiple subjects of the same experiment. The model is built as an auto-encoder, with an encoder shared across subjects, and a separate decoder per subject. The model is regularized so that the first layer of each decoder gives a representation close to a pre-computed manifold embedding. An optional regularization also constrains the shared representation to be similar across subjects.  The paper then proposes a series of experiments to demonstrate the benefits of the learn representation. The experiments consider the tasks of (a) projecting new test samples to the learned manifold, (b) classify some stimulus features from the embedding (decoding task), and (c) predict brain recordings of a new subject.