summaries: The paper presents a deep Markov factor analysis (DMFA), which is a generative model that employs Markov property in a chain of low dimensional temporal embeddings as well as spatial inductive assumptions in order to capture temporal dynamics in functional magnetic resonance imaging (fMRI). The paper shows that DMFA has capability to cluster fMRI data in its low temporal embedding in regards to subject and cognitive state variability. The advantages of DMFA are demonstrated both through synthetic and application data.

summaries: The paper Deep Markov Factor Analysis (DMFA): Towards Concurrent Temporal and Spatial Analysis of fMRI Data proposes a scalable approach to factorize the data of each subject into a product of a small number of spatial and temporal components. Following on the work of (Manning et al., Topographic Factor Analysis, Plos ONE 2014) on topographical factor analysis, each spatial component is given by a radial basis function. The temporal components are modeled from a set of latent satisfying the Markov property. The first latent is conditioned on the data class giving to the model clustering capabilities. Variational inference is used to derive a lower bound of the log-likelihood which is then optimized via stochastic gradient descent. In a first experiment on resting state data (Autism dataset), DMFA is used to obtain subject specific atlases that allow to seggregate ASD and control subjects with an accutacy of 60 % (it is definitely hard to read). On a second experiment (Depression dataset), MDD and controls are exposed to different musical stimuli and DMFA is able to seggregate well the data of different subjects without supervision. In some extend it gives some hint about the kind of stimuli and the conditions of the subjects but results are not so convincing (the authors are honest about this).  In a third experiment they show that DMFA gives better results than other methods in terms of held-out log-likelihood and prediction of the next time point given the history. The code for DMFA is available. Pre-trained model are given and the code to generate the figures from pre-trained model works.

summaries: This paper introduces Deep Markov Factor Analysis (DMFA). DMFA can model spatiotemporal BOLD fMRI time series by using time-varying latent factors that govern time-varying spatial maps consisting of spatial mixtures of RBFs. Latent time series are associated to a categorical clustering variable in order to enable grouping. Certain links in the Bayesian diagram are modeled using deep neural networks. Optimization is done by a variational approach. Experiments are performed on a simulated data set with known ground truth, the ABIDE autism data set, and a Depression study data set. Some clustering of relevant clinical outcomes is observed.

summaries: The authors present deep Markov factor analysis (DMFA) to capture temporal dynamics in functional magnetic resonance imaging data. The method chosen by the authors is relevant since existing methods overlook the highly nonlinear and complex temporal dynamics of neural processes when factorizing their imaging data. The authors claim that their method (DMFA) is able to cluster fMRI data in its low dimensional temporal embedding, which enables validation of fMRI related hypotheses.

summaries: The paper proposes deep Markov factor analysis (DMFA) that uses the Markov process to capture the temporal dynamics in the fMRI dataset and maps the high spatial dimensions to a low-dimensional feature space. DMFA can cluster fMRI responses into “low dimensional temporal embedding” based on subject or cognitive state. The empirical studies on synthetic and real fMRI datasets illustrate that DMFA generates better performance in comparison with state-of-the-art techniques --- such as NTFA and HTFA.

abstractText: Factor analysis methods have been widely used in neuroimaging to transfer high dimensional imaging data into low dimensional, ideally interpretable representations. However, most of these methods overlook the highly nonlinear and complex temporal dynamics of neural processes when factorizing their imaging data. In this paper, we present deep Markov factor analysis (DMFA), a generative model that employs Markov property in a chain of low dimensional temporal embeddings together with spatial inductive assumptions, all related through neural networks, to capture temporal dynamics in functional magnetic resonance imaging (fMRI) data, and tackle their high spatial dimensionality, respectively. Augmented with a discrete latent, DMFA is able to cluster fMRI data in its low dimensional temporal embedding with regard to subject and cognitive state variability, therefore, enables validation of a variety of fMRI-driven neuroscientific hypotheses. Experimental results on both synthetic and real fMRI data demonstrate the capacity of DMFA in revealing interpretable clusters and capturing nonlinear temporal dependencies in these high dimensional imaging data.

authors:
Alexandre Abraham
Michael P Milham
Adriana Di Martino
R Cameron Craddock
Dimitris Samaras
Bertrand Thirion
Gael Varoquaux

title: Deriving reproducible biomarkers from multi-site resting-state data: An autism-based

venue: example. NeuroImage,

authors:
Bahar Azari
Christiana Westlin
Ajay B Satpute
J Benjamin Hutchinson
Philip A Kragel
Katie Hoemann
Zulqarnain Khan
Jolie B Wormwood
Karen S Quigley
Deniz Erdogmus

title: Comparing supervised and unsupervised approaches to emotion categorization in the human brain, body, and subjective experience

venue: Scientific reports,

authors:
Mohammad Taha Bahadori
Qi Rose Yu
Yan Liu

title: Fast multivariate spatio-temporal analysis via low rank tensor learning

venue: In Advances in neural information processing systems,

authors:
Philipp Becker
Harit Pandya
Gregor Gebhardt
Cheng Zhao
C James Taylor
Gerhard Neumann

title: Recurrent kalman networks: Factorized inference in high-dimensional deep feature spaces

venue: In International Conference on Machine Learning,

authors:
Christian F Beckmann
Stephen M Smith

title: Tensorial extensions of independent component analysis for multisubject fmri analysis

authors:
Samuel R Bowman
Luke Vilnis
Oriol Vinyals
Andrew M Dai
Rafal Jozefowicz
Samy Bengio

title: Generating sentences from a continuous space

venue: In 20th SIGNLL Conference on Computational Natural Language Learning,

authors:
Yongjie Cai
Hanghang Tong
Wei Fan
Ping Ji
Qing He

title: Facets: Fast comprehensive mining of coevolving high-order time series

venue: In Proceedings of the 21th ACM SIGKDD International Conference on Knowledge Discovery and Data Mining,

authors:
Laurent Charlin
Rajesh Ranganath
James McInerney
David M Blei

title: Dynamic poisson factorization

venue: In Proceedings of the 9th ACM Conference on Recommender Systems,

authors:
Po-Hsuan Chen
Janice Chen
Yaara Yeshurun
Uri Hasson
James V Haxby
Peter J Ramadge

title: A reduceddimension fmri shared response model

venue: In NIPS,

authors:
Xinyu Chen
Zhaocheng He
Yixian Chen
Yuhuan Lu
Jiawei Wang

title: Missing traffic data imputation and pattern discovery with a bayesian augmented tensor factorization model

venue: Transportation Research Part C: Emerging Technologies,

authors:
Junyoung Chung
Kyle Kastner
Laurent Dinh
Kratarth Goel
Aaron C Courville
Yoshua Bengio

title: A recurrent latent variable model for sequential data

venue: Advances in neural information processing systems,

title: Tensor decompositions: new concepts in brain data analysis

venue: Journal of the Society of Instrument and Control Engineers,

authors:
Pierre Comon
Christian Jutten
Jeanny Herault

title: Blind separation of sources, part ii: Problems statement

venue: Signal processing,

authors:
Cameron Craddock
Yassine Benhajali
Carlton Chu
Francois Chouinard
Alan Evans
András Jakab
Budhachandra Singh Khundrakpam
John David Lewis
Qingyang Li
Michael Milham

title: The neuro bureau preprocessing initiative: open sharing of preprocessed neuroimaging data and derivatives

venue: Frontiers in Neuroinformatics,

authors:
Nicha C Dvornek
Daniel Yang
Pamela Ventola
James S Duncan

title: Learning generalizable recurrent neural networks from small task-fmri datasets

venue: In International Conference on Medical Image Computing and Computer-Assisted Intervention,

authors:
Joset A Etzel
Valeria Gazzola
Christian Keysers

title: An introduction to anatomical roi-based fmri classification analysis

venue: Brain research,

authors:
Amirreza Farnoosh
Hamid Soltanian-Zadeh

title: Information theoretic evaluation of brain connectivity compared to cross correlation using simulated resting state fmri data

venue: Iranian Conference on Electrical Engineering (ICEE),

authors:
Amirreza Farnoosh
Bahar Azari
Sarah Ostadabbas

title: Deep switching auto-regressive factorization: Application to time series forecasting

venue: In Proceedings of the AAAI Conference on Artificial Intelligence,

authors:
Marco Fraccaro
Simon Kamronn
Ulrich Paquet
Ole Winther

title: A disentangled recognition and nonlinear dynamics model for unsupervised learning

venue: In Advances in Neural Information Processing Systems,

authors:
Soham Gadgil
Qingyu Zhao
Adolf Pfefferbaum
Edith V Sullivan
Ehsan Adeli
Kilian M Pohl

title: Spatiotemporal graph convolution for resting-state fmri analysis

title: Overview of functional magnetic resonance imaging

venue: Neurosurgery Clinics,

authors:
R Devon Hjelm
Eswar Damaraju
Kyunghyun Cho
Helmut Laufs
Sergey M Plis
Vince D Calhoun

title: Spatio-temporal dynamics of intrinsic networks in functional magnetic imaging data using recurrent neural networks

venue: Frontiers in neuroscience,

authors:
Matthew D Hoffman
David M Blei
Chong Wang
John Paisley

title: Stochastic variational inference

venue: The Journal of Machine Learning Research,

authors:
Samuel B Hopkins
Jonathan Shi
David Steurer

title: Tensor principal component analysis via sum-of-square proofs

venue: In Conference on Learning Theory,

authors:
Jocelyn V Hull
Lisa B Dokovna
Zachary J Jacokes
Carinna M Torgerson
Andrei Irimia
John Darrell Van Horn

title: Resting-state functional connectivity in autism spectrum disorders: A review

venue: Frontiers in psychiatry,

authors:
Aapo Hyvarinen
Hiroshi Morioka

title: Nonlinear ica of temporally dependent stationary sources

venue: In Artificial Intelligence and Statistics,

authors:
Peiguang Jing
Yuting Su
Xiao Jin
Chengqian Zhang

title: High-order temporal correlation model learning for time-series prediction

venue: IEEE transactions on cybernetics,

authors:
Fikret Işik Karahanoğlu
Dimitri Van De Ville

title: Transient brain activity disentangles fmri resting-state dynamics in terms of spatially and temporally overlapping networks

venue: Nature communications,

authors:
Maximilian Karl
Maximilian Soelch
Justin Bayer
Patrick van der Smagt

title: Deep variational bayes filters: Unsupervised learning of state space models from raw data

venue: In 5th International Conference on Learning Representations,

authors:
Amirali Kazeminejad
Roberto C Sotero

title: Topological properties of resting-state fmri functional networks improve machine learning-based autism classification

authors:
Rahul Krishnan
Uri Shalit
David Sontag

title: Structured inference networks for nonlinear state space models

authors:
Jong-Hwan Lee
Te-Won Lee
Ferenc A Jolesz
Seung-Schik Yoo

title: Independent vector analysis (iva): multivariate approach for fmri group study

authors:
Rebecca J Lepping
Ruth Ann Atchley
Evangelia Chrysikou
Laura E Martin
Alicia A Clair
Rick E Ingram
W Kyle Simmons
Cary R Savage

title: Neural processing of emotional musical and nonmusical stimuli in depression

venue: PloS one,

authors:
Jeremy R Manning
Rajesh Ranganath
Waitsang Keung
Nicholas B Turk-Browne
Jonathan D Cohen
Kenneth A Norman
David M Blei

title: Hierarchical topographic factor analysis

venue: In 2014 International Workshop on Pattern Recognition in Neuroimaging,

authors:
Jeremy R Manning
Rajesh Ranganath
Kenneth A Norman
David M Blei

title: Topographic factor analysis: a bayesian model for inferring brain networks from neural data

authors:
Jeremy R Manning
Xia Zhu
Theodore L Willke
Rajesh Ranganath
Kimberly Stachenfeld
Uri Hasson
David M Blei
Kenneth A Norman

title: A probabilistic approach to discovering dynamic full-brain functional connectivity patterns

authors:
Sarah Parisot
Sofia Ira Ktena
Enzo Ferrante
Matthew Lee
Ricardo Guerrerro Moreno
Ben Glocker
Daniel Rueckert

title: Spectral graph convolutions for population-based disease prediction

venue: In International conference on medical image computing and computer-assisted intervention,

title: Liii. on lines and planes of closest fit to systems of points in space

venue: The London, Edinburgh, and Dublin Philosophical Magazine and Journal of Science,

title: Region of interest analysis for fmri

venue: Social cognitive and affective neuroscience,

authors:
Maria Giulia Preti
Thomas AW Bolton
Dimitri Van De Ville

title: The dynamic functional connectome: State-of-the-art and perspectives

authors:
Rajesh Ranganath
Chong Wang
Blei David
Eric Xing

title: An adaptive learning rate for stochastic variational inference

authors:
Emile Richard
Andrea Montanari

title: A statistical model for tensor pca

authors:
Hugo Richard
Luigi Gresele
Aapo Hyvärinen
Bertrand Thirion
Alexandre Gramfort
Pierre Ablin

title: Modeling shared responses in neuroimaging studies through multiview ica

venue: In Advances in Neural Information Processing

authors:
Mark Rogers
Lei Li
Stuart J Russell

title: Multilinear dynamical systems for tensor time series

authors:
Eli Sennesh
Zulqarnain Khan
Yiyu Wang
J Benjamin Hutchinson
Ajay Satpute
Jennifer Dy
Jan-Willem van de Meent

title: Neural topographic factor analysis for fmri data

venue: Advances in Neural Information Processing Systems,

authors:
Chandan Singh
Beilun Wang
Yanjun Qi

title: A constrained, weighted-l1 minimization approach for joint discovery of heterogeneous neural connectivity

venue: graphs. CoRR,

authors:
John Z Sun
Dhruv Parthasarathy
Kush R Varshney

title: Collaborative kalman filtering for dynamic matrix factorization

venue: IEEE Transactions on Signal Processing,

authors:
Lei Sun
Yanfang Xue
Yining Zhang
Lishan Qiao
Limei Zhang
Mingxia Liu

title: Estimating sparse functional connectivity networks via hyperparameter-free learning model

venue: Artificial Intelligence in Medicine,

authors:
Lijun Sun
Xinyu Chen

title: Bayesian temporal factorization for multidimensional time series prediction

authors:
Koh Takeuchi
Hisashi Kashima
Naonori Ueda

title: Autoregressive tensor factorization for spatio-temporal predictions

venue: IEEE International Conference on Data Mining (ICDM),

title: Functional interactions as big data in the human

venue: brain. Science,

authors:
Erik-Jan Van Kesteren
Rogier A Kievit

title: Exploratory factor analysis with structured residuals for brain network data

venue: Network Neuroscience,

authors:
M Alex O Vasilescu
Demetri Terzopoulos

title: Multilinear independent components analysis

venue: IEEE Computer Society Conference on Computer Vision and Pattern Recognition (CVPR’05),

authors:
Lebo Wang
Kaiming Li
Xu Chen
Xiaoping P Hu

title: Application of convolutional recurrent neural network for individual recognition based on resting state fmri data

authors:
Mingliang Wang
Jiashuang Huang
Mingxia Liu
Daoqiang Zhang

title: Modeling dynamic characteristics of brain functional connectivity networks using resting-state functional mri

venue: Medical image analysis,

authors:
Manuel Watter
Jost Springenberg
Joschka Boedecker
Martin Riedmiller

title: Embed to control: A locally linear latent dynamics model for control from raw images

authors:
Liang Xiong
Xi Chen
Tzu-Kuo Huang
Jeff Schneider
Jaime G Carbonell

title: Temporal collaborative filtering with bayesian probabilistic tensor factorization

venue: In Proceedings of the 2010 SIAM international conference on data mining,

authors:
Weizheng Yan
Vince Calhoun
Ming Song
Yue Cui
Hao Yan
Shengfeng Liu
Lingzhong Fan
Nianming Zuo
Zhengyi Yang
Kaibin Xu

title: Discriminating schizophrenia using recurrent neural network applied on time courses of multi-site fmri

venue: data. EBioMedicine,

authors:
Hsiang-Fu Yu
Nikhil Rao
Inderjit S Dhillon

title: Temporal regularized matrix factorization for high-dimensional time series prediction

heading: 1 Introduction

text:
Functional magnetic resonance imaging (fMRI) has been extensively used in cognitive neuroscience to study brain structures and their interactions at rest or during a cognitive task [Glover, 2011]. fMRI also provides insights into how brain’s functional connectivity changes during different experimental conditions [Preti et al., 2017, Azari et al., 2020]. However, due to high dimensional nature of fMRI data (tens of thousands of voxels in few seconds-long sessions), analyzing functional connectivity of brain could get very challenging [Turk-Browne, 2013]. Most classical methods employ region of interest (ROI)-based approaches to reduce size of data for processing [Poldrack, 2007, Etzel et al., 2009, Farnoosh and Soltanian-Zadeh, 2017, Gadgil et al., 2020]. However, averaging across many voxels within each region could wash out signals from small number of task-relevant voxels with noise from non-activated voxels, therefore results in a loss of information [Poldrack, 2007].
35th Conference on Neural Information Processing Systems (NeurIPS 2021)
Recently, few approaches have been proposed, based on probabilistic generative models, for topographic factorization of fMRI data into a weighted summation of few localized activation sources (i.e., temporal weights and topographic spatial factors), among which topographic factor analysis (TFA) [Manning et al., 2014b] and its multi-subject extension, hierarchical TFA (HTFA) [Manning et al., 2018], and neural TFA (NTFA) [Sennesh et al., 2020] are the most noted ones. This factorization serves as a necessary preparation step for subsequent statistical analysis that can effectively characterize subject- and stimulus-level variations and reveal task- or cognitive state-related networks in brain. However, TFA approaches assume a prior in which temporal weights are conditionally independent as a function of time, which means they do not encode temporal dynamics. Given the non-linearity and complex time-dependencies inherent in fMRI, a model is required that can capture and represent these dependencies.
In this paper, we propose deep Markov factor analysis (DMFA)1, a Bayesian model for factorization of fMRI data that learns a deep generative Markovian prior to reason about nonlinear temporal dynamics. This is realized by a chain of low dimensional temporal embeddings related through neural networks. This prior is further augmented by a discrete latent for multimodal dynamical estimation, and clustering subject- and task-level variations directly in its low dimensional temporal embedding. To accommodating high spatial dimensionality, DMFA generatively parameterize spatial factors from a low dimensional spatial latent through neural networks.
We evaluate the performance of DMFA on a synthetic and two real large-scale fMRI datasets. Our experiments demonstrate that DMFA uncovers meaningful clusters in these data and achieves better predictive performance for unseen data relative to the state-of-the-art.

heading: 2 Related Work

text:
Factor Analysis in fMRI: Factor analysis in neuroimaging includes a wide range of approaches for reducing data dimensionality to facilitate their interpretability and computational tractability. Principal component analysis (PCA) [Pearson, 1901] and independent component analysis (ICA) [Comon et al., 1991] are among the most well-known classical factor analysis methods. To accommodate tensor data and mitigate scalability issues, multilinear versions of PCA and ICA have been proposed in Vasilescu and Terzopoulos [2005], Beckmann and Smith [2005], Cichocki [2011], Richard and Montanari [2014], Hopkins et al. [2015]. Specifically, for multi-subject fMRI study, Lee et al. [2008] proposed independent vector analysis (IVA) and Richard et al. [2020] developed MultiView ICA to model shared responses. Likewise, Chen et al. [2015] developed a shared response model (SRM) for aggregating multi-subject fMRI data and highlighting group differences, and Van Kesteren and Kievit [2021] incorporated structured residuals into the exploratory factor analysis (EFA) framework. Karahanoğlu and Van De Ville [2015] deconvolved hemodynamic response from rest fMRI time series and then performed temporal clustering on the resulting whole-brain innovation signals to recover the corresponding spatial patterns. However, spatial factors obtained by these methods are unstructured, often have the same size as the images in the original dataset, and may include many small and large voxel clusters across the brain, therefore are not directly interpretable [Manning et al., 2014b]. More important, these methods are permutation invariant along temporal dimension (i.e., do not assume any relationships between temporal dimensions), therefore do not model temporal dynamics [Yu et al., 2016].
To enhance spatial interpretability, Manning et al. [2014b] introduced topographic factor analysis (TFA), a probabilistic technique that casts each brain image as a weighted sum of Gaussian spatial factors. In this way, each Gaussian blob (a.k.a topographic factor) can be thought of as a brain region for which activation levels are estimated over time, and is easily interpreted through its set of parameters. Later on, authors proposed a hierarchical version of TFA (HTFA) [Manning et al., 2018] to incorporate data from multiple subjects and enable hypothesis testing across subjects by applying hierarchical Gaussian priors for global and subject-specific parameters of temporal and topographic factors. These linearly-dependent hierarchical Gaussian priors, though, favor estimation of unimodal distributions and limit the model’s expressivity. Recently, Sennesh et al. [2020] proposed neural TFA (NTFA) for task fMRI, which extends TFA by incorporating neural networks onto its framework. NTFA assumes separate latent embeddings for participants and stimuli and map them into the temporal and spatial latents with neural networks. However, methods in this line of work
1The source code is available at https://github.com/ostadabbas/DMFA
essentially assume a prior in which temporal weights are conditionally independent as a function of time, which means that they do not encode temporal dynamics.
Dynamical Factorization: A number of matrix/tensor factorization approaches have been proposed for modeling temporal dynamics in sequential data. These methods employ linear-Gaussian statespace models [Sun et al., 2014], autoregressive temporal regularizer [Bahadori et al., 2014, Yu et al., 2016, Takeuchi et al., 2017], and multilinear dynamical systems [Rogers et al., 2013, Cai et al., 2015, Jing et al., 2018]. In contrast to these methods which provide point estimates, Bayesian dynamical factorization has been proposed in Xiong et al. [2010], Charlin et al. [2015], Sun and Chen [2019], in which linear temporal dynamics are applied to factor latents. While some of these methods have been successful in dynamical modeling of sequential data, they are less effective for high dimensional spatial data like fMRI as they do not explicitly adopt any structural constraints for spatial factors. Moreover, their linear dynamical assumptions lack the capacity to characterize nonlinear dependencies.
Motivated by recent advances in deep learning, several studies have implemented neural networks into Gaussian state space models for nonlinear dynamical modelling [Krishnan et al., 2015, Watter et al., 2015, Chung et al., 2015, Karl et al., 2017, Krishnan et al., 2017, Fraccaro et al., 2017, Becker et al., 2019, Farnoosh et al., 2021]. A common practice is to learn a temporal latent model followed by a mapping to data space, a.k.a. a decoder, and an encoder for amortized variational inference. However, this encoding/decoding framework is not tractable in very high dimensional spatial data like fMRI. To be more specific, it is computationally intensive to feed high dimensional data for amortized estimation, or map to them directly from a latent space (see Section 4.1 for a detailed discussion). Also, several studies have employed recurrent neural networks (RNNs) for temporal analysis of fMRI data [Hjelm et al., 2018, Dvornek et al., 2018, Yan et al., 2019, Wang et al., 2019]. RNN-based methods are not probabilistic and can only process reduced dimensional ROI data. Although convolutional RNN models can capture both temporal and spatial correlations, by construction they do not provide the familiar spatial correlation maps that we expect in neuroimaging analysis and that are necessary for neuroscientific interpretability. In addition, RNN model are harder to train on long sequences due to the vanishing/exploding gradient problem as they feed parameters sequentially. While this problem is much alleviated in LSTM, GRU, and Transformer models, these architectures still result in a huge computational graph and intense GPU memory consumption on long sequences, which leads to very long training times. In contrast, DMFA, as will be explained below, samples from a distinct posterior distribution at each time point and can fit arbitrarily long sequences without any issues (by allowing parallel computations over time), and at the same time infers smooth spatial factors automatically.

heading: 3 Deep Markov Factor Analysis (DMFA)

text:
Let’s assume a corpus of N fMRI data {Yn |n ∈ [1:N ]} for Yn ∈ RT×V , where T and V are the number of time points and voxels, respectively. DMFA defines a hierarchical dynamical deep generative model over this corpus in a factorization framework:
Yn ∼ Norm(W>n Fn, σYI), Wn ∼ Norm(µWθ (ZWn ), σWθ (ZWn )),
Fn = [ fk,v := RBFk(v; ρk, γk) ] ,
ρ, γ ∼ Norm(µFθ(zFn), σFθ(zFn)), zFn ∼ Norm(0, I), ZWn ∼ pθ(ZWn |Cn), Cn ∼ Cat(π).
where Fn ∈ RK×V are K V spatial factors and Wn ∈ RK×T are their associated temporal weights, consistent with a factorization framework. pθ(ZWn |Cn) is a deep generative Markovian prior over a set of low dimensional temporal latents ZWn = {zWn,t | t ∈ [0:T ]}
conditioned on a discrete latent Cn for clustering. The temporal weights Wn are sampled from a Gaussian distribution whose mean and covariance are a function of ZWn and are parameterized by neural networks µWθ and σ W θ , respectively. The spatial factors Fn are determined by a set of K radial basis functions (RBFs) whose parameters {ρ, γ} are sampled from a Gaussian distribution parameterized by neural networks µFθ and σ F θ from a low dimensional spatial latent z F n. All networks have parameters, which are collectively denoted by θ. Finally, σY denotes observation noise. The graphical model for DMFA is depicted in Fig. 1.
Clustering Latent: We assume that each fMRI data Yn belongs to a specific cluster s ∈ [1:S] which is declared by its associated discrete latent Cn ∼ Cat(π) sampled from a categorical prior with assignment probability vector π. This discrete latent conditions the first temporal latent zWn,0 resulting in a Gaussian mixture distribution over this temporal latent:
pθ(z W n,0|Cn = s) = Norm(µs,Σs),
where µs and Σs are determined by cluster assignment.
Deep Generative Markovian prior: We assume that the temporal latents ZWn are related through a Markov chain, for which the transition probability pθ(zWn,t|zWn,t−1) = Norm(µZθ(zWn,t−1), σZθ(zWn,t−1)) is a Gaussian distribution whose mean and covariance are parameterized by neural networks from zWn,t−1. We further blend the estimated mean from neural network with a linear transformation of zWn,t−1 to support both nonlinear and linear transitions (we dropped n for brevity):
µZθ(z W t−1) = (1− g) µ Z,L θ (z W t−1) + g µ Z,NL θ (z W t−1)
where µZ,Lθ (·) represents a linear mapping, µ Z,NL θ (·) is the nonlinear mapping of neural network, and g ∈ [0, 1] is a weighting vector estimated from zWt−1 with another neural network. Temporal Weights & Spatial Factors: As with the transition model, we assumed Gaussian distributions for temporal weights, such that their means and covariances are a function of their associated temporal latents:
wn,t ∼ Norm(µWθ (zWn,t), σWθ (zWn,t)), where µWθ and σ W θ are neural network functions. The high spatial dimensionality in fMRI data encourages the need for a hierarchical analysis that summarizes spatial factors with fewer parameters. Therefore, consistent with Manning et al. [2014b,a, 2018], we represent each spatial factor as a radial basis function (i.e., a Gaussian blob):
fkv := RBFk(v; ρk, γk) = exp ( − ‖pv − ρk‖ 2
exp (γk)
) ,
where pv ∈ R3 denotes the position of voxel with index v, and ρk, γk ∈ R3 denote the center and extent of the Gaussian blob, respectively. The set of all RBF parameters {ρk, γk | k ∈ [1:K]} ∼ Norm(µFθ(z F n), σ F θ(z F n)) are sampled from a Gaussian distribution whose mean and covariance are a function of the shared spatial latent zFn. Introducing z F n as a low dimensional spatial embedding encourages estimation of a multi-modal distribution among spatial factors.

heading: 3.1 Learning DMFA with Variational Inference

text:
We train DMFA with stochastic variational methods [Hoffman et al., 2013, Ranganath et al., 2013]. These methods approximate the posterior pθ(W,ZW, zF,C, ρ, γ|Y ) using a variational distribution qφ(W,Z
W, zF,C, ρ, γ), where φ denotes parameters of the variational model, by maximizing a lower bound on the log-likelihood of the data:
L(θ, φ) = log pθ(Y )− KL ( qφ(W,Z W, zF,C, ρ, γ) || pθ(W,ZW, zF,C, ρ, γ|Y ) ) .
By maximizing this bound with respect to the parameters θ and φ, we learn a deep generative model and perform Bayesian inference, respectively.
Parameterizing Variational Distribution: We assume a fully factorized variational distribution and introduce trainable variational parameters for each data as mean and covariance of a diagonal Gaussian:
qφ(W,Z W, zF,C, ρ, γ) = N∏ n=1 q(Cn)q(zWn,0)q(z F n) K∏ k=1 q(ρn,k, γn,k) T∏ t=1 q(zWn,t)q(wn,t),
where q(·) for continuous latents are parametric Gaussian distributions. Note that due to the high dimensionality of fMRI data (V N , i.e., number of voxels are much larger than number of data), we do not utilize an amortized inference (where a single parametric function maps each data to a set of variational parameters) as it could be computationally very expensive and result in over-fitting (see Section 4.1). Although we can define variational parameters for categorical distributions q(Cn), we approximate it with posterior p(Cn|zWn,0) to compensate information loss induced by the mean-field approximation:
q(Cn = s) ' p(Cn = s|zWn,0) = p(Cn = s)p(zWn,0|Cn = s)∑S s=1 p(Cn = s)p(z W n,0|Cn = s)
This approximation has a two-fold advantage: spares the model additional trainable parameters for variational distribution, and further links variational parameters of qφ(zWn,0) to generative parameters of pθ(zWn,0) and pθ(c), hence results in a more robust learning and inference algorithm.
Derivation of Evidence Lower Bound (ELBO): Once the generative model and variational distribution are determined, we can derive the ELBO which contains a reconstruction term and regularization terms for each latent to bring their variational posterior as close as possible to their generative prior: Ln(θ, φ) = Eqφ(Wn)qφ(ρn,γn) [ log pθ(Yn|Wn, Fn) ] (Likelihood)
− ∑ Cn qφ(Cn) KL ( qφ(z W n,0)||pθ(zWn,0|Cn) ) − KL ( qφ(Cn)||p(C) ) (Clusters)
− ∑ t Eqφ(zWn,t−1) [ KL ( qφ(z W n,t)||pθ(zWn,t|zWn,t−1) )] (Markov Transitions)
− ∑ t Eqφ(zWn,t) [ KL ( qφ(wn,t)||pθ(wn,t|zWn,t) )] (Weights)
− ∑ k Eqφ(zFn) [ KL ( qφ(µn,k, γn,k)||pθ(µn,k, γn,k|zFn) )] − KL ( qφ(z F n)||p(zF) ) . (Factors)
We compute Monte Carlo estimate of gradient of the ELBO using a reparameterized sample from the variational distribution of continuous latents. For the discrete latent, Cn, we compute the expectations over qφ(Cn) by summing over all the possibilities, hence no sampling is performed. We analytically calculate the KL terms of ELBO for both multivariate Gaussian and categorical distributions, which leads to lower variance gradient estimates and faster training as compared to e.g., noisy Monte Carlo estimates often used in the literature.

heading: 4 Training Details

text: We described the network architectures of DMFA in Table 1, where Dt and De are the dimensions of hidden layers for (µZθ, σ Z θ) and (µ W θ , σ W θ ), respectively. We did all the programming in PyTorch v1.3 and used Adam optimizer with learning rate of 0.01. We initialized all the parameters randomly except for locations of Gaussian blobs for which we set initial values to local extrema in their averaged fMRI data. We clipped Gaussian blob parameters to the confines of brain if needed. We used a linear KL annealing schedule, [Bowman et al., 2016], which increased from 0.01 to 1 over 100 epochs. We learned/tested all models on an Intel Core i7 CPU @3.7 GHz with 8 Gigabytes of RAM, which proves tractability of the learning process. Per-epoch training time varied from 0.1 minutes in small datasets to 6.0 minutes in larger experiments and 200 epochs sufficed for most experiments.

heading: 4.1 Discussion of Parameter Count for DMFA

text:
The number of learnable parameters for variational distribution in DMFA is O(NTK). DMFA has O(KDe+Dt) parameters for its temporal generative model and O(K) parameters for its spatial generative model. Note that the clustering latent, C, does not impose additional parameters to the variational distribution, while only adds O(S) parameters to the temporal generative model.
While DMFA introduces extra features and more complex modeling assumptions for fMRI experiments compared to TFA methods of Manning et al. [2014b,a, 2018], (i.e., infers nonlinear temporal dynamics and performs clustering), we emphasize that it has the same order of parameters as these methods. TFA methods similarly have O(NTK) parameters as they employ a fully factorized variational distribution. We also want to highlight that DMFA is tractable in very high dimensional spatial data like fMRI in contrast to nonlinear state-space models of Watter et al. [2015], Krishnan et al. [2017], Fraccaro et al. [2017], Karl et al. [2017], Becker et al. [2019]. The encoder/decoder structure in these works, i.e., qφ(Wn|Yn)/pθ(Yn|Wn), immediately scales both generative and variational parameters to at least O(KV), where V ∼ 105 NT in fMRI data, hence causes extensive computational burden and more importantly overfitting. Furthermore, these methods do not learn a generative model for spatial factors, i.e., pθ(ρ, γ), and as a result are not able to reason about subject-level variabilities in this respect. We overcome these challenges in DMFA by carefully designing our non-amortized variational inference and imposing functional form assumptions on spatial factors in a factorization framework. The proposed learning and inference algorithm keep generative parameters in O (KDe + Dt) O(KV), and variational parameters in O(NTK), where NT ∼ 102 − 105, yield an observation to parameter ratio of O( NTVNTK ) = O( V K ) for all the experiments, therefore permit an efficient learning process.

heading: 5 Experimental Results

text: We analysed the performance of DMFA on a synthetic and two large-scale fMRI datasets. In the first experiment, we verified DMFA’s capability in recovering the true clusters and capturing the underlying temporal dynamic in a synthesized fMRI dataset (visualized in Fig. 2). Next, we evaluated the performance of DMFA on a large scale resting-state fMRI data, Autism dataset [Craddock et al., 2013], and a task fMRI data, Depression dataset [Lepping et al., 2016]. We assessed the clustering feature of DMFA on both datasets in terms of disease and cognitive state separation tasks, visualized in Fig. 3 and Fig. 4. We further provided a quantitative comparison with two state-of-the-art Bayesian generative models for fMRI data, HTFA [Manning et al., 2018] and NTFA [Sennesh et al., 2020], and a deep state-space model, RKN [Becker et al., 2019], in terms of synthesis quality of the generative models on both datasets by computing held-out log-likelihood and prediction accuracy, respectively, in Table 2.

heading: 5.1 Synthetic fMRI

text: We generated synthetic fMRI data using a MATLAB package provided by Manning et al. [2014b], which is known to be useful for analysing fMRI models. The synthesized brain image for each trial (time point) is a weighted summation of a number of radial basis functions (spatial factors) randomly located in brain. This synthesized fMRI data is then convolved with a hemodynamic response function (HRF) and zero-mean Gaussian noise with a medium-level signal-to-noise ratio is added. Here, we considered 30 activation sources (spatial factors) randomly located in a standard MNI-152-3mm brain template with roughly V = 270, 000 voxels and 150 trials. We randomly split these 30 activation sources into 3 groups, each having 10 of the Gaussian blobs. These three groups of sources are periodically activated according to some random weights, one after another, for 5 trials. We generated non-overlapping sequences of T = 5 time points from this synthetic fMRI data. This resulted in 10 data points for each activation group (i.e., N = 30). To train DMFA, we set T = 5, K = 30, Dt = 2, De = 8, S = 3, and σY = 0.01. As depicted in Fig. 2 (a), DMFA was able to successfully recover the 3 clusters of activation that were present in this dataset totally unsupervised. The dynamical trajectory of each cluster mean in the temporal latent (i.e., µzWt |µzWt−1 ,C) is predicted sequentially from the learned generative model and is visualized in the bottom-right of Fig. 2 (a), and appears to be partitioned into three consecutive rotational dynamics. This is consistent with the periodic activations of sources in data clusters (which come in tandem). Predictions of the learned generative model for a selected activation source are visualized in Fig. 2 (d) for the next 50 time points, estimated as follows: wt ∼ p(wt|zt), where zt ∼ p(zt|zt−1) for t = {151, . . . , 200}. These predicted samples perfectly follow hemodynamic response function, confirming DMFA’s capacity in capturing the underlying nonlinear HRF by using neural networks in its deep temporal generative model.

heading: 5.2 Autism Dataset

text:
We used the publicly available preprocessed resting state fMRI (rs-fMRI) data from the Autism Brain Imaging Data Exchange (ABIDE) collected at 16 international imaging sites [Craddock et al., 2013]. This dataset includes rs-fMRI imaging from 408 individuals suffering from Autism Spectrum
Disorder (ASD), and 476 typical controls. Each scan has T = 145 ∼ 315 time points at TR = 2, and V = 271, 633 voxels. We split the signals into sequences of 75 time points. We take two approaches to evaluate the performance of DMFA in separating ASD from control: (i) Cluster data directly in the low dimensional latent, ZW, using the clustering feature of DMFA (called DMFA), (ii) Extract functional connectivity matrices, [Hull et al., 2017], from learned temporal weights, W , followed by a 10-fold SVM for classification (called DMFA+SVM). As baselines, we performed a 10-fold SVM classification on extracted connectivity matrices from (i) averaged signals of 116 ROIs in automatic anatomical labeling (AAL) atlas [Kazeminejad and Sotero, 2019] (called ROI+SVM) and (ii) a hundred time courses obtained with nonlinear spatial ICA of Hyvarinen and Morioka [2017] (called ICA+SVM). Several studies have been done on this dataset to differentiate ASD group from control [Abraham et al., 2017, Parisot et al., 2017, Singh et al., 2017, Kazeminejad and Sotero, 2019, Wang et al., 2021, Sun et al., 2021], all of which using supervised methods, and could achieve accuracies up to 72% using the signals extracted from anatomically labeled regions in the brain by carefully splitting data to be as homogeneous as possible and reducing site-related variability.
We set T = 75, K = 100, De = 15, Dt = 5, S = 2, σY = 0.01, and trained DMFA for 200 epochs on the entire dataset (Full), and also datasets from 9 imaging sites (with more balanced datasets) separately: Caltech, Leuven, MaxMun, NYU, SBL, Stanford, UM, USM, Yale. As shown in Fig. 3 (Right), DMFA and DMFA+SVM outperformed baselines in Caltech, SBL, Stanford, and Yale (at p-value < 0.005, corrected for multiple comparisons), while ROI+SVM and ICA+SVM only performed better in NYU dataset. DMFA+SVM performed slightly better than baselines on the entire dataset (Please note that DMFA is a clustering approach, hence, no error bars are provided for it in Fig. 3). Clustering results for Caltech, Maxmun, SBL, and Stanford are shown in Fig. 3 (Left) in which ASD and control seems to be partially separable (see Fig. S1 in supplementary for more visualization results).

heading: 5.3 Depression Dataset

text:
In this dataset [Lepping et al., 2016], 19 individuals with major depressive disorder (MDD) and 20 never-depressed control participants listened to standardized positive and negative emotional musical and nonmusical stimuli during fMRI scanning. Each participant underwent 3 musical, and 2 nonmusical runs each for 105 time points at TR=3 with V = 353, 600 voxels. During each run, each stimulus type (positive, and negative) was presented for 33 seconds (∼ 11 time points) interleaved with instances of neutral tone of the same length. We discarded instances of neutral tone, and split each run into non-overlapping sequences of T = 6 time points in agreement with stimuli design (each stimuli block is split into two sequences). In other words, each run has 4 sequences associated with “positive stimuli”, and 4 with “negative stimuli” resulting in a total of 8 data points for each run.
In the first experiment, we trained DMFA on the entire musical runs (N = 39 × 3 × 8 = 936) by setting T = 6, K = 100, De = 15, Dt = 5, σY = 0.001 for 200 epochs. The results are shown in Fig. 4 (a, Left). We observed that DMFA fully separated the data points associated with each subject into distinct clusters across the low dimensional temporal latent space. In other words, DMFA was able to re-unite pieces of signals associated with each subject without any supervision. More importantly, DMFA was able to partially separate the data points associated with MDD group from control. As seen in Fig. 4 (a, Left), the MDD group data points are fairly populated in the center of temporal latent while control group are dispersed across the latent space. However, DMFA was not able to meaningfully separate “negative” and “positive” music pieces in its low-dimensional latent space from a subject-level perspective, since the variation between runs of a subject dominates the stimulus-level variation. For this reason, in a second experiment, we focused on 5 subjects, and their first musical run from both MDD and control group and trained DMFA respectively. Again, as expected, data points from each subject were distinctly clustered in the latent space (see middle and right columns in Fig. 4 (a)). Additionally, DMFA was able to fit two partially separating Gaussians to “positive”, and “negative” stimuli per subject. However, since the number of data points for each subject and run is limited, the significance of these clusters are not conclusive. A dataset with longer runs could possibly answer that.

heading: 5.4 Comparison with State-of-the-Art

text:
We further evaluated DMFA against HTFA [Manning et al., 2018], an established probabilistic generative model for multi-subject fMRI analysis, which uses unimodal Gaussian priors for both temporal weights and spatial factor parameters, and its neural network-based extension NTFA [Sennesh et al., 2020] in terms of held-out log-likelihood, and a state-of-the-art deep state-space model, recurrent Kalman networks (RKN) [Becker et al., 2019], in terms of next-time-point prediction accuracy of test set (see Table 2). Held-out likelihood is an established metric for evaluating generative models (in lack of a ground-truth) and measures how probable an unseen data is under the generative model. For next-time-point prediction, we adopt a rolling prediction scheme as in Chen et al. [2019], and predict next time point on the test set sequentially from historical data using the generative model and spatial factors learned on the train set.
We split our data to train and test as follows. We used the Caltech site subset from autism dataset and split each subject’s fMRI time series into two half (each with T = 70). We trained the models on the first half, and tested on the second half. For the depression dataset, we considered 4 sequences from each subject’s run for training and tested on the remaining 4 sequences. We set K = 100 spatial factors for all the models and trained each with their default hyperparameters. For the RKN model, we fit AAL atlas ROI signals as it is intractable to feed raw high-dimensional data. We reported average normalized prediction error (NRMSE%) for both RKN and DMFA on AAL atlas regions to make a fair comparison. For the synthetic fMRI, we set K = 30 and picked half of the data for test, and used the 30 Gaussian blob activations as ROI signals for reporting prediction accuracy in RKN and DMFA. The results are shown in Table 2, which proves that DMFA resulted in models with higher likelihood on the test sets, hence it is a better fit when compared to NTFA and HTFA models. DMFA also resulted in better predictions on the test sets of autism, depression, and synthetic datasets when compared with RKN.

heading: 5.5 Ablation Study

text: We performed two ablation studies to further evaluate (and quantify) the impact of deep neural networks and Markov temporal apriori. First, we trained a version of DMFA in which we removed nonlinear activation functions (i.e., PReLU), denoted as DMFA. As reported in Table 2, in our datasets, this resulted in a decrease in held-out likelihood. Also, the inferred temporal latents lacked the interpretable patterns seen in Fig. 3 and Fig. 4 in terms of clinical/cognitive separability as latents collapsed into the conditional Gaussian priors. Second, we trained a version of DMFA in which we removed the temporal connections (i.e., temporal transition model), denoted as DMFA. Again, this resulted in a drop in held-out likelihood as reported in Table 2, however, the temporal latents preserved their interpretable patterns.

heading: 6 Conclusion

text: We presented deep Markov factor analysis, DMFA, a new probabilistic model for robust factor analysis of high dimensional fMRI data. We employed a chain of low dimensional Markovian latents connected by deep neural networks and conditioned on a discrete latent as a state-space embedding for temporal weights to account for nonlinear dynamics, enable data clustering in a low dimensional space, and provide informative visualizations about data. To tackle high spatial dimensionality in fMRI, we employed a low dimensional spatial embedding to generatively parameterize spatial factors. DMFA proves fast and capable on large-scale fMRI data.

heading: Acknowledgment

text: This work was supported by award numbers 1944964 and 1835309 from the national science foundation (NSF). The authors would also like to thank the anonymous reviewers for their insightful feedback.

title: Deep Markov Factor Analysis: Towards Concurrent Temporal and Spatial Analysis of fMRI Data


summaries: The paper tackles the problem of learning flow-based generative models. The proposed approach, a new instance of continuous normalizing-flow methods, models the density as a difference between the prior and a divergence term, parameterized by a neural network.  The claims are that: * MF constitutes a universal density approximator, * MF is more efficient to train than alternative CNFs; this is demonstrated empirically and the intuition is that MF does not require invoking or backpropagating through an ODE during training (note however that solving an ODE is required for sampling), * MF improves upon alternative CNFs in terms of density estimation, sample quality, and training complexity, * MF demonstrates for the first time the use of flow models for sampling from general curved surfaces.

summaries: This paper proposes an alternative way to train a continuous normalizing flow (CNF), inspired by the Dacorogna-Moser transport. The proposed training does not require numerical integration; instead, the likelihood of the flow can be directly evaluated by computing the divergence of the mass flow rate. The idea is tested on boundaryless manifolds, such as the flat torus for toy data and implicit surfaces.

summaries: The paper introduces a new type of continuous normalizing flow (CNF) for manifold-valued data. Unlike in previous works, this Moser Flow (MF) models a density as a base density minus the divergence of a neural network. The authors prove that under certain assumptions this defines a probability distribution over the manifold. This approach has a different trade-off between inference and generation from other approaches: while the model density can be evaluated very cheaply, sampling from this model requires an ODE solver. The approach is qualitatively demonstrated on toy datasets on tori and the Stanford bunny, and quantitatively evaluated on a few datasets on 2-spheres.

summaries: This paper presents Moser Flow a new class of Continuous Normalizing Flows for Riemannian Manifolds. The main approach exploits results in differential geometry (Moser 1965) to construct a vector field that satisfies the normalizing equation given by the pullback of the volume form under $\Phi$. The approach is both sensible and provides direct advantages over Riemmanian CNF's as training does not require solving the ODE using a solver. The authors also consider Moser Flows for Euclidean Submanifolds under the Induced metric and show that Moser Flows can learn arbitrary distributions on manifolds of interest (---i.e. compact, boundryless etc...). Finally, the experiments conducted are both interesting in their variety and illuminating in that they show the effectiveness of the proposed method. Overall, this paper represents a non-trivial step forward---manifold or otherwise---for Continuous Normalizing Flow research.

summaries: The authors define a continuous normalizing flow on an arbitrary Riemannian manifold through the divergence of a neural network using the Moser theorem. They motivate their approach and show how this flow can be represented and learned. By applying their method to a toy distributions and a real world earth and climate dataset, they demonstrate the effectiveness of the method and that they can outperform competing procedures. Furthermore, they show that their method is much faster and more efficient than FFJORD when learning a complex two dimensional density.

abstractText: We are interested in learning generative models for complex geometries described via manifolds, such as spheres, tori, and other implicit surfaces. Current extensions of existing (Euclidean) generative models are restricted to specific geometries and typically suffer from high computational costs. We introduce Moser Flow (MF), a new class of generative models within the family of continuous normalizing flows (CNF). MF also produces a CNF via a solution to the change-of-variable formula, however differently from other CNF methods, its model (learned) density is parameterized as the source (prior) density minus the divergence of a neural network (NN). The divergence is a local, linear differential operator, easy to approximate and calculate on manifolds. Therefore, unlike other CNFs, MF does not require invoking or backpropagating through an ODE solver during training. Furthermore, representing the model density explicitly as the divergence of a NN rather than as a solution of an ODE facilitates learning high fidelity densities. Theoretically, we prove that MF constitutes a universal density approximator under suitable assumptions. Empirically, we demonstrate for the first time the use of flow models for sampling from general curved surfaces and achieve significant improvements in density estimation, sample quality, and training complexity over existing CNFs on challenging synthetic geometries and real-world benchmarks from the earth and climate sciences.

authors:
A.H. Al-Mohy
N.J. Higham

title: A New Scaling and Squaring Algorithm for the Matrix Exponential

venue: SIAM Journal on Matrix Analysis and Applications,

authors:
A.J. Bose
A. Smofsky
R. Liao
P. Panangaden
W.L. Hamilton

title: Latent Variable Modelling with Hyperbolic Normalizing Flows

authors:
T.B. Brown
B. Mann
N. Ryder
M. Subbiah
J. Kaplan
P. Dhariwal
A. Neelakantan
P. Shyam
G. Sastry
A Askell

title: Language models are few-shot learners. arXiv preprint arXiv:2005.14165

authors:
M. Chen
B. Tu
B. Lu

title: Triangulated manifold meshing method preserving molecular surface topology

venue: Journal of Molecular Graphics and Modelling,

authors:
R.T. Chen
Y. Rubanova
J. Bettencourt
D. Duvenaud

title: Neural ordinary differential equations. arXiv preprint arXiv:1806.07366

authors:
R.T.Q. Chen
J. Behrmann
D. Duvenaud
Jacobsen
J.-H

title: Residual flows for invertible generative modeling

authors:
B. Dacorogna
J. Moser

title: On a partial differential equation involving the jacobian determinant

venue: In Annales de l’Institut Henri Poincare (C) Non Linear Analysis,

authors:
P. Dhariwal
A. Nichol

title: Diffusion models beat gans on image synthesis

venue: arXiv preprint arXiv:2105.05233

authors:
L. Dinh
J. Sohl-Dickstein
S. Bengio

title: Density estimation using real nvp

venue: arXiv preprint arXiv:1605.08803

authors:
A. Dooley
N. Wildberger

title: Harmonic analysis and the global exponential map for compact Lie groups

venue: Functional Analysis and Its Applications,

authors:
L. Falorsi
P. Forré

title: Neural Ordinary Differential Equations on Manifolds. arXiv:2006.06663 [cs, stat

authors:
S. Gerber
T. Tasdizen
P.T. Fletcher
S. Joshi
R. Whitaker
Initiative
A.D. N

title: Manifold modeling for brain population analysis

authors:
W. Grathwohl
R.T.Q. Chen
J. Bettencourt
I. Sutskever
D. Duvenaud

title: Ffjord: Free-form continuous dynamics for scalable reversible generative models

authors:
A. Gropp
L. Yariv
N. Haim
M. Atzmon
Y. Lipman

title: Implicit geometric regularization for learning shapes

authors:
H. Hoppe
T. DeRose
T. Duchamp
J. McDonald
W. Stuetzle

title: Surface reconstruction from unorganized points

venue: In Proceedings of the 19th annual conference on computer graphics and interactive techniques,

authors:
Huang
C.-W
R.T.Q. Chen
C. Tsirigotis
A. Courville

title: Convex potential flows: Universal probability distributions with optimal transport and convex optimization

authors:
M. Kazhdan
M. Bolitho
H. Hoppe

title: Poisson surface reconstruction

venue: In Proceedings of the fourth Eurographics symposium on Geometry processing,

authors:
D.P. Kingma
J. Ba

title: Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980

authors:
M. Kumar
M. Babaeizadeh
D. Erhan
C. Finn
S. Levine
L. Dinh
D. Kingma

title: Videoflow: A flow-based generative model for video

venue: arXiv preprint arXiv:1903.01434,

title: Smooth manifolds. In Introduction to Smooth Manifolds, pages

authors:
W.E. Lorensen
H.E. Cline

title: Marching cubes: A high resolution 3d surface construction algorithm

venue: ACM siggraph computer graphics,

authors:
A. Lou
D. Lim
I. Katsman
L. Huang
Q. Jiang
Lim
S.-N
C. De Sa

title: Neural manifold ordinary differential equations

authors:
E. Mathieu
M. Nickel

title: Riemannian continuous normalizing flows. arXiv preprint arXiv:2006.10605

title: On the volume elements on a manifold

venue: Transactions of the American Mathematical Society,

authors:
G. Papamakarios
E. Nalisnick
D.J. Rezende
S. Mohamed
B. Lakshminarayanan

title: Normalizing flows for probabilistic modeling and inference. arXiv preprint arXiv:1912.02762

authors:
D. Rezende
S. Mohamed

title: Variational inference with normalizing flows

venue: In International Conference on Machine Learning,

authors:
D.J. Rezende
G. Papamakarios
S. Racaniere
M. Albergo
G. Kanwar
P. Shanahan
K. Cranmer

title: Normalizing flows on tori and spheres

title: A Jacobian Inequality for Gradient Maps on the Sphere and Its Application to Directional Statistics. Communications in Statistics - Theory and Methods, 42(14):2525–2542

authors:
M. Tancik
P.P. Srinivasan
B. Mildenhall
S. Fridovich-Keil
N. Raghavan
U. Singhal
R. Ramamoorthi
J.T. Barron
R. Ng

title: Fourier features let networks learn high frequency functions in low dimensional domains

heading: 1 Introduction

text:
The major successes of deep generative models in recent years are primarily in domains involving Euclidean data, such as images (Dhariwal and Nichol, 2021), text (Brown et al., 2020), and video (Kumar et al., 2019). However, many kinds of scientific data in the real world lie in non-Euclidean spaces specified as manifolds. Examples include planetary-scale data for earth and climate sciences (Mathieu and Nickel, 2020), protein interactions and brain imaging data for life sciences (Gerber et al., 2010; Chen et al., 2012), as well as 3D shapes in computer graphics (Hoppe et al., 1992; Kazhdan et al., 2006). Existing (Euclidean) generative models cannot be effectively applied in these scenarios as they would tend to assign some probability mass to areas outside the natural geometry of these domains.
An effective way to impose geometric domain constraints for deep generative modeling is to design normalizing flows that operate in the desired manifold space. A normalizing flow maps a prior (source) distribution to a target distribution via the change of variables formula (Rezende and Mohamed, 2015; Dinh et al., 2016; Papamakarios et al., 2019). Early work in this direction proposed invertible architectures for learning probability distributions directly over the specific manifolds defined over spheres and tori (Rezende et al., 2020). Recently, Mathieu and Nickel (2020) proposed to extend continuous normalizing flows (CNF) (Chen et al., 2018) for generative modeling over Riemannian manifolds wherein the flows are defined via vector fields on manifolds and computed as the solution to an associated ordinary differential equation (ODE). CNFs have the advantage that the neural
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
network architectures parameterizing the flow need not be restricted via invertibility constraints. However, as we show in our experiments, existing CNFs such as FFJORD (Grathwohl et al., 2018) and Riemannian CNFs (Mathieu and Nickel, 2020) can be slow to converge and the generated samples can be inferior in capturing the details of high fidelity data densities. Moreover, it is a real challenge to apply Riemannian CNFs to complex geometries such as general curved surfaces.
To address these challenges, we propose Moser Flows (MF), a new class of deep generative models within the CNF family. An MF models the desired target density as the source density minus the divergence of an (unrestricted) neural network. The divergence is a local, linear differential operator, easy to approximate and calculate on manifolds. By drawing on classic results in differential geometry by Moser (1965) and Dacorogna and Moser (1990), we can show that this parameterization induces a CNF solution to the change-of-variables formula specified via an ODE. Since MFs directly parameterize the model density using the divergence, unlike other CNF methods, we do not require to explicitly solve the ODE for maximum likelihood training. At test-time, we use the ODE solver for generation. We derive extensions to MFs for Euclidean submanifolds that efficiently parameterize vector fields projected to the desired manifold domain. Theoretically, we prove that Moser Flows are universal generative models over Euclidean submanifolds. That is, given a Euclidean submanifold M and a target continuous positive probability density µ overM, MFs can push arbitrary positive source density ν overM to densities µ̄ that are arbitrarily close to µ. We evaluate Moser Flows on a wide range of challenging real and synthetic problems defined over many different domains. On synthetic problems, we demonstrate improvements in convergence speed for attaining a desired level of details in generation quality. We then experiment with two kinds of complex geometries. First, we show significant improvements of 49% on average over Riemannian CNFs (Mathieu and Nickel, 2020) for density estimation as well as high-fidelity generation on 4 earth and climate science datasets corresponding to global locations of volcano eruptions, earthquakes, floods, and wildfires on spherical geometries. Next and last, we go beyond spherical geometries to demonstrate for the first time, generative models on general curved surfaces.

heading: 2 Preliminaries

text:
Riemannian manifolds. We consider an orientable, compact, boundaryless, connected ndimensional Riemannian manifoldM with metric g. We denote points in the manifold by x, y ∈M. At every point x ∈ M, TxM is an n-dimensional tangent plane toM. The metric g prescribes an inner product on each tangent space; for v, u ∈ TxM, their inner product w.r.t. g is denoted by 〈v, u〉g. X(M) is the space of smooth (tangent) vector fields toM; that is, if u ∈ X(M) then u(x) ∈ TxM, for all x ∈ M, and if u written in local coordinates it consists of smooth functions. We denote by dV the Riemannian volume form, defined by the metric g over the manifoldM. In particular, V (A) = ∫ A dV is the volume of the set A ⊂M.
We consider probability measures overM that are represented by strictly positive continuous density functions µ, ν :M→ R>0, where µ by convention represents the target (unknown) distribution and ν represents the source (prior) distribution. µ, ν are probability densities in the sense their integral w.r.t. the Riemannian volume form is one, i.e., ∫ M µdV = 1 = ∫ M νdV . It is convenient to consider the volume forms that correspond to the probability measures, namely µ̂ = µdV and ν̂ = νdV . Volume forms are differential n-forms that can be integrated over subdomains ofM, for example, pν(A) = ∫ A ν̂ is the probability of the event A ⊂M.
Continuous Normalizing Flows (CNF) on manifolds operate by transforming a simple source distribution through a map Φ into a highly complex and multimodal target distribution. A manifold CNF, Φ :M→M, is an orientation preserving diffeomorphism from the manifold to itself (Mathieu and Nickel, 2020; Lou et al., 2020; Falorsi and Forré, 2020). A smooth map Φ :M→M can be used to pull-back the target µ̂ according to the formula:
(Φ∗µ̂)z(v1, . . . , vn) = µ̂Φ(z)(DΦz(v1), . . . , DΦz(vn)), (1)
where v1, . . . , vn ∈ TzM are arbitrary tangent vectors, DΦz : TzM→ TΦ(z)M is the differential of Φ, namely a linear map between the tangent spaces toM at the points z and Φ(z), respectively. By pulling-back µ̂ according to Φ and asking it to equal to the prior density ν, we get the manifold version of the standard normalizing equation:
ν̂ = Φ∗µ̂. (2)
If the normalizing equation holds, then for an event A ⊂M we have that pν(A) = ∫ A ν̂ = ∫ A Φ∗µ̂ = ∫ Φ(A) µ̂ = pµ(Φ(A)).
Therefore, given a random variable z distributed according to ν, then x = Φ(z) is distributed according to µ, and Φ is the generator.
One way to construct a CNF Φ is by solving an ordinary differential equation (ODE) (Chen et al., 2018; Mathieu and Nickel, 2020). Given a time-dependent vector field vt ∈ X(M) with t ∈ [0, 1], a one-parameter family of diffeomorphisms (CNFs) Φt : [0, 1]×M→M is defined by
d dt Φt = vt(Φt), (3)
where this ODE is initialized with the identity transformation, i.e., for all x ∈ M we initialize Φ0(x) = x. The CNF is then defined by Φ(x) = Φ1(x).
Example: Euclidean CNF. Let us show how the above notions boil down to standard Euclidean CNF for the choice ofM = Rn, and the standard Euclidean metric; we denote z = (z1, . . . , zn) ∈ Rn. The Riemannian volume form in this case is dz = dz1∧dz2∧· · ·∧dzn. Furthermore, µ̂(z) = µ(z)dz and ν̂(z) = ν(z)dz. The pull-back formula (equation 1) in coordinates (see e.g., Proposition 14.20 in Lee (2013)) is
Φ∗µ̂(z) = µ(Φ(z)) det(DΦz)dz,
where DΦz is the matrix of partials of Φ at point z, (DΦz)ij = ∂Φ i
∂zj (z). Plugging this in equation 2 we get the Euclidean normalizing equation:
ν(z) = µ(Φ(z)) det(DΦz). (4)

heading: 3 Moser Flow

text: Moser (1965) and Dacorogna and Moser (1990) suggested a method for solving the normalizing equation, that is equation 2. Their method explicitly constructs a vector field vt, and the flow it defines via equation 3 is guaranteed to solve equation 2. We start by introducing the method, adapted to our needs, followed by its application to generative modeling. We will use notations introduced above.

heading: 3.1 Solving the normalizing equation

text:
Moser’s approach to solving equation 2 starts by interpolating the source and target distributions. That is, choosing an interpolant αt : [0, 1]×M→ R>0, such that α0 = ν, α1 = µ, and ∫ M αtdV = 1 for all t ∈ [0, 1]. Then, a time-dependent vector field vt ∈ X(M) is defined so that for each time t ∈ [0, 1], the flow Φt defined by equation 3 satisfies the continuous normalization equation:
Φ∗t α̂t = α̂0, (5)
where α̂t = αtdV is the volume form corresponding to the density αt. Clearly, plugging t = 1 in the above equation provides a solution to equation 2 with Φ = Φ1. As it turns out, considering the continuous normalization equation simplifies matters and the sought after vector field vt is constructed as follows. First, solve the partial differential equation (PDE) over the manifoldM
div(ut) = − d
dt αt, (6)
where ut ∈ X(M) is an unknown time-dependent vector field, and div is the Riemannian generalization to the Euclidean divergence operator, divE = ∇·. This manifold divergence operator is defined by replacing the directional derivative of the Euclidean space with its Riemannian version, namely the covariant derivative,
div(u) = n∑ i=1 〈∇eiu, ei〉g , (7)
where {ei}ni=1 is an orthonormal frame according to the Riemannian metric g, and ∇ξu is the Riemannian covariant derivative. Note that here we assume thatM is boundaryless, otherwise we need ut to be also tangent to the boundary ofM. Second, define
vt = ut αt . (8)
Theorem 2 in Moser (1965) implies:
Theorem 1 (Moser). The diffeomorphism Φ = Φ1, defined by the ODE in equation 3 and vector field vt in equation 8 solves the normalization equation, i.e., equation 2.
The proof of this theorem in our case is provided in the supplementary for completeness. A simple choice for the interpolant αt that we use in this paper was suggested in Dacorogna and Moser (1990):
αt = (1− t)ν + tµ. (9)
The time derivative of this interpolant, i.e., ddtαt = µ − ν, does not depend on t. Therefore the vector field can be chosen to be constant over time, ut = u, and the PDE in equation 6 takes the form div(u) = ν − µ, (10) and consequently vt takes the form
vt = u
(1− t)ν + tµ . (11)
Figure 1 shows a one dimensional illustration of Moser Flow.

heading: 3.2 Generative model utilizing Moser Flow

text:
We next utilize MF to define our generative model. Our model (learned) density µ̄ is motivated from equation 10 and is defined by µ̄ = ν − div(u), (12) where u ∈ X(M) is the degree of freedom of the model. We model this degree of freedom, u, with a deep neural network, more specifically, Multi-Layer Perceptron (MLP). We denote θ ∈ Rp the learnable parameters of u. We start by noting that, by construction, µ̄ has a unit integral overM: Lemma 1. IfM has no boundary, or u|∂M ∈ X(∂M), then ∫ M µ̄dV = 1.
This lemma is proved in the supplementary and a direct consequence of Stokes’ Theorem. If µ̄ > 0 overM then, together with the fact that ∫ M µ̄dV = 1 (Lemma 1), it is a probability density overM. Consequently, Theorem 1 implies that µ̄ is realized by a CNF defined via vt:
Corollary 1. If µ̄ > 0 overM then µ̄ is a probability distribution overM, and is generated by the flow Φ = Φ1, where Φt is the solution to the ODE in equation 3 with the vector field vt ∈ X(M) defined in equation 11.
Since µ̄ > 0 is an open constraint and is not directly implementable, we replace it with the closed constraint µ̄ ≥ , where > 0 is a small hyper-parameter. We define
µ̄+(x) = max { , µ̄(x)} ; µ̄−(x) = −min { , µ̄(x)} .
As can be readily verified: µ̄+, µ̄− ≥ 0, and µ̄ = µ̄+ − µ̄−. (13)
We are ready to formulate the loss for training the generative model. Consider an unknown target distribution µ, provided to us as a set of i.i.d. observations X = {xi}mi=1 ⊂ M. Our goal is to maximize the likelihood of the data X while making sure µ̄ ≥ . We therefore consider the following loss:
`(θ) = −Eµ log µ̄+(x) + λ ∫ M µ̄−(x)dV (14)
where λ is a hyper-parameter. The first term in the loss is approximated by the empirical mean computed with the observations X , i.e.,
Ex∼µ log µ̄+(x) ≈ 1
m m∑ i=1 log µ̄+(xi).
This term is merely the negative log likelihood of the observations.
The second term in the loss penalizes the deviation of µ̄ from satisfying µ̄ ≥ . According to Theorem 1, this measures the deviation of µ̄ from being a density function and realizing a CNF. One point that needs to be verified is that the combination of these two terms does not push the minimum away from the target density µ. This can be verified with the help of the generalized Kullback–Leibler (KL) divergence providing a distance measure between arbitrary positive functions f, g :M→ R>0:
D(f, g) = ∫ M f log ( f g ) dV − ∫ M fdV + ∫ M gdV. (15)
Using the generalized KL, we can now compute the distance between the positive part of our model density, i.e., µ̄+, and the target density:
D(µ, µ̄+) = Eµ log ( µ
µ̄+ ) − ∫ M µdV + ∫ M µ̄+dV
= Eµ logµ− Eµ log µ̄+ + ∫ M µ̄−dV
where in the second equality we used Lemma 1. The term Eµ logµ is the negative entropy of the unknown target distribution µ. The loss in equation 14 equals D(µ, µ̄+) − Eµ logµ + (λ − 1) ∫ M µ̄−dV . Therefore, if λ ≥ 1, and minx∈M µ(x) > (we use the compactness ofM to infer existence of such a minimal positive value), then the unique minimum of the loss in equation 14 is the target density, i.e., µ̄ = µ. Indeed, the minimal value of this loss is −Eµ logµ and it is achieved by setting µ̄ = µ. Uniqueness follows by considering an arbitrary minimizer µ̄. Since such a minimizer satisfies D(µ, µ̄+) = 0 and ∫ M µ̄−dV = 0, necessarily µ̄ = µ. We proved: Theorem 2. For λ ≥ 1 and sufficiently small > 0, the unique minimizer of the loss in equation 14 is µ̄ = µ. Variation of the loss. Lemma 1 and equation 13 imply that ∫ M µ̄+dV = ∫ M µ̄−dV +1. Therefore, an equivalent loss to the one presented in equation 14 is:
`(θ) = −Eµ log µ̄+(x) + λ− ∫ M µ̄−dV + λ+ ∫ M µ̄+dV (16)
with λ− + λ+ ≥ 1. Empirically we found this loss favorable in some cases (i.e., with λ+ > 0).
Integral approximation. The integral ∫ M µ̄−dV in the losses in equation 16 and 14 is approximated by considering a set Y = {yj}lj=1 of i.i.d. samples according to some distribution η overM and taking a Monte Carlo estimate ∫
M µ̄−dV ≈
1
l l∑ j=1 µ̄−(yj) η(yj) . (17)
∫ M µ̄+dV is approximated similarly. In this paper we opted for the simple choice of taking η to be the uniform distribution overM.

heading: 4 Generative modeling over Euclidean submanifolds

text:
In this section, we adapt the Moser Flow (MF) generative model to submanifolds of Euclidean spaces. That is we consider an orientable, compact, boundaryless, connected n-dimensional submanifold M⊂ Rd, where n < d. Examples include implicit surfaces and manifolds (i.e., preimage of a regular value of a smooth function), as well as triangulated surfaces and manifold simplicial complexes. We
denote points in Rd (and therefore inM) with x,y ∈ Rd. As the Riemannian metric ofM we take the induced metric from Rd; that is given arbitrary tangent vectors v,u ∈ TxM the metric is defined by 〈v,u〉g = 〈v,u〉, where the latter is the Euclidean inner product. We denote by π : Rd →M the closest point projection onM, i.e., π(x) = miny∈M ‖x− y‖, with ‖y‖2 = 〈y,y〉 the Euclidean norm in Rd. Lastly, we denote by Px ∈ Rd×d the orthogonal projection matrix on the tangent space TxM; in practice if we denote by N ∈ Rd×k the matrix with orthonormal columns spanning NxM = (TxM)⊥ (i.e., the normal space toM at x) then, Px = I −NNT . We parametrize the vector field u required for our MF model (in equation 12) by defining a vector field u ∈ X(Rd) such that u|M ∈ X(M). We define
u(x) = Pπ(x)vθ(π(x)), (18)
where vθ : Rd → Rd is an MLP with Softplus activation (β = 100) and learnable parameters θ ∈ Rp. By construction, for x ∈M, u(x) ∈ TxM. To realize the generative model, we need to compute the divergence div(u(x)) for x ∈ M with respect to the Riemannian manifoldM and metric g. The vector field u in equation 18 is constant along normal directions to the manifold at x (since π(x) is constant in normal directions). If n ∈ NxM, then in particular
d
dt ∣∣∣ t=0 u(x + tn) = 0. (19)
We call vector fields that satisfy equation 19 infinitesimally constant in the normal direction. As we show next, such vector fields u ∈ X(M) have the useful property that their divergence along the manifoldM coincides with their Euclidean divergence in the ambient space Rd: Lemma 2. If u ∈ X(Rd), u|M ∈ X(M) is infinitesimally constant in normal directions ofM, then for x ∈M, div(u(x)) = divE(u(x)), where divE denotes the standard Euclidean divergence.
This lemma simply means we can compute the Euclidean divergence of u in our implementation. Given a set of observed data X = {xi}mi=1 ⊂ M ⊂ Rd, and a set of uniform i.i.d. samples Y = {yj}lj=1 ⊂M overM, our loss (equation 14) takes the form
`(θ) = − 1 m m∑ i=1 log max { , ν(xi)− divEu(xi)}+ λ′− l l∑ j=1 ( −min { , ν(yj)− divEu(yj)} ) ,
where λ′− = λ−V (M). We note the volume constant can be ignored by considering an un-normalized source density V (M)ν ≡ 1, see supplementary for details. The loss in equation 16 is implemented similarly, namely, we add the empirical approximation of λ+ ∫ M µ̄+dV .
We conclude this section by stating that the MF generative model over Euclidean submanifolds (defined with equations 12 and 18) is universal. That is, MFs can generate, arbitrarily well, any continuous target density µ on a submanifold manifoldM⊂ Rd. Theorem 3. Given an orientable, compact, boundaryless, connected, differentiable n-dimensional submanifoldM ⊂ Rd, n < d, and a target continuous probability density µ : M → R>0, there exists for each > 0 an MLP vθ and a choice of weights θ so that µ̄ defined by equations 12 and 18 satisfies
max x∈M
|µ(x)− µ̄(x)| < .

heading: 5 Experiments

text:
In all experiments, we modeled a manifold vector field as a multi-layer perceptron (MLP) uθ ∈ X(M), with parameters θ. All models were trained using Adam optimizer (Kingma and Ba, 2014), and in all neural networks the activation is Softplus with β = 100. Unless stated otherwise, we set λ+ = 0. We used an exact calculation of the divergence divE(u(x)). We experimented with two kinds of manifolds.
Flat Torus. To test our method on Euclidean 2D data, we usedM as the flat torus, that is the unit square [−1, 1]2 with opposite edges identified. This defines a manifold with no boundary which is locally isometric to the Euclidean plane. Due to this local isometry the Riemannian divergence on the flat
torus is equivalent to the Euclidean divergence, div = divE . To make uθ a well defined smooth vector field inM we use periodic positional encoding, namely uθ(x) = vθ(τ(x)), where vθ is a standard MLP and τ : R2 → R4k is defined as τ(x) = (cos(ω1πx), sin(ω1πx), ..., cos(ωkπx), sin(ωkπx)), where wi = i, and k is a hyper-parameter that is application dependent. Since any periodic function can be approximated by a polynomial acting on eiπx, even for k = 1 this is a universal model for continuous functions on the torus. As described by Tancik et al. (2020), adding extra features can help with learning higher frequencies in the data. To solve an ODE on the torus we simply solve it for the periodic function and wrap the result back to [−1, 1]2. Implicit surfaces. We experiment with surfaces as submanifolds of R3. We represent a surface as the zero level set of a Signed Distance Function (SDF) f : R3 → R. We experimented with two surfaces. First, the sphere, represented with the SDF f(x) = ‖x‖ − 1, and second, the Stanford Bunny surface, representing a general curved surface and represented with an SDF learned with (Gropp et al., 2020) from point cloud data. To model vector fields on an implicit surface we follow the general equation 18, where for SDFs
π(x) = x− f(x)∇f(x), and Px = I −∇f(x)∇f(x)T .
In the supplementary, we detail how to replace the global projection π(x) with a local one, for cases the SDF is not exact.

heading: 5.1 Toy distributions

text: First, we consider a collection of challenging toy 2D datasets explored in prior works (Chen et al., 2020; Huang et al., 2021). We scale samples to lie in the flat torus [−1, 1]2 and use k = 1 for the positional encoding. Figure 2 depicts the input data samples, the generated samples after training, and the learned distribution µ̄. In the top six datasets, the MLP architecture used for Moser Flows consists of 3 hidden layers with 256 units each, whereas in the bottom two we used 4 hidden layers with 256 neurons due to the higher complexity of these distributions. We set λ− = 2.

heading: 5.2 Choice of hyper-parameter λ

text:
We test the effect of the hyper-parameter λ ≥ 1 on the learned density. Figure 3 shows, for different values of λ, our density estimation µ+, the push-forward density Φ∗ν, and their absolute difference. To evaluate Φ∗ν from the vector field vt, we solve an ODE as advised in Grathwohl et al. (2018).
As expected, higher values of λ lead to closer modeled density µ̄+ and Φ∗ν. This is due to the fact that a higher value of λ leads to a lower value of ∫ M µ̄−, meaning µ̄ is a better representation of a probability density. Nonetheless, even for λ = 1 the learned and generated density are rather consistent.

heading: 5.3 Time evaluations

text: To compare our method to Euclidean CNF methods, we compare Moser Flow with FFJORD (Grathwohl et al., 2018) on the flat torus. We consider a challenging density with high frequencies obtained via 800x600 images (Figure 4). We generate a new batch every iteration by sampling each pixel location with probability which is proportional to the pixel intensity. The architectures of both vθ and the vector field defined in FFJORD are the same, namely an MLP with 4 hidden layers of 256 neurons each. To capture the higher frequencies in the image we use a positional encoding with k = 8 for both methods. We used a batch size of 10k. We used learning rate of 1e-5 for Moser Flow and 1e-4 for FFJORD. We used λ− = 2. Learning was stopped after 5k seconds. Figure 4 presents the results. Note that Moser Flow captures high-frequency details better than FFJORD. This is also expressed in the graph on the top right, showing how the NLL decreases faster for MF than FFJORD. Furthermore, as can be inspected in the per iteration time graph on the bottom-right, MF per iteration time does not increase during training, and is roughly 1-2 order of magnitudes faster than FFJORD iteration.

heading: 5.4 Earth and climate science data

text: We evaluate our model on the earth and climate datasets gathered in Mathieu and Nickel (2020). The projection operator π in this case is simply π(x) = x‖x‖ . We parameterize vθ as an MLP with 6 hidden layers of 512 neurons each. We used full batches for the NLL loss and batches of size 150k for our integral approximation. We trained for 30k epochs, with learning rate of 1e-4. We used λ− = 100. The quantitative NLL results are reported in Table 1 and qualitative visualizations in 5. Note that we produce NLL scores smaller than the runner-up method by a large margin.

heading: 5.5 Curved surfaces

text: We trained an SDF f for the Stanford Bunny surfaceM using the method in Gropp et al. (2020). To generate uniform (ν) and data (µ) samples overM we first extract a meshM′ from f using the Marching Cubes algorithm (Lorensen and Cline, 1987) setting its resolution to 100x100x100. Then, to randomly choose a point uniformly fromM′ we first randomly choose a face of the mesh with probability proportional to its area, and then randomly choose a point uniformly within that face. For target µ we used clamped manifold harmonics to create a sequence of densities with increased complexity. To that end, we first computed the k-th eigenfunction of the Laplace-Beltrami operator overM′ (we provide details on this computation in the supplementary), for the frequencies (eigenvalues) k ∈ {10, 50, 500}. Next, we sampled the eigenfunctions at the faces’ centers, clamped their negative values, and normalized to get discrete probability densities over the faces ofM′. Then, to sample a point, we first choose a face at random based on this probability, and then random a point uniformly within that face. We take 500k i.i.d. samples of this distribution as our dataset. We take vθ to be an MLP with 6 hidden layers of dimension 512. We use batch size of 10k for both the NLL loss and for the integral approximation; we ran for 1000 epochs with learning rate of 1e-4. We used λ− = λ+ = 1. Figure 6 depict the results. Note that Moser Flow is able to learn the surface densities for all three frequencies.

heading: 6 Related Work

text:
In the following, we discuss related work on normalizing flows for manifold-valued data. On a high level, such methods can be divided into Parametric versus Riemannian methods. Parametric methods consist of a normalizing flow in the Euclidean space Rn, pushed-forward onto the manifold through an invertible map ψ : Rn → M. However, to globally represent the manifold, ψ needs to be a homeomorphism implying thatM and Rn are topologically equivalent, limiting the scope of that
approach. Existing methods in this class are often based on the exponential map expx : TxM ∼= Rn → M of a manifold. This leads to the so called wrapped distributions. This approach has been taken, for instance, by Falorsi et al. (2019) and Bose et al. (2020) to parametrize probability distributions on Lie groups and hyperbolic space. However, Parametric methods based on the exponential map often lead to numerical and computational challenges. For instance, in compact manifolds (e.g., spheres or the SO(3) group) computing the density of wrapped distributions requires an infinite summation. On the hyperboloid, on the other hand, the exponential map is numerically not well-behaved far away from the origin (Dooley and Wildberger, 1993; Al-Mohy and Higham, 2010).
In contrast to Parametric methods, Riemannian methods operate directly on the manifold itself and, as such, avoid numerical instabilities that arise from the mapping onto the manifold. Early work in this class of models proposed transformations along geodesics on the hypersphere by evaluating the exponential map at the gradient of a scalar manifold function (Sei, 2011). Rezende et al. (2020) introduced discrete Riemannian flows for hyperspheres and torii based on Möbius transformations and spherical splines. Mathieu and Nickel (2020) introduced continuous flows on general Riemannian manifolds (RCNF). In contrast to discrete flows (e.g., Bose et al., 2020; Rezende et al., 2020), such time-continuous flows alleviate the previous topological constraints by parametrizing the flow as the solution to an ODE over the manifold (Grathwohl et al., 2018). Concurrently to RCNF, Lou et al. (2020) and Falorsi and Forré (2020) proposed related extensions of neural ODEs to smooth manifolds. Moser Flow also generates a CNF, however by limiting the flow space (albeit, not the generated distributions) it allows expressing the learned distribution as the divergence of a vector field.

heading: 7 Discussion and limitations

text:
We introduced Moser Flow, a generative model in the family of CNFs that represents the target density using the divergence operator applied to a vector valued neural network. The main benefits of MF stems from the simplicity and locality of the divergence operator. MFs circumvent the need to solve an ODE in the training process, and are thus applicable on a broad class of manifolds. Theoretically, we prove MF is a universal generative model, able to (approximately) generate arbitrary positive target densities from arbitrary positive prior densities. Empirically, we show MF enjoys favorable computational speed in comparison to previous CNF models, improves density estimation on spherical data compared to previous work by a large margin, and for the first time facilitate training a CNF over a general curved surface.
One important future work direction, and a current limitation, is scaling of MF to higher dimensions. This challenge can be roughly broken to three parts: First, the model probabilities µ̄ should be computed/approximated in log-scale, as probabilities are expected to decrease exponentially with the dimension. Second, the variance of the numerical approximations of the integral ∫ M µ̄−dV will increase significantly in high dimensions and needs to be controlled. Third, the divergence term, div(u), is too costly to be computed exactly in high dimensions and needs to be approximated, similarly to other CNF approaches. Finally, our work suggests a novel generative model, and similarly to other generative models can be potentially used for generation of fake data and amplify harmful biases in the dataset. Mitigating such harms is an active and important area of ongoing research.

heading: Acknowledgments

text: NR is supported by the European Research Council (ERC Consolidator Grant, "LiftMatch" 771136), the Israel Science Foundation (Grant No. 1830/17), and Carolito Stiftung (WAIC).

title: Moser Flow: Divergence-based Generative Modeling on Manifolds