summaries: The main focus/goal of the submitted paper is the parallelization of the Gated Recurrent Unit (GRU) (the authors focus on classification problems). The authors describe the incorporation of a multigrid reduction in time (MGRIT) solver to speed-up and better parallelize the  application of forward and back propagation of information. The proposed technique seems to provide a speedup of about an order of  magnitude (at most) when implemented on distributed/shared memory hybrid computing environments.

summaries: The paper proposed to parallelize the inference and training of GRU networks (a type of recurrent neural networks) at the `time` dimension.  The main contribution is the application of multigrid reduction in time (MGRIT) solver, and a new GRU architecture (Implicit GRU) that handles the stiffness in the architecture. As a result, the evaluation shows 6.5 times faster training time.

summaries: The paper describes a technique for evaluating GRU networks based on the multigrid reduction in time (MGRIT) technique. These techniques are not new, in general or to neural network training, but the contribution here is their application to GRU layers. After presenting some of the theory behind ordinary differential equation (ODE) representations of the GRU and laying out the mathematical framework for the MGRIT method, the method is evaluated on two datasets. The results show impressive scalability up to 32 processes (CPU-only) at the cost of a slight loss of accuracy compared to the traditional, sequential GRU

summaries: This paper aims to address the limitations of existing approaches for training Gated Recurrent Unit (GRU) given long sequence in terms of both training time and model accuracy. To tackle this challenge, author propose a novel parallel training scheme (called parallel-in-time) for GRU based on a multigrid reduction in time (MGRIT) solver.  Specifically, the key to achieving speedup is a hierarchical correction of the hidden state to accelerate end-to-end communication in both the forward and backward propagation. Authors gives experimental results on two public datasets to demonstrate the performance improvement in the long sequence scenario.

abstractText: Parallelizing Gated Recurrent Unit (GRU) networks is a challenging task, as the training procedure of GRU is inherently sequential. Prior efforts to parallelize GRU have largely focused on conventional parallelization strategies such as dataparallel and model-parallel training algorithms. However, when the given sequences are very long, existing approaches are still inevitably performance limited in terms of training time. In this paper, we present a novel parallel training scheme (called parallel-in-time) for GRU based on a multigrid reduction in time (MGRIT) solver. MGRIT partitions a sequence into multiple shorter sub-sequences and trains the sub-sequences on different processors in parallel. The key to achieving speedup is a hierarchical correction of the hidden state to accelerate end-to-end communication in both the forward and backward propagation phases of gradient descent. Experimental results on the HMDB51 dataset, where each video is an image sequence, demonstrate that the new parallel training scheme achieves up to 6.5× speedup over a serial approach. As efficiency of our new parallelization strategy is associated with the sequence length, our parallel GRU algorithm achieves significant performance improvement as the sequence length increases.

authors:
Davide Anguita
Alessandro Ghio
Luca Oneto
Xavier Parra
Jorge Luis Reyes-Ortiz

title: A public domain dataset for human activity recognition using smartphones

venue: In Esann,

authors:
Tal Ben-Nun
Torsten Hoefler

title: Demystifying parallel and distributed deep learning: An in-depth concurrency analysis

venue: ACM Computing Surveys (CSUR),

title: Multi-level adaptive solutions to boundary-value problems

venue: Mathematics of computation,

authors:
William L Briggs
Van Emden Henson
Steve F McCormick

title: A multigrid tutorial

authors:
Joao Carreira
Andrew Zisserman

title: Quo vadis, action recognition? a new model and the kinetics dataset

venue: In proceedings of the IEEE Conference on Computer Vision and Pattern Recognition,

authors:
Bo Chang
Minmin Chen
Eldad Haber
Ed H Chi

title: Antisymmetricrnn: A dynamical system view on recurrent neural networks

authors:
Tian Qi Chen
Yulia Rubanova
Jesse Bettencourt
David Duvenaud

title: Neural ordinary differential equations

venue: arXiv preprint arXiv:1806.07366,

authors:
Kyunghyun Cho
Bart Van Merriënboer
Caglar Gulcehre
Dzmitry Bahdanau
Fethi Bougares
Holger Schwenk
Yoshua Bengio

title: Learning phrase representations using rnn encoder-decoder for statistical machine translation

venue: arXiv preprint arXiv:1406.1078,

authors:
Eric C Cyr
Stefanie Günther
Jacob B Schroder

title: Multilevel initialization for layer-parallel deep neural network training

venue: arXiv preprint arXiv:1912.08974,

authors:
Edward De Brouwer
J Simm
Ádám Arany
Y Moreau

title: Gru-ode-bayes: Continuous modeling of sporadically-observed time series

venue: Proc. of Advances in Neural Information Processing Systems (NeurIPS),

authors:
Jia Deng
Wei Dong
Richard Socher
Li-Jia Li
Kai Li
Li Fei-Fei

title: Imagenet: A large-scale hierarchical image database

authors:
V.A. Dobrev
Tz. Kolev
N.A. Petersson
J.B. Schroder

title: Two-level convergence theory for multigrid reduction in time (mgrit)

venue: SIAM Journal on Scientific Computing,

authors:
Jeffrey Donahue
Lisa Anne Hendricks
Sergio Guadarrama
Marcus Rohrbach
Subhashini Venugopalan
Kate Saenko
Trevor Darrell

title: Long-term recurrent convolutional networks for visual recognition and description

venue: In Proceedings of the IEEE conference on computer vision and pattern recognition,

title: Multigrid methods in science and engineering

venue: IEEE Computational Science and Engineering,

authors:
Robert D Falgout
Stephanie Friedhoff
Tz V Kolev
Scott P MacLachlan
Jacob B Schroder

title: Parallel time integration with multigrid

title: 50 years of time parallel time integration. In Multiple shooting and time domain decomposition methods

authors:
Rohit Girdhar
Joao Carreira
Carl Doersch
Andrew Zisserman

title: Video action transformer network

venue: In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition,

authors:
Stefanie Gunther
Lars Ruthotto
Jacob B Schroder
Eric C Cyr
Nicolas R Gauger

title: Layerparallel training of deep residual neural networks

venue: SIAM Journal on Mathematics of Data Science,

authors:
Mansura Habiba
Barak A Pearlmutter

title: Neural ordinary differential equation based recurrent neural network model

venue: In 2020 31st Irish Signals and Systems Conference (ISSC),

authors:
Kaiming He
Xiangyu Zhang
Shaoqing Ren
Jian Sun

title: Deep residual learning for image recognition

authors:
Zhiheng Huang
Geoffrey Zweig
Michael Levit
Benoit Dumoulin
Barlas Oguz
Shawn Chang

title: Accelerating recurrent neural network training via two stage classes and parallelization

venue: IEEE Workshop on Automatic Speech Recognition and Understanding,

authors:
Ian D Jordan
Piotr Aleksander Sokół
Il Memming Park

title: Gated recurrent units viewed through the lens of continuous time dynamical systems

venue: Frontiers in computational neuroscience,

authors:
Andrew Kirby
Siddharth Samsi
Michael Jones
Albert Reuther
Jeremy Kepner
Vijay Gadepally

title: Layer-parallel training with gpu concurrency of deep residual neural networks via nonlinear multigrid

venue: IEEE High Performance Extreme Computing Conference (HPEC). IEEE,

authors:
Hildegard Kuehne
Hueihan Jhuang
Estı́baliz Garrote
Tomaso Poggio
Thomas Serre

title: Hmdb: a large video database for human motion recognition

venue: In 2011 International conference on computer vision,

title: Training with multiple gpus using model parallelism

venue: Technical report,

authors:
Benjamin W Ong
Jacob B Schroder

title: Applications of time parallelization

venue: Computing and Visualization in Science,

title: Parallelizing over artificial neural network training runs with multigrid

venue: arXiv preprint arXiv:1708.02276,

authors:
Nitish Srivastava
Elman Mansimov
Ruslan Salakhudinov

title: Unsupervised learning of video representations using lstms

venue: In International conference on machine learning,

authors:
Xian-He Sun
Lionel M Ni

title: Scalable problems and memory-bounded speedup

venue: Journal of Parallel and Distributed Computing,

title: Training recurrent neural networks

venue: University of Toronto Toronto, Canada,

authors:
Xiaolong Wang
Ali Farhadi
Abhinav Gupta

title: Actions ̃ transformations

venue: In Proceedings of the IEEE conference on Computer Vision and Pattern Recognition,

authors:
Gerhard Wanner
Ernst Hairer

title: Solving ordinary differential equations II, volume 375

authors:
Yonghui Wu
Mike Schuster
Zhifeng Chen
Quoc V Le
Mohammad Norouzi
Wolfgang Macherey
Maxim Krikun
Yuan Cao
Qin Gao
Klaus Macherey

title: Google’s neural machine translation system: Bridging the gap between human and machine translation

venue: arXiv preprint arXiv:1609.08144,

authors:
Yang You
Jonathan Hseu
Chris Ying
James Demmel
Kurt Keutzer
Cho-Jui Hsieh

title: Largebatch training for lstm and beyond

venue: In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis,

authors:
Joe Yue-Hei Ng
Matthew Hausknecht
Sudheendra Vijayanarasimhan
Oriol Vinyals
Rajat Monga
George Toderici

title: Beyond short snippets: Deep networks for video classification

heading: 1 INTRODUCTION

text:
Recently, the model complexity of Deep Neural Networks (DNN) have grown to keep pace with the increasing scale of data. Hence, parallelization of DNNs is becoming important to achieve high quality models in tractable computational run times (Ben-Nun & Hoefler, 2019). Sequence modeling is one learning task that suffers from substantial computational time due to data size. The goal is to capture temporal dependencies within a sequence. Several different deep learning approaches addressing sequence tasks have been proposed, including 3D-CNNs (Carreira & Zisserman, 2017), LSTM (Srivastava et al., 2015), and Transformers (Girdhar et al., 2019). Focusing on sequence classification, this paper proposes a parallelization strategy for a recurrent neural network (RNN) architecture known as the gated recurrent unit (GRU) developed by Cho et al. (2014).
Parallel algorithms for training neural networks are generally broken into two classes; dataparallelism distributes a batch across processors, and model-parallelism distributes the architecture (Ben-Nun & Hoefler, 2019). Data-parallelism inherent in mini-batch stochastic gradient descent (SGD) is the primary strategy for parallelizing GRU (You et al., 2019; Huang et al., 2013). However, the accuracy of mini-batch SGD degrades as the number of sequences in each mini-batch increases. Moreover, although multiple mini-batches are processed in parallel, the training procedure of each sequence (within a mini-batch) is invariably sequential. Alternatively, the GRU model can be parallelized across hidden layers (Wu et al., 2016; MXNet, 2018). This model-parallelism approach is beneficial when the network requires a large number of layers. However, even largescale RNNs such as the Google Neural Machine Translation (GNMT) developed by Wu et al. (2016) have only 8 LSTM layers. Therefore, the scalability of layer-based model-parallelism is limited. A common problem in data-parallel and model-parallel algorithms for GRU is that the training execution time increases with the sequence length. At the heart of this difficulty is a fundamental need to process sequences in a serial fashion, yielding forward and backward propagation algorithms that are inherently sequential, and thus hard to parallelize.
We propose a parallel-in-time (PinT) training method for GRU networks with long sequences. To achieve this, we adapt a multigrid reduction in time (MGRIT) solver for forward and backward prop-
agation. Within the numerical methods community, a resurgence in PinT has paralleled increasing computational resources (Gander, 2015; Ong & Schroder, 2020). These efforts have even been applied to the training of neural Ordinary Differential Equations (ODEs) (Schroder, 2017; Gunther et al., 2020; Kirby et al., 2020; Cyr et al., 2019), where inexact forward and back propagation is exploited to achieve parallel speedups. To date, these techniques have not been applied to RNNs.
Following Jordan et al. (2021), GRU networks can be written as an ODE to facilitate application of MGRIT. Different from existing parallel training algorithms for RNNs, our MGRIT parallel training scheme partitions a sequence into multiple shorter sub-sequences and distributes each sub-sequence to different processors. By itself, this provides local improvement of the hidden states, yet global errors remain. To correct these errors, propagation is computed on a coarse representation of the input sequence requiring less computational work while still providing an improvement to the original hidden states. This process is iterated to the accuracy required for training. Applying this algorithm to the classic GRU networks will achieve parallelism. However, this training algorithm will not lead to accurate networks due to stability problems on the coarse grid. This emphasizes the challenges of choosing a proper coarse grid model for neural networks, and multigrid algorithms in general. To alleviate this, we develop a new GRU architecture, or discretization of the ODE, which we refer to as Implicit GRU that handles stiff modes in the ODE. This is required for application of the MGRIT algorithm where coarse sub-sequences providing the correction, correspond to discretizations with more strict numerical stability restrictions. We also compare the accuracy of serial versus parallel inference. This ensures that the network is not compensating for the error introduced by the MGRIT procedure, and is a study that has not been considered in prior work.
Two datasets are chosen for evaluation of the PinT algorithm. The first is the UCI-HAR dataset for the human activity recognition using smartphones (Anguita et al., 2013). This small dataset is used to demonstrate the algorithm can obtain similar accuracy to classical training approaches, while also demonstrating speedups. The second learning problem uses the HMDB51 dataset for human activity classification in videos (Kuehne et al., 2011). As the size of each image in the video is too large to directly use it as an input for a RNN-based model, we use a pre-trained CNN model to generate low-dimensional input features for the GRU network. This approach has the advantage of fully capturing spatio-temporal information in videos (Yue-Hei Ng et al., 2015; Donahue et al., 2015; Wang et al., 2016). In our experiments, extracted image features from the ResNets pre-trained on ImageNet are used as input sequences for GRU networks. Using this architecture on the HMDB51 dataset, the MGRIT algorithm demonstrates a 6.5× speedup in training over the serial approach while maintaining accuracy.
The paper is organized as follows. Section 2 presents background on GRUs and their formulation as ODEs. Section 3 details PinT training for GRU using the MGRIT method. Section 4 compares the performance of the proposed approach on sequence classification tasks. We conclude in Section 5. The appendix provides greater detail on the MGRIT method including a serial implementation and parameters used in the experiments.

heading: 2 GATED RECURRENT UNIT

text:
The GRU network, proposed by Cho et al. (2014), is a type of RNN and a state-of-the-art sequential model widely used in applications such as natural language processing and video classification. Each GRU cell at time t takes input data xt and the previous hidden state ht−1, and computes a new hidden state ht. The cell resets or updates information carefully regulated by gating units in response to prior hidden state and input sequence data stimulation as follows:
rt = σ(Wirxt + bir +Whrht−1 + bhr)
zt = σ(Wizxt + biz +Whzht−1 + bhz)
nt = ϕ(Winxt + bin + rt (Whnht−1 + bhn)) ht = zt ht−1 + (1− zt) nt.
(1)
We refer to this method as Classic GRU in the text.
GRUs to ODEs Recent work has shown how RNNs can be modified to obtain an ODE description (Chang et al., 2019; Habiba & Pearlmutter, 2020). Of particular interest here is those works for
GRU networks (De Brouwer et al., 2019; Jordan et al., 2021) where the ODE is ∂tht = −(1− zt) ht︸ ︷︷ ︸
stiff mode
+(1− zt) nt. (2)
Jordan et al. (2021) present a detailed description of the network behavior. For our purposes, note that the first term can generate stiff dynamics that drive the hidden state h towards a manifold specified by the data. The second term injects changes from the data. The term zt controls the decay rate towards the manifold for each component of the hidden state. If zt is near 1, then the rate is slow. If zt is close to zero the decay rate is rapid and the first term can be stiff, constraining the choice of time step (for numerical treatment of stiff ODE’s see Wanner & Hairer (1996)).
The classic GRU method in Eq. 1 is rederived using an explicit Euler scheme with a unit time step to integrate Eq. 2. Adding and subtracting ht−1 on the right hand side of the hidden state update, and for notational purposes multiplying by γ = 1, yields
ht = Φγ(xt, ht−1) := ht−1 + γ (−(1− zt) ht−1 + (1− zt) nt) . (3) The notation Φγ , introduced here, denotes the classic GRU hidden state update if γ = 1.

heading: 3 PARALLEL-IN-TIME TRAINING

text:
Training GRU networks is typically done with the BPTT algorithm (Sutskever, 2013). BPTT performs a forward evaluation of the neural network, followed by backpropagation to compute the gradient. Both forward and backward propagation are computationally expensive algorithms that do not naturally lend themselves to parallelization in the time domain. Despite this, we develop a methodology achieving the parallel distribution shown in Fig. 1. This distribution is enabled by computing the forward and backward propagation inexactly, trading inexactness for parallelism and ultimately run time speedups. MGRIT, a multigrid method, is used to compute the hidden states for both forward and backward propagation.
Multigrid methods are used in the iterative solution of partial differential equations (Brandt, 1977; Douglas, 1996; Briggs et al., 2000; Trottenberg et al., 2000). Recent multigrid advances have yielded the MGRIT method achieving parallel-in-time evolution of ODEs like Eq. 2 (see for instance Falgout et al. (2014); Dobrev et al. (2017); Gunther et al. (2020)). Thus, multigrid and MGRIT methods are reasonably well developed; however, they remain quite complex. Here, we briefly outline the algorithm for forward propagation, and present a brief motivation for the method in Appendix A.1. In Section 3.3, an implicit formulation for GRU that handles the stiff modes is proposed. This is the key advance that enables the use of MGRIT for training GRU networks.

heading: 3.1 NOTATION AND PRELIMINARIES

text:
Denote a sequence by ~u = {ut}Tt=0 where the length T will be clear from context, and u0 = 0 by assumption. This notation is used for the input data sequence ~x and the sequence of hidden states ~h. The GRU forward propagator in Eq. 3 can be written as the system of nonlinear algebraic equations
f(~h) = {ht − Φ1(xt, ht−1)}Tt=1 . (4)
When f(~h) = 0, then ~h is the hidden state for the GRU networks with input sequence ~x.
A two-level multigrid algorithm uses a fine sequence on level l = 0, corresponding to a data input of length T . The coarse level, defined on level l = 1, has data of length T/cf where cf is an integer coarsening factor. Sequences are transferred from the fine to coarse level using a restriction operator
~U = R(~u) := {ucf t} T/cf t=1 . (5)
Applying R to ~x gives a coarse sub-sequence of the original. Time points on the fine level that are included in the coarse level are referred to as coarse points (see Fig. 2). Transfers from coarse to fine grids are denoted by the prolongation operator
~u = P(~U) := {Ubt/cfc} T t=1. (6)
Prolongation fills intermediate values using replication of the prior coarse point. The nonlinear system from Eq. 4 is modified using ∆t = cf , yielding the residual for the coarse network:
fc( ~H) = { Ht − Φcf (Xt, Ht−1) }T/cf t=1 . (7)
Notice that in addition to a larger time step, the length of sequence ~H is smaller than the fine grid’s.
Finally, we introduce the notion of relaxation of the hidden state. Here, we define this formally, while in Fig. 2 a schematic view of these algorithms is taken. Relaxation applies propagation to an initial condition for the hidden state in a parallel fashion. F -relaxation is defined in Algorithm 1 as ΨF (cf , ~x,~h) and demonstrated in the top of Fig. 2. F -relaxation uses the hidden state at the coarse points, and fills in the fine points up to the next coarse point (exclusive). FC-relaxation, ΨFC , is identical to F -relaxation except the next coarse point is included (the bottom row in Fig. 2, where closed circles fall on the coarse grid point). Combining F - and FC-relaxation yields FCF - relaxation used in MGRIT to ensure sequence length independent convergence
~h′ = ΨFCF (cf , ~x,~h) := ΨF (cf , ~x,ΨFC(cf , ~x,~h)). (8)
The available parallelism for all relaxation methods is equal to the number of coarse points.
Algorithm 1: ΨF (γ, ~x,~h) - F -relaxation 1. parallel for tc ← 1 to T/γ 2. t = γtc 3. h′t = ht 4. for τ ← t to t+ γ − 1 5. h′τ+1 = Φγ(xτ , h ′ τ )
6. return ~h′

heading: 3.2 MULTI-GRID IN TIME

text:
Algorithm 2 presents pseudocode for the two-level MGRIT method for forward propagation. A related schematic is presented in Fig. 3 where line numbers are included in parentheses. In Fig. 3, the fine grid on the top corresponds to an input sequence of length 16. It is depicted as being distributed across 4 processors. The coarse grid representation is a sub-sequence of length 4, its propagation is limited to a single processor. Red arrows in Fig. 3 indicate grid transfers, while green arrows indicate relaxation and propagation through the network. Appendix A.1 provides a brief motivation for MGRIT, and a serial implementation is included in the supplementary materials.
Algorithm 2: MGProp(~x,~h) - Forward propagation using multigrid // Relax the initial guess
Line 1 in Algorithm 2 relaxes the initial guess, and improves the fine grid hidden state. FCF - relaxation is typically used to obtain sequence length independent convergence, though often on the fine level F -relaxation is used. Regardless, relaxation is applied in parallel yielding a processor local improved hidden state. The fine level residual given the relaxed hidden state is
~r = f(~h′). (9)
Let ~h∗ be the forward propagated hidden state that satisfies f(~h∗) = 0. This is easily computed by serial forward propagation. Rewriting ~r in terms of ~h′ and ~h∗ gives
~r = f(~h′)− f(~h∗). (10) Coarsening the residual using the restriction operator from Eq. 5, we arrive at an equation for the coarse grid approximation of ~h∗
fc(~h ∗ c) = fc(R(~h′))−R(~r). (11)
The coarse grid approximation is computed using Eq. 11 in Line 4 of the algorithm. This computation has reduced run time relative to the fine grid because the input sequence is shorter. Line 5 transfers a coarse approximation to the fine grid. The final step in the algorithm is F -relaxation on Line 6 that propagates the corrected solution from the coarse points to the fine points. Remark 1 (Multiple Iterations). MGRIT is an iterative algorithm. One pass through Algorithm 2 may not provide sufficient accuracy in the hidden state for use in neural network training. However, in Gunther et al. (2020), relatively few iterations were needed for successful training. Remark 2 (Beyond Two-Level). As presented, this is a two level algorithm. The coarse grid correction in line 4 provides an end-to-end correction at a reduced run-time of T/cf . However, for large T , this may be prohibitively expensive. The solution is to recognize that the coarse grid problem itself can be solved using MGRIT in a recursive fashion. Thus, new levels are added to reduce the run time of the coarse problem. The run time of level l with T time steps scales as c(l) = T/P+c(l+1), where P is the number of processors. Therefore, the total run time for L levels scales as
c(0) = L−1∑ l=0 T Pclf ≤ T P ∞∑ l=0 1 clf ≤ T P cf cf − 1 . (12)

heading: 3.3 COARSE GRID PROPAGATORS

text:
The discussion above assumes that all levels use the classic GRU propagator in Eq. 3 (forward Euler with a time step of ∆t = 1). This presents a problem for coarsening. Consider only the stiff term in the dynamics from Eq. 2 and assume that z ∈ [0, 1] is constant for all time. This results in the model ODE ∂th = −(1− z) h. Discretizing with forward Euler, stable time steps satisfy
|1−∆t(1− z)| ≤ 1⇒ 0 ≤ ∆t ≤ 2(1− z)−1. (13)
The model ODE shares the stiff character of the GRU formulation, and the time step is restricted to be ∆t < 2 assuming the worse case z = 0 (z = 1 implies there is no time step restriction). On the fine grid with ∆t = 1, forward Euler is not unstable. However, on coarse grids, where the time step is equal to clf , stability of forward and back propagation is a concern. In practice, we found instabilities on coarse grids materialize as poor training, and generalization.
To alleviate problems with stiffness, the traditional approach is to apply an implicit method. For instance, applying a backward Euler discretization to the model ODE yields a stability region of
|1 + ∆t(1− z)|−1 ≤ 1. (14)
Hence, the time step restriction is unlimited for the model problem (for 1 − z < 0 the ODE is unstable).
A backward Euler discretization of Eq. 2 requiring a nonlinear solver is possible; however, it is likely needlessly expensive. We hypothesize, considering the model problem, that the stiffness arises from the first term in Eq. 2. The second term incorporates the input sequence data, changing the hidden state relatively slowly to aggregate many terms over time. Finally the nonlinearity in z may be problematic; however, this term primarily controls the decay rate of the hidden state.
As a result of this discussion, we propose an implicit scheme to handle the term labeled “stiff mode” in Eq. 2. The discretization we use computes the updated hidden state to satisfy
(1 + ∆t(1− zt)) ht = ht−1 + ∆t(1− zt) nt. (15)
We refer to this method as Implicit GRU. Here, the term (1 − zt) on the left hand side and the right hand side are identical and evaluated at the current data point xt and the old hidden state ht−1. Compared to classic GRU, this method inverts the left-hand side operator to compute the new hidden state ht. This is easily achieved because (1 + ∆t(1− zt)) is a diagonal matrix. Thus, the evaluation cost of Eq. 15 is not substantially different from Eq. 3. Yet the (linear) stiffness implied by first term in the ODE is handled by the implicit time integration scheme. To simplify notation, Eq. 15 is rewritten by defining an update operator Φγ (similar to Eq. 3)
ht = Φγ(xt, ht−1; ) := (1 + γ(1− zt))−1 (ht−1 + γ(1− zt) nt) . (16)
The subscript γ is the time step size, and the learned parameters are identical to classic GRU. This propagator replaces the classic GRU propagator in all levels of the MGRIT algorithm.

heading: 3.4 TRAINING WITH MGRIT

text:
Training using MGRIT follows the traditional approaches outlined in Goodfellow et al. (2016). In fact, the MGRIT algorithm has been implemented in PyTorch (Paszke et al., 2019), and the Adam implementation is used without modification. The difference comes in the forward propagation, and backward propagation steps. These are replaced with the MGRIT forward and backward propagation. We provide a more detailed description of this approach in Appendix A.2. However, the iterations used for these steps are much fewer than would be required to achieve zero error in the propagation. We found that using 2 MGRIT iterations for the forward propagation, and 1 iteration for back propagation was all that is required to reduce the training loss.
Remark 3 (Serial vs. Parallel Inference). As training with MGRIT is inexact, a question arises: does inference with the network evaluated in parallel differ from when the network is evaluated in serial? Our results below indicate consistency between serial and parallel inference even when training is done with MGRIT. This suggests that error introduced by MGRIT is not a severe limitation.

heading: 4 EXPERIMENTAL EVALUATION

text:
In this section, we evaluate the performance of our parallel-in-time training approach on sequence classification tasks such as human activity recognition and video classification. Experiments are run on a multi-node supercomputer, named , with each node comprised of 2 sockets with a 2.3 GHz Intel Xeon processor, each processor has 18 cores yielding 36 cores per node. Inter-process communication is handled using the Message Passing Interface (MPI) (Gropp et al., 1999) which has been optimized for the Intel Omni-Path interconnect. For the video classification dataset, 9 OpenMP threads were used per MPI rank, yielding 4 ranks per node.
We used two publicly available datasets – UCI-HAR (Anguita et al., 2013) and HMDB51 (Kuehne et al., 2011). The UCI-HAR dataset is a collection of 10299 sequences of human activities labeled with 6 activities such as walking, walking upstairs, walking downstairs, sitting, standing, and laying. Each sequence consists of 128 consecutive input data samples of size 9 measuring three different types of motion recorded by smartphone accelerometers. The HMDB51 dataset contains 6726 video clips recording 51 human activities. To limit computational costs, in several of our experiments a subset of classes are used (see Table 2 in the appendix). For consistency, each video clip is converted (truncated/padded) into 128 frames. Each image in the sequence is of size 240 × 240 pixels, and is passed through ResNet18, 34 or 50 networks pre-trained on ImageNet to extract features for the GRU networks (Deng et al., 2009; He et al., 2016). See Appendix A.3.2 for architectural details.
The software implementation of the MGRIT GRU propagation algorithm is a combination of the C++ library and PyTorch. The combined capability is referred to as 1. The implementation leverages performance enhancements developed for classic GRU within PyTorch. While Implicit GRU has been carefully implemented, further optimizations are certainly possible to achieve greater computational performance. A GPU implementation of the MGRIT algorithm is under development. Finally, it is important to note that the time parallel implementation does not interfere substantially with traditional data and model forms of parallelism.

heading: 4.1 UCI-HAR DATASET

text:
In Fig. 4, the left image shows the test accuracy variability over 16 training runs for the UCI-HAR dataset. The set of parameters is detailed in Table 1 of the appendix. Solid lines are the mean accuracy, while the colored shadows indicate one standard deviation. The classic GRU (black) achieves 90% accuracy. Implicit GRU (blue line) does as well though with a more rapid increase with epochs. When training with MGRIT, three levels with a coarsening rate of 4 are used for each sequence of length 128. This yields a hierarchy of grids of size 128, 32, and 8. Here, we see that the test accuracy using parallel inference (red line) is modestly lower than the serial case. Inference with the serial algorithm suffers a little more. The plot on the right demonstrates a roughly 2× speedup over the classic GRU method in serial (black dashed line) when using the MGRIT algorithm. This dataset is very small, and likely suffers from the effects of Amdahls’ law (see Sun & Ni (1993) for discussions of Amdahl’s law). Thus, even achieving 2× speedup is notable. Beyond 32 MPI ranks there is no further speedup due to communication swamping potential gains from parallelism.
1Available at https://github.com/

heading: 4.2 HMDB51 DATASET

text:
Fig. 5 presents parameter studies elucidating the scalability and performance of the MGRIT algorithm. The left image in Fig. 5 shows the residual norm as a function of MGRIT iterations for both forward (solid lines) and backward (dashed lines) propagation. The colors differentiate coarsening factors. This demonstrates that the MGRIT algorithm converges in iterations marginally faster for smaller cf . Note that the deeper hierarchies implied by cf = 2 can be more expensive, while cf = 8 may not provide sufficient robustness. As a result, cf = 4 is used in the remainder. Based on the convergence plots and experimentation, only 2 iterations for forward propagation and 1 iteration of backward are used for training (sufficient accuracy indicated in the figure). Further, iterations do not improve training; a similar observation was made in Gunther et al. (2020) and Kirby et al. (2020). The right image in Fig. 5 shows the parallel performance over multiple sequence lengths. While there is some degradation in the weak scaling, the run time for a 128 timestep dataset is roughly only twice that of a dataset with sequences of length 8. Further, MGRIT performance improves for all lengths if the number of timesteps per MPI rank is 4 or greater. When timesteps per rank are too few, the computational work is dominated by parallel communication yielding no speedup. Sequences of length 128 are used in the remaining studies in this section.
Fig. 6 (left) shows test accuracy achieved for a subset of the HMDB51 dataset (see dataset details in Table 2). The Adam optimizer is used with a fixed learning rate of 10−3. The accuracy for the parallel algorithms are nearby the result obtained for classic GRU. Note that the use of different processor counts does not yield a different algorithm. In the right image, the differences between the schemes are examined after training 16 times. Here, we see that the standard deviation (shaded regions) and mean (solid lines) for the different approaches overlap each other. In addition, both serial and parallel inference of the Implicit GRU yield similar results, supporting Remark 3.
To further study the parallel performance of MGRIT training of GRU, we consider now the full HMDB51 dataset. Fig. 7 shows the run time per batch (with a batch size of 100) when different pre-trained ResNet architectures are used. In all cases, implicit GRU networks are trained with 32 MPI ranks and 9 OpenMP threads (using a total of 288 CPUs), achieving roughly 6.5× speedup over the classic GRU algorithm trained with 9 OpenMP threads. The largest amount of speedup occurs after only 4 MPI ranks. Increasing parallelism to 64 MPI ranks did not result in any speedup. Note
that in this case, each rank is responsible for computing only 2 time steps. Based on Fig. 5 (right), we hypothesize that at this scale the amount of computational work is too small to see benefits from increased parallelism. Fig. 8 (left) shows a breakdown of run times for the ResNet18 case. This log plot shows consistent gains with parallelism for the total (red), forward (green), backward (blue) and disk compute times (orange). The expense of the forward simulation dominated by the evaluation of the ResNet at each time step. The backward propagation cost is reduced because sensitivities with respect to the ResNet are not required. Finally, the disk time can be reduced by simply storing these in memory or overlapping computation and disk read. Note that both of these options are available to the MGRIT training as well. The robustness of MGRIT training applied to implicit GRU using ResNet34 on 32 MPI ranks is investigated with respect to initialization with 7 different random seeds in Fig. 8 (right). A learning rate schedule starting from 10−3 and being reduced by 1/2 after 5 epochs is used to achieve roughly 65% accuracy. The variation is relatively small, and the overall performance is robust. For more details on these studies, including a full listing of hyper-parameters, see Appendix A.3.2.

heading: 5 CONCLUSION

text:
We presented a parallel-in-time approach for forward and backward propagation of a GRU network. Parallelism is achieved by admitting controlled errors in the computation of the hidden states by using the MGRIT algorithm. A key development to make this possible is a new Implicit GRU network based on the ODE formulation of the GRU methodology. Applying MGRIT to other types of RNN (LSTM’s or those proposed in Chang et al. (2019)) is possible with an implicit solver.
The implementation of this approach leverages existing automatic differentiation frameworks for neural networks. PyTorch and its Adam optimizer was used in this work. Comparisons between PyTorch’s GRU implementation and the parallel algorithm were made on CPUs using OpenMP demonstrating a 6.5× speedup for HMDB51. A GPU version of the MGRIT technology is currently underdevelopment, though related work has demonstrated potential on Neural ODEs (Kirby et al., 2020). Other forms of data and model parallelism are complementary to the proposed technique, though undoubtedly there are unexplored optimizations.

heading: A APPENDIX

heading: A.1 WHY DOES MGRIT WORK?

text:
To help the reader who may not be familiar with the details of the MGRIT algorithm, we include this brief discussion to provide some intuition about how it works. For brevity, we focus only on MGRIT forward propagation. As discussed above, there are a number of high quality introductory texts on multigrid in general that we recommend for a more in-depth exploration. The monograph “A Multigrid Tutorial” by Briggs et al. (2000) is a really excellent starting place. For MGRIT specificaly see (Falgout et al., 2014) and Dobrev et al. (2017). Moreover, a serial implementation of the MGRIT algorithm written in python using NumPy is included in the supplementary material (see Section A.4 for a brief overview). The ODE used in this code is
∂th(t) = − 1
2 h(t) + σ(Ah(t) +Bd(t) + b). (17)
Here, h is the hidden state, d represents the user data, andA,B, b are randomly selected weights and biases independent of time. In the code, the width of the hidden state is a user defined parameter, while in the discussion that follows h ∈ R3. This ODE is meant as a simplified representation of an RNN/GRU network. Finally, realize that in the code and in this discussion only forward propagation is considered, no training is performed.
MGRIT is a divide and conquer algorithm where the problem of forward propagation is partitioned into reducing high-frequency error, and low frequency error separately. To this end, the error is written as a decomposition
~h− ~h∗ = ~e = ~ehigh-freq + ~elow-freq (18)
where ~h∗ is the exact solution, ~h is the approximation, ~e is the error, and the remaining two terms are the high-frequency and low-frequency errors. The high-frequency error is reduced in MGRIT by using FCF -relaxation. To demonstrate this we consider the ODE in Eq. 17, and choose our initial approximation by selecting random values at each time point to define ~h. This seeds the initial error with a number of high-frequency modes. This is evident in Fig. 9 plots (A) and (B), which show the initial error plotted in time and the Fourier coefficients plotted as a function of wave number. Each color represents a different component of the hidden state. The high frequencies are large and significant, even if the low frequency mode is also substantial. Plots (C) and (D) show the error and its Fourier coefficients after 4 sweeps of FCF -relaxation. Clearly, from plot (C), the solution has improved in early times in a dramatic fashion. This is expected as FCF -relaxation propagates exactly from the initial condition. In addition, the error is reduced in magnitude through out the
time domain. Not so evident from the error plot is the reduction in high frequency errors. This is clearly demonstrated in plot (D), where the low frequency modes remain relatively large while the high frequency modes are substantially damped. Thus, FCF -relaxation can be used to reduce high-frequency errors throughout the time domain, and is the first component of the MGRIT divide and conquer strategy.
The second part of the divide and conquer strategy is the coarse grid correction, where a coarse representation of the ODE is used to provide a correction based on the residual from the fine grid. The residual implicitly contains the error in the current guess of the solution on the fine grid. The final pair of images, (E) and (F), shows how the error from the fine grid behaves after it is restricted to the coarse grid, in this case with a coarsening factor of 4. Here, the highest frequency modes from the fine grid are removed from the error, leaving only modes that are low-frequency on the fine grid. The goal of the coarse grid is then to compute a correction to these low frequency modes. If this is the coarsest level, then this computation is a serial forward propagation, which exactly resolves the coarse grid error and the correction is transferred to the fine grid with the prolongation operator. For a multi-level scheme, another round of the divide and conquer algorithm is initiated recursively (applying the MGRIT algorithm again, as discussed in Remark 2). Here, the key is to recognize that parts of the low-frequency error from the fine grid are aliased onto the coarse grid, and become, relative to the coarse grid, high-frequency. This can be seen as a spreading of the
Fourier spectrum image (F) relative to image (D). Again, the FCF -relaxation algorithm will damp the high-frequency error modes on this level, and the coarse grid correction will correct the low frequency error modes. Regardless, when the recursion returns to the coarse level, the error in both the high-frequency modes and low-frequency modes will have been reduced. Thus, a correction transferred to the fine grid will provide further error reduction.
This presentation has focused on the forward propagation. However, using an adjoint ODE the motivation and implementation are nearly identical. The most notable difference in the implementation is the reversal of the time evolution.

heading: A.2 TRAINING WITH MGRIT

text:
Incorporating the time-parallel approach into a standard stochastic gradient descent algorithm requires replacing the forward and backward propagation steps with the equivalent MGRIT computations. Algorithm 3 presents pseudo-code demonstrating how that works. This algorithm requires additional notation for the parameters in the neural network, here denoted by Θ. The structure of the algorithm follows the batched gradient descent approach. In the first part, lines 3-5, the MGRIT algorithm is applied to perform forward propagation. In the second block, lines 6-9, MGRIT is called again, this time to compute the backpropagation through the neural network. Line 10 computes the gradients of the loss with respect to the weights. Note that the sensitivities computed by backpropagation are used to compute the full gradient with respect to the parameters in a parallel fashion. The final line of the algorithm is simply the gradient descent step with a specified learning rate.
Algorithm 3: MGRITTraining(lr,Batches,ForwardIters,BackwardIters) // This is a training loop over epochs, and mini-batches
1. for e ∈ 1 . . .Epochs 2. for each ~x ∈ Batches
// Iterate with MGRIT to perform forward propagation
3. ~h = ~0 4. for i ∈ 1 . . .ForwardIters 5. ~h = MGProp(~x,~h)
// Iterate with MGRIT to perform backward propagation
6. ~w = ~0 7. wT = ∇hTL 8. for i ∈ 1 . . .BackwardIters 9. ~w = MGBackProp(~x,~h, ~w)
// Compute the gradient with respect to parameters 10. gi = wi∂Θihi ∀i 11. ~Θ = ~Θ− lr~g
The algorithm specified shows MGRIT embedded within a broader stochastic gradient descent algorithm. However, the code described is already relatively modular. Lines 3-5 are simply forward propagation. Lines 6-10 use back propagation to compute the gradient. In our own PyTorch implementation, we were able to encapsulate the forward propagation in a torch.Module evaluation call. Similarly, the gradient computation is implemented to be compatible with PyTorch’s autograd capability. As a result, utilization of PyTorch’s optimizers, like Adam, is straightforward.
A final point needs to be made about the application of MGBackProp in Algorithm 3. This is meant to imply that the MGRIT algorithm is applied in a parallel-in-time fashion to the backpropagation algorithm. In our opinion, the most transparent way to do this is by considering the adjoint of Eq. 15 and its discretization. This is the approach taken in the Neural ODE paper (Chen et al., 2018).

heading: A.3 COMPUTATIONAL EXPERIMENTS

heading: A.3.1 UCI-HAR DATASET

text: The UCI-HAR dataset is used as an initial evaluation of the effectiveness of the MGRIT training for GRU. The parameters used for this study are listed in Table 1. The subsets of labeled Training Parameters are values used for both the computational resources and algorithmic details of the optimizers. Note that the number of threads used per MPI Rank is only 1 as this dataset is relatively small. This implies, on our architecture that contains 36 cores per node, only one node is used for this simulation. Those labeled GRU Parameters are network architecture specific parameters. All of these apply to Classic GRU, serial Implicit GRU, and parallel Implicit GRU. Note that a time step of ∆t = 1 is always chosen to ensure alignment with the Classic GRU architecture. Finally, the MGRIT parameters have been chosen by following the recommendations in prior works of Gunther et al. (2020) and Kirby et al. (2020).

heading: A.3.2 HMDB51 DATASET

text:
The HMDB51 dataset is composed of 51 classes, and 67262 different videos. This full dataset was used in the generation of the results in Fig. 7 and Fig. 8. The subset of the data was used to generate Fig. 5 and Fig. 6. The breakdown of training to testing videos and the specific classes used are detailed in Table 2. The images in the video are cropped and padded so the image sizes are all 240×240. Further, to focus on the performance of the MGRIT algorithm, we truncate/pad (with blank frames) the video sequences to a specified length.
Table 3 shows the parameters used for the simulations. These follow the same grouping convention as the parameters for the UCI-HAR dataset. This dataset is substantially larger and more computationally expensive than UCI-HAR dataset. As a result, more computational resources are used. In the table, this is demonstrated by the use of 9 OpenMP threads per MPI rank. On the platform used, this implies there are 4 MPI ranks per node (each of the 36 cores on a node gets assigned a thread). While we experimented with different configurations, we were unable to get more performance by increasing the number of threads per MPI rank. We suspect that this has to do with the configuration of the memory, and software stack on our system. The table also shows the types of studies and parameter variations we have performed and documented in the present work. Parameter names denoted with an asterisk are varied in our studies. For instance, Adam is used as the optimizer
2This differs from the documented number somewhat. At the time we downloaded these videos, several were unavailable.
and ∆t = 1 on the fine grid in all experiments, while the learning rate is varied in the experiment producing Fig. 8.
Images from the video sequence are preprocessed during training and inference into reduced features (size 1000, also specifying the GRU hidden size) using pre-trained ResNet18, ResNet34 or ResNet50 networks. In training, the weights and biases of the ResNet’s remain fixed. Further, the MGRIT algorithm computes the application of the ResNet only once on the fine grid. The coarse grids make use of those precomputed features to save computational time. In the MGRIT backpropagation step, the features are also reused.
The experiments producing Fig. 5 and Fig. 6 used the smaller subset dataset. For the residual convergence study, the MGRIT coarsening rate was varied. For cf = 2, the number of levels used was 6, with 128, 64, 32, 16, 8 and 4 timesteps on each level (from fine to coarse). For cf = 4 (the baseline), the number of levels used was 3, with 128, 32 and 8 timesteps on each level. Finally, for cf = 8, the number of levels used was 2, with 128 and 16 timesteps on each level. A minimum number of timesteps of 4 on the coarse level was used that defined the depth of the MGRIT hierarchy. For the performance scaling study in the right image of Fig. 5, varying sequence lengths were used, and, as a result, the time domain parameter T also varied. The remaining parameters followed the baseline from Table 3, notably using cf = 4, and a pre-trained ResNet18.
For Fig. 6, ResNet18 was used and all parameters followed the baseline. In addition to the test accuracy, we include an analogous figure with the loss in Fig. 10. In all cases, the loss decreases steadily regardless of the number of processors. Varying training over multiple initial states yields a similar reduction in loss for the MGRIT algorithm.
Figures 7 and 8 use the full data set. All performance simulations use the baseline parameter set with a sequence length of 128. The plot of the timings breakdown uses ResNet18. Finally, in the full dataset accuracy results using ResNet34 in Fig. 8 (right), the baseline parameters are used but
the learning rate is varied. The learning starts from 10−3, and then is reduced by 0.5 every 5 epochs, terminating at 1.25× 10−4.

heading: A.4 SUPPLEMENTAL MATERIAL: SERIAL MGRIT CODE

text: To help the reader new to MGRIT, we have included a serial implementation of the two-level algorithm using python with minimal prerequisites (only NumPy). It has also been included in our library as an example and a teaching tool. This example does no training, but produces only one forward solve of the ODE in Eq. 17. By default, the ODE has dimension 10, α = 1.0, and the number of steps is 128. The result of running this example is:

heading: Two Level MG

text:
error = 1.6675e-01, residual = 1.1287e-01 error = 3.1409e-03, residual = 1.9173e-03 error = 5.4747e-05, residual = 3.9138e-05 error = 3.5569e-07, residual = 2.6001e-07 error = 3.2963e-09, residual = 2.9624e-09 error = 5.5450e-11, residual = 3.6914e-11 error = 5.7361e-13, residual = 3.9258e-13 error = 5.1118e-15, residual = 4.5498e-15
As the weights and biases are randomly selected, the specific values may be different depending on the seed. However, in general, the error reduction should be similar. The code itself is commented. Further, the implementation of the two level algorithm is labeled with the line numbers used in Algorithm 2.


summaries: 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.

summaries: The paper proposes a new neural architecture named MRMD-AE that can be applied on noisy fMRI data in different tasks. Subject-specific decoders are used to more directly recover individual signals, while the encoder is shared across every subject under the assumption that every person will share common low-dimensional features for the same stimuli. A key component of this architecture is the usage of a regularisation term named PHATE (previously introduced in the literature) which allows the latent space to not be split into individual embeddings and be extendable for unseen data. The paper empirically shows how this architecture improves metrics on classification tasks when compared to previously used techniques.

summaries: This paper proposes a neural network-based modeling strategy to learn a common latent space from multi-subject fMRI data. In addition to capturing a useful common latent space, the proposed technique is further able to disentangle common representational patterns from subject-specific variations through the use of subject-specific decoders. The authors impose meaningful and desirable priors on the latent embedding, like geometric regularization and cross-subject embedding alignment. Unlike other manifold learning techniques, the proposed deep neural network modeling framework lends itself well to extendibility to new data (stimuli) since the PHATE embeddings are only required at training time in the geometric regularization loss. The proposed framework is tested on two large fMRI datasets and an improved stimulus decoding (from the shared space) and cross-subject translation accuracy is achieved over competitive baselines.

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

abstractText: 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.

title: Metric multidimensional scaling (mds): analyzing distance matrices

venue: Encyclopedia of measurement and statistics,

authors:
Matthew Amodio
David Van Dijk
Krishnan Srinivasan
William S Chen
Hussein Mohsen
Kevin R Moon
Allison Campbell
Yujiao Zhao
Xiaomei Wang
Manjunatha Venkataswamy

title: Exploring single-cell data with deep multitasking neural networks

venue: Nature methods,

authors:
Mikhail Belkin
Partha Niyogi

title: Laplacian eigenmaps for dimensionality reduction and data representation

venue: Neural computation,

authors:
Jonathan Charles William Brooks
Olivia Kate Faull
Kyle TS Pattinson
Mark Jenkinson

title: Physiological noise in brainstem fmri

venue: Frontiers in human neuroscience,

authors:
Nathan Brugnone
Alex Gonopolskiy
Mark W Moyle
Manik Kuchroo
David van Dijk
Kevin R Moon
Daniel Colon-Ramos
Guy Wolf
Matthew J Hirn
Smita Krishnaswamy

title: Coarse graining of data via inhomogeneous diffusion condensation

venue: IEEE International Conference on Big Data (Big Data),

authors:
Erica L. Busch
Lukas Slipski
Ma Feilong
J. Swaroop Guntupalli
Jeremy F. Huckins
Samuel A. Nastase
M. Ida Gobbini
Tor D. Wager
James V. Haxby

title: Hybrid hyperalignment: A single high-dimensional model of shared information embedded in cortical patterns of response and functional connectivity

authors:
Janice Chen
Yuan Chang Leong
Christopher J. Honey
Chung H. Yong
Kenneth A. Norman
Uri Hasson

title: Shared memories reveal shared structure in neural activity across individuals

venue: Nature neuroscience,

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

title: A reduced-dimension fmri shared response model

venue: Advances in Neural Information Processing Systems,

authors:
Katherine Minjee Chung
Jaffarguriqbal Singh
Lauren Lawres
Kimberly Judith Dorans
Cathy Garcia
Daniel B Burkhardt
Rebecca Robbins
Arjun Bhutkar
Rebecca Cardone
Xiaojian Zhao

title: Endocrine-exocrine signaling drives obesity-associated pancreatic ductal adenocarcinoma

authors:
Andrés F Duque
Sacha Morin
Guy Wolf
Kevin Moon

title: Extendable and invertible manifold learning with geometry regularized autoencoders

authors:
Siyuan Gao
Gal Mishne
Dustin Scheinost

title: Non-linear manifold learning in fmri uncovers a low-dimensional space of brain

venue: dynamics. bioRxiv,

authors:
Scott Gigante
Adam S Charles
Smita Krishnaswamy
Gal Mishne

title: Visualizing the phate of neural networks

authors:
Michael Hanke
Florian J. Baumgartner
Pierre Ibe
Falko R. Kaule
Stefan Pollmann
Oliver Speck
Wolf Zinke
Jörg Stadler

title: A high-resolution 7-tesla fmri dataset from complex natural stimulation with an audio movie

venue: Scientific Data,

authors:
James V Haxby
J Swaroop Guntupalli
Andrew C Connolly
Yaroslav O Halchenko
Bryan R Conroy
M Ida Gobbini
Michael Hanke
Peter J Ramadge

title: A common, high-dimensional model of the representational space in human ventral temporal cortex

authors:
CO Husler
M Hanke

title: An annotation of cuts, depicted locations, and temporal progression in the motion picture "forrest gump" [version 1; peer review: 3 approved

authors:
Kui Jia
Lin Sun
Shenghua Gao
Zhan Song
Bertram E Shi

title: Laplacian auto-encoders: An explicit learning of nonlinear data

venue: manifold. Neurocomputing,

authors:
Manik Kuchroo
Jessie Huang
Patrick Wong
Jean-Christophe Grenier
Dennis Shung
Alexander Tong
Carolina Lucas
Jon Klein
Daniel Burkhardt
Scott Gigante

title: Multiscale phate exploration of sars-cov-2 data reveals multimodal signatures of disease

venue: bioRxiv,

authors:
Leland McInnes
John Healy
James Melville

title: Umap: Uniform manifold approximation and projection for dimension reduction

venue: arXiv preprint arXiv:1802.03426,

authors:
Kevin R. Moon
David van Dijk
Zheng Wang
Scott Gigante
Daniel B. Burkhardt
William S. Chen
Kristina Yim
Antonia van den Elzen
Matthew J. Hirn
Ronald R. Coifman
Natalia B. Ivanova
Guy Wolf
Smita Krishnaswamy

title: Visualizing structure and transitions in high-dimensional biological data

venue: Nature Biotechnology,

authors:
Michael Moor
Max Horn
Bastian Rieck
Karsten Borgwardt

title: Topological autoencoders

authors:
Jeanette A Mumford
Russell A Poldrack

title: Modeling group fmri data

venue: Social cognitive and affective neuroscience,

authors:
Samuel A Nastase
Valeria Gazzola
Uri Hasson
Christian Keysers

title: Measuring shared responses across subjects using intersubject correlation

venue: Social Cognitive and Affective Neuroscience,

authors:
Jenna L Pappalardo
Le Zhang
Maggie K Pecsok
Kelly Perlman
Chrysoula Zografou
Khadir Raddassi
Ahmad Abulaban
Smita Krishnaswamy
Jack Antel
David van Dijk

title: Transcriptomic and clonal characterization of t cells in the human central nervous system

venue: Science Immunology,

authors:
Bastian Rieck
Tristan Yates
Christian Bock
Karsten Borgwardt
Guy Wolf
Nicholas Turk-Browne
Smita Krishnaswamy

title: Uncovering the topology of time-varying fmri data using cubical persistence

authors:
Ayan Sengupta
Falko R. Kaule
J. Swaroop Guntupalli
Michael B. Hoffmann
Christian Häusler
Jörg Stadler
Michael Hanke

title: A studyforrest extension, retinotopic mapping and localization of higher visual areas

authors:
Stephen M Smith
Aapo Hyvärinen
Gaël Varoquaux
Karla L Miller
Christian F Beckmann

title: Group-pca for very large fmri

venue: datasets. Neuroimage,

authors:
Laurens Van der Maaten
Geoffrey Hinton

title: Visualizing data using t-sne

venue: Journal of machine learning research,

authors:
Kiran Vodrahalli
Po-Hsuan Chen
Yingyu Liang
Christopher Baldassano
Janice Chen
Esther Yong
Christopher Honey
Uri Hasson
Peter Ramadge
Kenneth A. Norman
Sanjeev Arora

title: Mapping between fmri responses to movies and their natural language annotations

authors:
Guy Wolf
Aviv Rotbart
Gil David
Amir Averbuch

title: Coarse-grained localized diffusion

venue: Applied and Computational Harmonic Analysis,

authors:
Wenchao Yu
Guangxiang Zeng
Ping Luo
Fuzhen Zhuang
Qing He
Zhongzhi Shi

title: Embedding with autoencoder regularization

venue: In Joint European Conference on Machine Learning and Knowledge Discovery in Databases,

authors:
Yujiao Zhao
Matthew Amodio
Brent Vander Wyk
Bram Gerritsen
Mahesh M Kumar
David van Dijk
Kevin Moon
Xiaomei Wang
Anna Malawista
Monique M Richards

title: Single cell immune profiling of dengue virus patients reveals intact immune responses to zika virus with enrichment of innate immune signatures

venue: PLoS neglected tropical diseases,

heading: 1 INTRODUCTION

text:
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.

heading: 2 RELATED WORK

text:
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).

heading: 3 PRELIMINARIES

text:
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.

heading: 4 MULTI-DECODER MANIFOLD-REGULARIZED AUTOENCODER

heading: 4.1 PROBLEM SETUP

text:
Let X = {Xk : k = 1, . . . ,m} be a set of common stimuli for m timepoints, and let Yj = {Yj,k ∈ R
q : k = 1, . . . ,m}, j = 1, . . . , n 1, be observed fMRI neuronal activations for n subjects over these stimuli. We seek to learn a low dimensional latent space M ⊂ Rd (where d q) with associated map f : Rq → Rd and inverse maps gj : Rd → Rq, j = 1, . . . , n, such that: (a) The mappings corresponding to each stimulus Xk in subjects Yi and Yj are close to one another, i.e., ‖f(Yi,k) − f(Yj,k)‖2 ≤ . (b) Stimulus features (for each Xk) can be computed by a function f(Yj,k) 7→ f ′(Xk) of the low dimensional latent space. (c) Stimulus features can be inferred in an unseen/untrained subject Yn+1 from the same latent space via f(Yn+1,k) 7→ f ′(Xk), at any timepoint k. (d) Subject-specific responses can be predicted on unseen stimulus Xm+1 (e.g., at held-out timepoint m + 1) based on the shared latent space, i.e., gj(f(Yi,m+1)) = Yj,m+1 for subjects i, j. Note here that each Yi,k is a high dimensional measurement consisting of raw activation values at hundreds to thousands of voxels depending on the respective region of interest (ROI). By contrast, we expect f(Yi) to be vastly lower dimensional. We analyzed the spectral entropy of PHATE diffusion potential of the fMRI activation data used in the following experiments and found the intrinsic dimensions of the data ranging 10 to 40 (Appendix Table 1).

heading: 4.2 MRDM-AE ARCHITECTURE

text:
Single Encoder – Multiple Decoders To achieve the properties outlined in the previous section, we propose a manifold-regularized common-encoder, multiple-decoder model (Figure 1). The common encoder f(Yj,k) = Enc(Yj,k) takes any set of neuronal activities from any subject and maps them in a many-to-one fashion to the same latent space L, with the same stimulus prediction. Our model features several decoders gj(f(Y·,k)) = Dec[j](f(Y·,k)) ≈ Yj,k, to decode back to subject specific responses. The rationale for using a common encoder is for the model to learn a common fMRI language from which subject specific responses can be derived. While the fMRI data exhibits strong batch effect between subjects, the common encoder builds invariance in the latent space to individual noise in the encoded space. The individual decoders enable incorporating subjectdependent interpretation of the latent information from the common encoder. Thus this framework pushes individual noise into the decoding space. In particular, the decoders can learn to reconstruct the subject-specific signals and recreate batch effect to reconstruct data. The flexibility provided by individual decoders allow the single common encoder to only retain input-dependent variability between subject responses that is sufficient for the decoders to reconstruct the measurements. This base model is trained with a reconstruction error penalty for each subject seen at training time.
Lreconstructionf,g1,...,gn (Y1, . . . , Yn) = ∑ j,k ‖Yj,k − gj(f(Yj,k))‖2
1During preprocessing, functional images were aligned to the common template based on anatomical landmarks and one time point of a fMRI region of interest is a vector of length equal to the number of voxels.
Embedding Alignment Penalty Despite the common encoder (and individual decoders), the network may learn to separate subject-specific responses into their own latent sub-spaces for ease of decoding. To counteract this tendency we explicitly push the latent encodings of individual responses to common stimuli together by penalizing Euclidean distances in the encoded space between timepoints corresponding to common stimuli. Therefore, we implement the embedding alignment penalty as Lalignmentf,g1,...,gn(Y1, . . . , Yn) = ∑ i 6=j,k ‖f(Yi,k)− f(Yj,k)‖2.
Manifold Regularization Penalty In order to ensure a smooth embedding space that enables faithful interpolation, inference, and extendibility to new data, a geometric regularization should be applied to ensure the latent space learned by the autoencoder follows the data geometry.To this end, within each decoder, we apply a regularization in the bottleneck layer using the manifold embedding Pj that are learned with PHATE (Moon et al. (2019)) using the train data for each subject j. We note that the strong subject batch effects in fMRI data make the task of directly learning a subject-independent manifold with PHATE challenging. Indeed, without careful consideration, it often yields a group fMRI data geometry with variance dominated by batch effects, which would then submerge stimulus-related variance that is of actual interest. We therefore bypass this issue here by only considering individual-subject manifolds and letting the combined autoencoder deduce a latent representation that can easily be mapped to them. Thus, as shown in Figure 1,we split each decoder into gj = g′j ◦g′′j , where intuitively we expect the first part g′j to “implement batch effects” going from the common latent representation to a subject-dependent one that captures well the intrinsic structure of their stimuli responses. It is to this first layer of g′j we apply the manifold regularization. The second part g′′j will then map this intrinsic (manifold) representation to the observable measurements to finish reconstruction. In this way, we obtain an intermediary representation after the first layer of each decoder (where the single-layer limitation restricts the implemented “batch effects” to be relatively simple) that should match the PHATE embeddings, implemented in training via the manifold regularization as Lgeometricf,g′1,...,g′n(Y1, . . . , Yn) = ‖g ′ j(f(Yj,k))− Pj(Yj,k)‖2.
Combined Loss To combine three loss terms together, we introduce coefficients λ, µ controlling the amount of alignment and geometric regularization (correspondingly) applied to the hidden layers of the MRMD-AE, yielding the total loss
L = ∑ j,k ‖Yj,k − gj(f(Yj,k))‖2 +λ ∑ i6=j,k ‖f(Yi,k)− f(Yj,k)‖2 +µ‖g′j(f(Yj,k))−Pj(Yj,k)‖2 (1)
optimized over network functions f , gj = g′j ◦ g′′j , j = 1, . . . ,m, parameterized via neural networks and organized together as illustrated in Fig. 1.
Additional Cross-Subject Training Penalty The multiple decoder framework allows us to add carefully designed loss penalties for specific tasks. For example, we can explicitly train for crosssubject translation. We achieve the purpose by adding the following regularization to encourage translation,
Ltranslationf,g1,...,gn(Y1, . . . , Yn) = ∑ j,k ∑ l 6=j ‖Yj,k − gj(f(Yl,k))‖2. (2)
In most of the following experiments this cross-subject penalty was not used, as we did not aim to train for translation. In the last section of Results we explore the use of this penalty to facilitate learning to translate between subjects’ fMRI activities.

heading: 5 RESULTS

text:
To empirically validate MRDM-AE, we test it on tasks of manifold extension, stimulus classification, and cross-subject translation. The manifold-extension tasks test the ability of MRDM-AE to place points not seen during training appropriately within the learned data manifold in the embedded space. The classification tasks analyze the ability of MRDM-AE to capture stimulus-relevant signal from fMRI data, i.e., features of the movies from voxel activations. Cross-subject classification and translation tasks test the ability of MRDM-AE to capture a shared latent space that can be used for common encoding and individual decoding. These tasks are tested on two different datasets comprising of subjects watching different movies. We compare MRDM-AE to several baselines including dimensionality reduction methods (PCA, PHATE), the shared response model (SRM) commonly used in fMRI data, as well as ablations of our autoencoder method.
Datasets The Sherlock dataset is an open-access dataset of 16 participants watching a 48-minute clip of the BBC television series “Sherlock.” Data was collected in two runs of 946 and 1030 timepoints (repetition time; TRs) respectively. This data was downloaded from the DataSpace public repository. For full details on the stimulus, see the original publication of this dataset (Chen et al. (2017)). Annotations for the Sherlock movie stimulus were released with Vodrahalli et al. (2018) and available freely on Github. These annotations were created by humans who viewed the film and labeled each timepoint on a variety of metrics. In our classification tasks, we used two sets of binary labels that correspond with visual and auditory aspects of the stimulus: whether a given timepoint featured an indoor or outdoor scene (Indoor/Outdoor) and whether or not the timepoint featured music (Music Present).
As a second testbed, we used the StudyForrest open-access dataset downloaded from DataLad; full details for this can be found at the original publication of the movie dataset (Hanke et al. (2014)) and the localizer extension dataset (Sengupta et al. (2016)). Here we included data from 15 participants who completed both the movie-viewing and functional localizer tasks, where participants viewed 24 unique grayscale images of objects from 6 categories (faces, bodies, houses, small objects, outdoor scenes, and scrambled images) in a separate session to localize brain regions that reliably respond to a specific category. In the movie task, participants watched a 2-hour version of the movie Forrest Gump. In a series of classification tasks, we use labels for each timepoint of the movie indicating the time of day (2 levels), and the flow of time in the narrative (4 levels), annotated in (Husler & Hanke (2016)), and image categories for the localizer data (6 levels). Autoencoder Hyperparameters All autoencoder based models have three fully connected layers in the encoder and the decoder(s), with a bottleneck latent space layer in between. The hidden neurons have the configurations of 256-128-64-20-64-128-256. Leaky ReLU activations are applied on all of the layers. We used the Adam optimization algorithm with a learning rate of 0.001, and a batch size of 64. The neural networks were trained for 4000 epochs. When regularization is used, λ and µ are either 0.1 or 0.01 and the same set of values are used across subjects for a specific task. Model Comparison We compare our proposed model with a standard autoencoder (AE) with the same hidden dimensions as the proposed MRMD-AE, but with a single decoder for all subjects. We refer a standard autoencoder with manifold regularization applied to the bottleneck layer as MR-AE. Dimensionality Reduction Methods Non-linear dimensionality reduction methods such as PHATE (and other methods such as t-SNE and UMAP) typically provides fixed latent representations for the input data, and do not provide a natural embedding function that can project new data to an existing manifold. PHATE preserves manifold structure while tSNE and UMAP usually shatter data trajectories to create clusters even when they do not exist ( Moon et al. (2019), Kuchroo et al. (2020), Rieck et al. (2020)). We therefore compare with an adaption of PHATE to interpolate new datapoints via landmark approach used by Moon et al. (2019), a subsampling method that was used and shown to be effective in improving scalability of PHATE. We also compare with linear PCA approach. For within-subject analyses, the dimensionality reduction was learned for a subject on one half of the dataset and applied to unseen timepoints. For cross-subject analyses, the dimensionality reduction was learned on all but one subject and applied to the held-out subject, and repeated holding out each subject once. For cross-subject analyses, we also compared to SRM by learning a reduceddimensional shared signal on data from all but one subject and evaluating the fit of the shared space on the held-out subject, repeating for all subjects. Place Future fMRI Time Points on Learned Manifold We first show that MRMD-AE achieves efficient and accurate manifold extension by placing new fMRI time points at the right positions on learned manifold. Consider the fMRI time series contain two parts, including the first mtrain TRs that we receive as pilot data with which we learned individual manifolds Ptrain, a mtrain× d matrix, with each row being the PHATE coordinates of one TR. We would like to faithfully extend these learned manifolds to subsequently receive test data of future mtest fMRI TRs, so that the extended manifold Ptest closely resemble the corresponding time points ground truth manifold P , that is learned from all mtrain +mtest TRs.
Since the manifold coordinates are fixed to the input data, we cannot directly compare the Ptest with P[TRtest], the corresponding rows of P . Instead, we first compute the manifold distance matrices between each test and train time point pairs, where D̂(i, j) = ||Ptest(i, ) − Ptrain(j, )||, and the ground truth D(i, j) = ||P(i, )− P(j, )||, i ∈ TRtest, j ∈ TRtrain. Then we evaluate the similarity using the MSE between D̂ and D. Fig.2 shows that MRMD-AE achieves lower or matching MSE than the landmark approach, especially with lower percentages of train data. The extended manifold
generated by MRMD-AE more closely resemble the ground truth. MRMD-AE also outperforms MR-AE, suggesting the advantage of per-subject decoders in reconstructing subject-specific signals.
We also use the extended manifold embedding to predict movie stimulus labels. An support vector classifier (SVC) is trained on each subject’s extended manifold coordinates and the average accuracy and standard deviation via a 5-fold cross validation are reported. We found that MRMD-AE extended manifold achieves higher classification accuracies than the landmark approach, suggesting the embeddings retain more meaningful signals that can be used to recover signals of the movie stimulus (right panel of Fig. 2 and Table 1). Effect of Latent Space Alignment Penalty We visualize the raw data and the embedding f(Yj), j = 1, 2, . . . , 15 of the Forrest image localizer data using PHATE, and color the plots by subject ID in Figure 3. The raw data visualization exhibits clear batch effect. By increasing λ in Eq. 1 and penalizing latent space misalignment, the batch effect in the latent space is eliminated. We computed the average Earth Mover’s Distance (EMD) between subject embedding pairs and showed improved alignment between subjects in the latent space comparing to the raw data and the improved alignment by increasing alignment penalty. Stimuli Classification Our goal is to test how well our MRMD-AE generalizes to unseen times series of fMRI data. We evaluate the model by how well we can recover meaningful labels in the latent space. After training on the train half of data from all subjects, we applied the trained model on the test half of data and collected the latent space representation at the bottleneck layer. A SVC is trained within each subject on the test embedding and the average classification accuracy and standard deviation via a 5-fold cross validation are recorded. The average accuracy and standard deviation across subjects are reported.
Results are reported in table 1 for Sherlock movie annotation classification and table 2 for Forrest movie and image localizer classification. MRMD-AE effectively recovers the latent factors corresponding to the stimuli and achieve the best or second best performance in all combinations of ROI and stimulus annotations. We showed previously that a learned PHATE manifold via landmark does not generalize to out-of-sample data. Here we confirm that the classification accuracy of the test set decreases significantly from the train set. The linear projection operator learned by PCA also fails to generalize well to new data (Appendix Fig. 1 shows in-sample vs out-of-sample comparison for PCA and PHATE Landmark). MRMD-AE also outperforms per-subject individual MR-AE, suggesting the synergy between subjects offered by the shared encoder and multi-decoder framework benefits generalization and recovering meaningful signal from data. Appendix Figure 2 shows that across subjects, MDMR-AE consistently outperforms the other methods. Untrained Cross-Subject Translation Our MRMD-AE framework allows integration of data from multiple subjects, thereby enabling translation between subjects. Similar as before, we trained our MRMD-AE with the train half of all subjects’ data. Then we translate the test half timepoints of subject i to each of the other subjects by Ŷi→j = gj(f(Yi)). We evaluate the quality of translation by comparing the MSE and correlation between Ŷi→j and Yj( 3). It should be noted that the model was not trained for translation, thus in training, the decoders did not learn to translate other subjects data. Yet, the common encoder maps the neural activities to the same latent space and builds invariance to individual noise in the encoded space, allowing the decoders to perform the translation task. SRM was
developed to achieve this goal by mapping all subjects to a shared response, then project to individual spaces. We show that our model achieves best translation performance on various brain ROIs in at least one metric. By incorporating latent space alignment penalty (λ > 0 in Eq. 1), MRMD-AE learns a more aligned shared latent space and allows for better translation across subjects. More experimental details can be found in the supplementary materials.
Trained Cross Subject Translation Our MRMD-AE framework provides flexibility to specialize in different tasks by incorporating carefully designed loss penalties. By adding a cross-subject training penalty in Eq. 2 to the combined loss, we explicitly train our model for translating between subjects. We show in Figure 4 the improved translation performance by using different combinations of loss penalties. Combining the alignment penalty with translation penalty achieves the lowest translation MSE.

heading: 6 DISCUSSION

text:
Extracting meaningful brain mechanisms shared across subjects is a daunting task, chiefly due to stimulus-independent noise and the challenge of comparing subjects’ activity pattern in a common reference space.To address this, we proposed MRDM-AE, a multi-decoder manifold-regularized autoencoder that takes fMRI data from a group of subjects as input, first projects it to a common latent space, then to subject-specific decoders. We showed that the subject-specific decoders enabled the autoencoder to decouple shared information from individual variations such that the shared latent space offered higher-accuracy classifications of underlying stimuli, and better manifold extensions to unseen data. Finally, we showed that the individual decoders can potentially be used for subject-tosubject translation, setting up fMRI response predictions of subjects to unseen data based on training and calibration on common stimuli. Overall, we show the manifold learning and structural priors bring in improve the state of the art in shared latent space learning in fMRI from many subjects.
REPRODUCIBILITY STATEMENT
The results presented in the paper are reproducible. We will provide a link to an anonymous repository to the reviewers and area chairs. The content of the link will contain complete instructions on how to run the code to reproduce the results presented for the paper. In addition, all details for the models and training procedures are described in the main paper under Sec. 4 and the Appendix.

heading: A SUPPLEMENTARY MATERIALS

text:
fMRI pre-processing and region identification After all standard fMRI preprocessing steps, functional images from both the Sherlock and Forrest datasets were aligned to the common MNI template brain based on anatomical landmarks. In this anatomical template, we defined several regions of interest (ROIs) to target in our analyses: the early visual cortex, late visual cortex, early auditory cortex, and the posterior medial cortex (PMC). We selected these regions as they robustly respond to audiovisual stimuli (early visual, late visual, and early auditory) and PMC was targeted in the original sherlock analyses for its memory involvement (Chen et al. (2017)). We extracted timeseries data from the voxels in each of these regions (early visual: 307 voxels, late visual: 1008 voxels, early auditory: 1018 voxels, PMC: 481 voxels) and collapsed the spatial structure within this ROI into a [timepoints, voxels] matrix of activity for each voxel across time. Within this matrix, we z-scored each voxel across its timeseries to account for differences across voxels in mean activation.
Dimension reduction methods on out-of-sample data Dimensionality reduction methods do not generalize to out-of-sample data with decremented performance when extending trained models to new data points. In Figure 1 we show that the stimuli classification accuracy decreases when PCA and PHATE (via landmark approach) are applied to out-of-sample data. Experiment was conducted on Sherlock early-visual ROI data.
Appendix Figure 1: Classification accuracy comparison. PCA and PHATE both fail to generalize to out-of-sample data, shown as decremented accuracy. Landmark approach was used to extend learned PHATE embedding to new data.
Stimuli classification experimental details For 5-fold cross validation , we proceed following the standard procedures. We split the dataset into 5 groups at random, and for each group we take as a hold out or test dataset and use the remaining groups as a training data set. Fit a classifier on the training set and evaluate on the test set, retain the test accuracy and summarize after all folds. For each training set, we balance the samples to reach an even split between classes during training, so the classification tasks are entirely balanced. We repeat the procedure for each subject. The mean accuracy from the 5-fold cross validation within each subject is reported as the skill score of the subject. Then we report the mean and standard deviation across subjects. For SVM classifiers, we used the C-support vector classification from Scikit-learn package, with an RBF kernel, regularization=10, and gamma set to be ’scale’,1/(nfeatures ×X.var()) .
Stimuli classification comparison across subjects Since we perform classification within each subject, we further look into how our model compare to the other methods for each subject. We present the results on the Sherlock early visual ROI and show that higher classification accuracy was achieved with MRMD-AE manifold embedding in 13 out of 16 subjects comparing to the next best method that is a vanilla AE. Moreover, MRMD-AE outperformed individual MR-AE (third-best) in all subjects. Although there are not enough samples to quantify subject-level statistics, we show MRME-AE achieves better performance in almost all subjects( 2).
Intrinsic dimensions of fMRI data We analyzed the intrinsic dimensions of fMRI data by examining the spectral entropy of PHATE diffusion matrix of fMRI data. We refer the readers to Moon
Appendix Figure 2: Classification accuracy comparison across subjects. Classification by indooroutdoor label on Sherlock early visual ROI using the top 3 performing models in Table 1. Left: The mean accuracy from the 5-fold cross-validation within each subject was reported. Error bar was standard deviation from the 5-fold cross-validation. MRMD-AE outperforms the next best method, a vanilla AE, in 13/16 subjects and outperforms the third best method, individual MR-AE, in all subjects. Right: Box plot of the same set of accuracy values show better performance of MRMD-AE compared to the other methods.
et al. (2019) for how PHATE computes the diffusion operator and the relationship between diffusion operator to the data geometry. For the purpose of current analysis, one can consider the diffusion operator to represent the ambient data geometry and the spectral entropy analysis via Von Neumann Entropy similar to the analysis of principle components that account for most variation in PCA. We report the intrinsic data dimensions of various ROIs of the Sherlock and Forrest fMRI datasets that we use in the experiments of this paper in Appendix Table 1.
Dataset Sherlock Forrest ROI early-visual pmc early-auditory late-visual early-visual
Intrinsic Dimension 30 29 41 13 15
Appendix Table 1: fMRI ROI data intrinsic dimensions