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 submission proposes a new model for functional alignment of fMRI datasets from multiple subjects. It combines a one-in-many-out autoencoder with two regularization loss terms (one inspired by GRAE) to develop a model that can encode every subject's data to a shared latent space, from which each subject's data is then decoded by a separate decoder. Experiments are provided that demonstrate benefits to some downstream tasks.