Understanding Brain Functional Dynamics Through Neural Koopman Operator with Control Mechanism

Abstract

One of the fundamental scientific problems in neuroscience is to have a good understanding of how cognition and behavior emerge from brain function. Since the neuroscience concept of cognitive control parallels the notion of system control in engineering, many computational models formulate the dynamics neural process into a dynamical system, where the hidden states of the complex neural system are modulated by energetic simulations. However, the human brain is a quintessential complex biological system. Current computation models either use neural networks to approximate the underlying dynamics, which makes it difficult to fully understand the system mechanics, or compromise to simplified linear models with very limited power to characterize non-linear and self-organized dynamics along with complex neural activities. To address this challenge, we devise an end-to-end deep model to identify the underlying brain dynamics based on Koopman operator theory, which allows us to model a complex non-linear system in an infinite-dimensional linear space. In the context of reverse engineering, we further propose a biology-inspired control module that adjusts the input (neural activity data) based on feedback to align brain dynamics with the underlying cognitive task. We have applied our deep model to predict cognitive states from a large scale of existing neuroimaging data by identifying the latent dynamic system of functional fluctuations. Promising results demonstrate the potential of establishing a system-level understanding of the intricate relationship between brain function and cognition through the landscape of explainable deep models.

I. Introduction

Human brain is a complex system physically wired by massive bundles of white matter fibers [1]. On top of intertwined structural connectomes, ubiquitous neuronal oscillations emerge remarkable functional fluctuations, synchronizing across large-scale neural circuits, that support myriad high-level cognitive functions necessary for everyday living [2]. With the prevalence of functional neuroimaging technology in many neuroscience studies, the idea of understanding the mind forms an important concept that the dynamic nature of the human brain cannot be understood by thinking of the system as comprised of independent components [3]. In the broader scope of computational neuroscience, the overarching goal is to elucidate the system mechanism to the extent that the computational model can identify the governing equation of the underlying complex system. In this regard, there is a critical need to establish a system-level understanding of how the brain switches through different mental states and how the self-organized system behavior of functional fluctuations supports the cognitive states.

Tremendous efforts have been made to model the evolution of neural activity observations (such as fMRI) towards the accurate prediction of brain states associated with underlying cognitive tasks and find a mechanistic explanation for the brain’s underlying control in processing information as well as making decisions [4]. Since the neuroscience notion of neural process is analogous to the dynamical system used in engineering, many existing works model the brain functional activity as a dynamical system [5], [6], where the hidden states of the complex neural system are modulated by (external) energetic simulations. In [6], a linear system is used to model the temporal trajectory of BOLD (blood-oxygen-level-dependent) signal extracted from fMRI, where the transition of brain states is restricted to the wiring pattern of neuronal fiber bundles. Nevertheless, the human brain represents a fundamentally complex biological system. Current linear models lack sufficient capability to capture the non-linear and self-organized dynamics, as well as the intricate neural activities present in this complexity. Numerous deep neural networks have been proposed to infer the neuron population dynamics [7], [8] by projecting the neural activity observations to a latent low-dimensional embedding space in a layer-by-layer manner. However, the ‘black box’ nature of the learning mechanism renders current deep models lacking explainability to achieve a biologically meaningful understanding of functional dynamics.

To address these limitations, we present an explainable deep model BRICK (BRain dynamics Identification using Control and Koopman operator) by resolving the following three challenges.

• Challenge 1: How to capture complex non-linear brain dynamics while maintaining mathematical tractability for systematic analysis? We tackle this challenge by employing Koopman operator theory [9], [10]. Koopman theory [9] has been used to discover the underlying dynamics of nonlinear complex systems through an infinite-dimensional linear operator acting on the space of all possible measurement functions of the system. Specifically, Koopman operator theory focuses on the functions of the states (a.k.a. “observables” or “measurements”) instead of the states themselves. It claimed that any sufficiently regular nonlinear autonomous dynamical system can be made linear under a high-dimensional non-linear blow-up of the state-space [11]. This guarantee allows us to design/learn a large number of “measurements” to fit the non-linear dynamics with linear form. Furthermore, we introduce a neural Koopman operator to linearize the complex nonlinear brain dynamics, which allows us to design an end-to-end deep model for large-scale functional neuroimages.

• Challenge 2: How to find a series of measurements that can form a Koopman invariant subspace? In theory, Koopman operator seeks to express the full behavior of dynamical systems by spanning an infinite-dimensional space of measurement functions [12]. However, it is computationally impractical to find infinite measurement functions. In this regard, it is a common practice to approximate Koopman operator using a subset of measurement functions. For example, many existing methods use basis functions such as sin, cos, polynomials and radial basis function to model non-linear systems [13], [14]. Considering the self-organized behavior of function connectivity along with evolving cognitive states, we propose a permutation-equivariant spatiotemporal neural network that generates region-trackable measurements to approximate the Koopman invariant subspace. Specifically, we first calculate the region-to-region functional connectivity from the BOLD signals, where the wiring patterns form the backbone of a dynamical system supporting the functional fluctuations. The functional connectivity captures the temporal dependencies and is then processed by multiple permutation-equivariant neural networks, which leverage the spatial structure of brain regions while preserving anatomical interpretability in the brain, to generate a set of biologically meaningful measurements.

• Challenge 3: How to design an interpretable control module to regulate the evolution of the dynamics? To guarantee that the learned measurements can form the Koopman invariant subspace, we propose a novel control module, which can generate a tasks-specific controller for unseen individual subjects based on their observations and cognitive states. With this controller, the system can correct the evolution of behavior according to the observations, reducing the accumulated error while offering neuroscience insight into cognitive control.

The outcome is an end-to-end deep model with great scalability to large-scale neuroimaging data, which allows us to advance our understanding of brain function through data-driven approaches. The major novelties and contributions are summarized as follows:

• We integrate the Koopman operator and control into an end-to-end variational inference model, which allows us to uncover the working mechanism of dynamical system of brain function.

• The latent brain dynamics identified by our method enable the modeling of the communication between regions and the discovery of how the particular brain regions regulate the evolution of self-organized functional fluctuations.

• The comprehensive comparison experiments on classic physics systems and real world fMRI datasets demonstrate that the model achieves state-of-the-art performance. We also conduct extensive exploratory experiments.

II. Related Work

This section summarizes existing literature on modeling dynamical systems through data-driven methods, with a focus on extensions to neural processes of human brain.

Dynamic system identification in machine learning aims to capture the temporal behaviors of complex systems. One of the classic approaches for system identification is extended dynamic mode decomposition (EDMD) [15]. EDMD introduces a dictionary of nonlinear functions to map the system states into a higher-dimensional space. This mapping allows for a better approximation of complex, nonlinear dynamics by leveraging the Koopman operator theory. Another method called sparse identification of nonlinear dynamics (SINDy) [16] offers a novel approach to discovering governing equations from data. Unlike traditional methods, SINDy leverages sparse regression techniques to identify the most relevant terms from a library of candidate functions, effectively balancing model complexity with accuracy. While these methods work well on many systems, their performance heavily depends on the choice of dictionary functions, requiring significant domain expertise and the quality and quantity of measurement data. Furthermore, the requirement for clean, noise-free measurements is often difficult to meet in real-world applications.

Recently, a number of deep learning models have been proposed to identify the complex system dynamics. For example, piecewise-linear recurrent neural networks (PLRNN) [17] present a specialized state-space modeling approach for understanding neural dynamics from neuroimaging data. Unlike traditional neural networks, PLRNN combines the interpretability of linear dynamical systems with the expressiveness of nonlinear models through a piecewise-linear formulation. The dendritically expanded piecewise-linear recurrent neural network (dendPLRNN) [18] enhances the original PLRNN framework by incorporating principles from dendritic computation through linear spline basis expansion. This extension maintains the interpretability and mathematical tractability of PLRNNs while significantly improving their capacity to approximate complex nonlinear dynamics in lower-dimensional spaces. However, the training of these models relies on a traditional sequential scheme and scheduled sampling [19], which has been reported to be difficult to converge.

In the field of neuroscience, there are many parallel works that focus on identifying the latent dynamics of populations of neurons. Latent factor analysis via dynamical systems (LFADS) [8] addresses the challenge of understanding neural population dynamics from single-trial recordings. This deep learning approach aims to infer underlying neural dynamics from noisy, incomplete neural spiking data without requiring trial-to-trial averaging. LFADS employs a variational auto-encoder framework with a dynamical systems core to reconstruct latent neural trajectories, making it particularly powerful for analyzing motor cortical activity. Meanwhile, Poisson latent neural differential equations (PLNDE) [7] approach offers a novel window to identifying latent dynamics in neural populations by combining neural ordinary differential equations with Poisson observations. This framework specifically addresses the challenges of modeling spike train data through a low-dimensional nonlinear dynamical system. PLNDE stands out for its ability to accurately infer phase portraits and fixed points from neural data, making it particularly valuable for analyzing decision-making tasks and other cognitive processes. Nevertheless, the latent space representations of these methods, while powerful for prediction, may not always be easily interpretable in a neurobiological perspective. Furthermore, current state-of-the-art methods are limited to neuron populations of specific brain regions rather than the entire brain network dynamics.

Besides these methods, recently some works considered using Koopman operator theory to identify the brain dynamics in a linear space. For example, Gallos et.al. [20] proposed a multi-stage framework to identify the fMRI dynamics, whose application is concurrent to our work. it first utilized the diffusion map for dimension reduction. Then either feedforward neural networks with geometric harmonics or the Koopman operator framework is applied to the low-dimensional embedding and finally solves the pre-image problem. Another work [21] focused on ECoG data in a mismatch negativity experiment using a stand-alone Koopman operator approach based on kernel methods without control components. Deep-GraphDMD [22] employs an auto-encoder architecture to learn Koopman eigenfunctions for graph data, effectively embedding non-linear network dynamics into a latent linear space. However, these methods aim to obtain a common dynamics across different subjects, which is difficult due to the highly variable individual differences.

In addition to Koopman Operator theory, several promising neural operators were proposed for system identification. Multi-wavelet neural operator [23], [24] utilizes wavelet and orthogonal polynomials to efficiently represent the operator for single and coupled PDE. Pade Exponential Operators [25] was proposed to solve the initial value problem. Most of these methods focused on directly learn the complex non-linear dynamics of a system.

III. Method

In this section, we present our learning-based method BRICK for identifying and understanding the dynamical system of functional brain networks. In III-A, we first briefly introduce Koopman operator theory and utilize it to linearize the complex brain dynamics in a higher dimension measurement space. Then in III-B, we introduce how to generate measurements that leverage the spatio-temporal information of brain functional fluctuations. In III-C, we extend the proposed system to a closed-loop feedback control system with a novel control module. Finally, we describe the objective function and the corresponding optimization procedure in III-D. The variable notations we use are listed in Table I.

TABLE I:

Notations and Definitions

Notation	Definition
${\mathit{x}}_{1:T}$	The BOLD signal sequence with length $T$ and $N$ brain regions, regarded as the states and observations of the original brain system
${\mathit{g}}_{1:T}$	The measurements functions with length $T$ and dimension $M$ generated from BOLD signals
$N$	The number of brain regions (ROIs)
$H$	The feature dimension of each node (ROI)
$M$	The dimension/number of measurements. $M=H\times N$
$\mathit{K}$	Koopman operator with shape $M\times M$
$\mathit{C}$	Control matrix with shape $M\times M$
$\mathit{A}$	The functional connectivity with shape $N\times N$
$p_{\theta }(\mathit{x})$	The unknown distribution of observations ${\mathit{x}}_{1:T}$
$p_{\theta }\left({\mathit{g}}_0\right)$	The prior distribution of initial conditions ${\mathit{g}}_0$, assumed as a Gaussian
$p_{\theta }(\mathit{u})$	The prior distribution of control inputs ${\mathit{u}}_{1:T}$, assumed as a Gaussian at each step
$p_{\theta }\left({\mathit{x}}_{1:T}\mid {\mathit{u}}_{1:T},{\mathit{g}}_0\right)$	The posterior distribution of observations
$q_{\psi }\left({\mathit{g}}_0\mid {\mathit{x}}_{1:T}\right)$	The posterior distribution of initial conditions ${\mathit{g}}_0$
$q_{\psi }\left({\mathit{u}}_{1:T}\mid {\mathit{x}}_{1:T}\right)$	The posterior distribution of control inputs ${\mathit{u}}_{1:T}$

A. Linearizing Brain Dynamics with Koopman Operator

We formulate the brain network as a non-linear discrete-time dynamical system ${\mathit{x}}_{t+1}=\mathit{f}\left({\mathit{x}}_t\right)$, where $\mathit{f}:\mathcal{X}\to \mathcal{X}$ represents the state transition function mapping the state ${\mathit{x}}_t$ at time $t$ to its subsequent state ${\mathit{x}}_{t+1}$. The states ${\mathit{x}}_{1:T}\in {ℝ}^{N\times T}$ comprise the temporal evolution of BOLD signal snapshots across $N$ brain regions. The fundamental principle underlying our methodology leverages Koopman operator theory to transform the non-linear brain state dynamics into a linear evolution in a higher-dimensional measurement space. Specifically, Koopman operator theory establishes that a sufficiently regular nonlinear dynamical system admits representation as an infinite-dimensional linear system [9], [10], [26]. Formally, for any measurement function $g:\mathcal{X}\to ℂ$, the Koopman operator $\mathcal{K}$ is defined by:

$$\mathcal{K}g(\mathit{x})=g(\mathit{f}(\mathit{x}))$$ (1)

However, finding such an infinite dimensional $\mathcal{K}$ is computationally impractical. Therefore, an alternative solution to find a Koopman invariant subspace $\mathcal{G}$, s.t. $\forall g\in \mathcal{G}$, $\mathcal{K}g\in \mathcal{G}$. Then if the Koopman invariant subspace $\mathcal{G}$ is spanned by a finite number of basis functions $\mathit{g}=\left\{g_1,\mathrm{..},g_M\right\}$, we can use a restricted Koopman operator $\mathit{K}\in {ℝ}^{M\times M}$ to approximate the original infinite dimensional Koopman operator $\mathcal{K}$. Then we can get the approximated underlying brain network dynamics:

$$\begin{gathered} {\mathit{g}}_{t+1}=\mathit{K}{\mathit{g}}_t \\ {\widehat{\mathit{x}}}_{t+1}={\mathit{W}}_{\mathit{x}}{\mathit{g}}_{t+1} \end{gathered}$$ (2)

where the subscript $t$ means the time point and ${\mathit{W}}_{\mathit{x}}$ transform the measurements to the original observations. Nevertheless, it is impossible to determine a finite dimensional Koopman invariant subspace that includes the original state variables for any system with multiple fixed points or any more general attractors [26]. Therefore, in the next two subsections, we aim to learn a set of measurement functions $\mathit{g}$ as precise as possible (in III-B) and design an interpretable control module (in III-C) to reduce the accumulated error caused by the gap between true Koopman operator $\mathcal{K}$ and the approximated counterpart $\mathit{K}$.

B. Generating Initial Condition by Spatio-Temporal Measurements

In this section, we propose a method to generate basis measurement functions $\mathit{g}$ while simultaneously inferring the initial conditions of the dynamical system.

1). Functional Connectivity:

Brunton et al. [27] demonstrated that a Koopman invariant subspace can be approximated using a finite set of time-delay measurements. This concept aligns with brain functional connectivity principles, defined as the temporal dependency of neuronal activation patterns across anatomically separated brain regions [28]. Therefore, we begin by calculating the functional connectivity of the given data. Specifically, given a series of BOLD signal snapshots ${\mathit{x}}_{1:T}\in {ℝ}^{N\times T}$, where $N$ is the number of brain regions and $T$ denotes the length of the BOLD signal sequence (same as the states ${\mathit{x}}_{1:T}$ in III-A), we compute the Pearson correlation [29] between each pair of nodes. This yields the functional connectivity (FC) matrix $\mathit{A}\in {ℝ}^{N\times N}$, where each entry ${\mathit{A}}_{ij}$ is the Pearson correlation between node $i$ and $j$. Subsequently, we apply a row-wise MLP to reduce the dimensionality from $N$ to $H$, producing a temporal intermediate result ${\mathit{Z}}_0\in {ℝ}^{N\times H}$, where $H\ll N$.

2). Coupling with Spatio Information:

While the intermediate state ${\mathit{Z}}_0$ captures functional information, it does not account for the brain network’s spatial information. Therefore, we employ a neural network to model inter-node interactions and obtain the final spatio-temporal measurements. It is crucial to maintain awareness of region-to-region communications in the underlying dynamic state (i.e., the measurement space). To achieve this, the spatio encoder network $\text{Φ}$ (green box in Fig. 1) should be permutation-equivariant, meaning that $\forall \sigma \in {\mathit{S}}_N$, $\text{Φ}\left(\sigma \circ {\mathit{Z}}_0\right)=\sigma \circ \text{Φ}\left({\mathit{Z}}_0\right)$, where ${\mathit{S}}_N$ is the permutation group. This spatial encoder $\text{Φ}$ enables us to maintain correspondence between brain regions and their latent measurements. For fair comparison and ease of application, we adopt TransformerEncoder [30] as our encoder $\text{Φ}$, which provides a general permutation-equivariant architecture. Alternative models such as MPNN [31] or GCN [32], which specialize in graph structures, could also be employed. Finally, rather than directly using the output as the initial condition measurement, we consider the initial condition is drawn from a prior distribution ${\mathit{g}}_0\sim p_{\theta }\left({\mathit{g}}_0\right)$, $p_{\theta }\left({\mathit{g}}_0\right)=\mathcal{N}(0,ϵ\mathit{I})$, accounting for the stochastic nature and complexity of brain functional fluctuation dynamics. The approximate posterior distribution of ${\mathit{g}}_0$ is $q_{\psi }\left({\mathit{g}}_0\right|{\mathit{x}}_{1:T})$, a Gaussian with mean and covariance determined by two separate TransformerEncoder:

$${\mu }_{{\mathit{g}}_0}={\text{Φ}}_{\mu }\left({\mathit{Z}}_0\right)$$ (3)

$$\log {\text{Σ}}_{{\mathit{g}}_0}^2={\text{Φ}}_{\text{Σ}}\left({\mathit{Z}}_0\right)$$ (4)

$${\text{Σ}}_{{\mathit{g}}_0}=\exp \left(\frac{1}{2}\log {\text{Σ}}_{{\mathit{g}}_0}^2\right)$$ (5)

$${\mathit{g}}_0\sim \mathcal{N}\left({\mathit{g}}_0\mid {\mu }_{{\mathit{g}}_0},{\text{Σ}}_{{\mathit{g}}_0}\right)\in {ℝ}^{N\times H}$$ (6)

$${\mathit{g}}_0=\text{vec}\left({\mathit{g}}_0\right)\in {ℝ}^M$$ (7)

where $\text{vec}(\cdot )$ is the vectorization function and $M=N\times H$ is the number (or dimension) of measurements.

Fig. 1:

Overview of our deep model BRICK. In general, our model consists of three components. (a) Data processing. The input to our model includes the mean time course of BOLD signal x and functional connectivity matrix $\mathit{A}$. (b) System identification. The learning process for identifying the governing equation of the underlying functional fluctuations involves a neural Koopman operator for system linearization and a control module for system correction. These components work together collaboratively to uncover the system’s mechanism from the observed measurements. (c) Driving force of system identification. The objective function represents the overarching goal of accurately predicting future behavior and reconstructing the input based on the latent states.

C. Task-Related Control Module

In this section, we introduce a novel task-related control module, extending the proposed autonomous system to a closed-loop feedback control system, which allows the system receiving the observations and producing corresponding control signal to regulate the evolution of the dynamics. Specifically, we modify the system in Eq. 2 to the following form:

$${\mathit{g}}_{t+1}=\mathit{K}{\mathit{g}}_t+\mathit{C}{\mathit{u}}_{t+1}$$ (8)

where $\mathit{C}\in {ℝ}^{M\times M}$ is the control matrix and can be regarded as part of the system while ${\mathit{u}}_{1:T}$ is the external control inputs.

Following the autonomous dynamics captured by the Koopman operator $\mathit{K}$, which represents the shared, intrinsic brain dynamics across individuals, our BRICK framework addresses individual variability by proposing that the control matrix $\mathit{C}$ should be individual-specific and task-related. While we parametrize a common Koopman operator $\mathit{K}$ to maintain consistency of underlying neural dynamics across subjects, the personalization of the control module $\mathit{C}$ reflects the fact that different individuals may require distinct control strategies based on task-related state, even when their fundamental neural dynamics follow similar patterns.

To generate the individual-specific control matrix $\mathit{C}$, we employ an encoder-based architecture that learns from functional observations. Specifically, we first encode the observations ${\mathit{x}}_{1:T}$ into a latent embedding $\mathit{E}=\text{Encoder}\left({\mathit{x}}_{1:T}\right)\in {ℝ}^{M\times T}$ (shown at the top of Fig. 1). This embedding serves dual purposes: (1) it is used to generate the control matrix $\mathit{C}={\psi }_{\theta }(\mathit{e})$ and control input ${\mathit{u}}_{1:T}$; (2) it simultaneously feeds into a classifier that predicts the task state, i.e., $\widehat{\mathit{s}}=\text{Classifier}(\mathit{e})$, where $\mathit{e}$ is obtained by applying an average pooling layer to $E$ along the time dimension. This architecture ensures that the learned control matrices are grounded in functional activity patterns while maintaining task relevance, as verified by the classification objective. The Encoder we use here is a Transformer Encoder [30] followed by a MLP to adjust the dimension from $N$ to $M$. Notably, we specifically use a simple linear layer as the Classifier. This design choice ensures that the task-relevant information is primarily encoded in the embedding space rather than being extracted by a sophisticated classifier. As for the control matrix generation, we simplify the control matrix to a diagonal matrix to prevent over-fitting and reduce the computational overhead. Together, we have $\text{diag}(\mathit{C})=\tanh (\mathit{W}\mathit{e}+\mathit{b})$.

Next we use a variational inference scheme to generate feature-wise control input signal according to the observations. We assume the control inputs are govern by a Gaussian prior $p_{\theta }\left({\mathit{u}}_t\right)=\mathcal{N}(0,\mathit{I})$, then we use a posterior $q_{\varphi }\left({\mathit{u}}_t\mid {\mathit{x}}_{1:t}\right)$ to approximate this module by:

$$\begin{array}{l}\mathit{E}=\text{Transformer}\left({\mathit{x}}_{1:T}\right)\\ {\mu }_{{\mathit{u}}_{1:T}}=\text{ML}{\text{P}}_{\mu }(\mathit{E}),\\ \log {\text{Σ}}_{{\mathit{u}}_{1:T}}^2=\text{ML}{\text{P}}_{\text{Σ}}(\mathit{E})\\ {\text{Σ}}_{{\mathit{u}}_{1:T}}=\exp \left(\frac{1}{2}\log {\text{Σ}}_{{\mathit{u}}_{1:T}}^2\right)\\ q_{\varphi }\left({\mathit{u}}_t\mid {\mathit{x}}_{1:t}\right)=\mathcal{N}\left({\mathit{u}}_t\mid {\mu }_{{\mathit{u}}_t},{\text{Σ}}_{{\mathit{u}}_t}\right)\\ q_{\varphi }\left({\mathit{u}}_{1:T}\mid {\mathit{x}}_{1:T}\right)=\prod _{t=1}^T{q}_{\varphi }\left({\mathit{u}}_t\mid {\mathit{x}}_{1:t}\right)\end{array}$$ (9)

D. Training and Optimization

So far, we have modeled the brain network dynamics with following components: (1) The initial condition ${\mathit{g}}_0$; (2) The dynamical system ${\mathit{g}}_{t+1}=\mathit{K}{\mathit{g}}_t+\mathit{C}{\mathit{u}}_{t+1}$, and (3) the inferred control input ${\mathit{u}}_{1:T}$. We now introduce the training scheme and optimization objective. During the training stage, the complete sequence of observations ${\mathit{x}}_{1:T}$ is available, allowing us to leverage the parallel processing capabilities of the Transformer architecture by computing the embeddings and control signals for all time steps simultaneously, rather than sequentially. Then, as for the system dynamics ${\mathit{g}}_{t+1}=\mathit{K}{\mathit{g}}_t+\mathit{C}{\mathit{u}}_{t+1}$, we adopt the parallel scan algorithm and the eigenvalue initialization method introduced in [11].

The objective function consists of several components: (1) the embedding classification loss ${\mathcal{L}}_{cls}=\text{CrossEntropy}(\mathit{s},\widehat{\mathit{s}})$, where $\mathit{s}$ and $\widehat{\mathit{s}}$ are the true task-states and the one predicted by the embedding $\mathit{e}$ respectively. (2) We want to maximize the log-likelihood of the given observations. For the compactness and readability, we replace ${\mathit{x}}_{1:T}$, ${\mathit{g}}_{1:T}$, ${\mathit{u}}_{1:T}$ with $\mathit{x}$, $\mathit{g}$, $\mathit{u}$

$$\begin{gathered} \log p_{\theta }(\mathit{x})\ge {\mathcal{L}}_{ELBO}={\mathcal{L}}_{likelihood}-{\mathcal{L}}_{\mathbb{K}\mathbb{L}} \\ {\mathcal{L}}_{likelihood}={\mathbb{E}}_{q_{\psi }(\mathit{u}\mid \mathit{x}),q_{\psi }\left({\mathit{g}}_0\mid \mathit{x}\right)}\left(\sum _{t=1}^T\log \left(\left(\log p_{\theta }\left({\mathit{x}}_t\mid {\mathit{u}}_t,{\mathit{g}}_0\right)\right)\right)\right) \\ {\mathcal{L}}_{\mathbb{K}\mathbb{L}}={\mathcal{D}}_{\mathbb{K}\mathbb{L}}\left(\mathcal{N}\left({\mathit{g}}_0\mid {\mu }_{{\mathit{g}}_0},{\text{Σ}}_{{\mathit{g}}_0}\right)\Vert p_{\theta }\left({\mathit{g}}_0\right)\right) \\ +{\mathcal{D}}_{\mathbb{K}\mathbb{L}}\left(q_{\psi }(\mathit{u}\mid \mathit{x})\Vert p_{\theta }(\mathit{u})\right) \end{gathered}$$ (10)

where ${\mathcal{L}}_{likelihood}$ is the log-likelihood of the reconstructed $\widehat{x}$ in comparison with the observed data $x$ under Gaussian noise with variance $\lambda$, while ${\mathcal{L}}_{\mathbb{K}\mathbb{L}}$ represents the KL divergence penalties that regularize both the initial latent state $g_0$ and the inferred inputs $u$ to stay close to their respective prior distributions.

Taken together, we formulate the optimization objective as:

$$\mathcal{L}=\beta {\mathcal{L}}_{cls}+(1-\beta ){\mathcal{L}}_{ELBO}$$ (11)

where $\beta$ is a balance hyperparameter.

IV. Experiments

To evaluate the power of our model BRICK for the system identification task, we fit our model on a synthetic dataset, which is a nonlinear dynamical system. We generated the trajectories of the system by (1) randomly selecting different initial values and system parameters. (2) randomly sampling trajectories with fixed lengths in the systems. Next, we conduct experiments on the real datasets HCP-A and ABIDE, where a detailed description will be provided in the corresponding section. We split the data into training, validation, and test sets with a ratio of 7:1:2. We use Grid Search on the training set and validation set to find the optimal hyperparameters combination and finally evaluate the models on the test set. The baselines include traditional data-driven system identification methods and deep learning methods. All experiments are conducted on a NVIDIA RTX 6000 Ada GPU with 48 GB memory. Here are the details:

• EDMD [15]. The main idea of extended dynamics mode decomposition is to find a finite-dimensional linear map $\mathbf{K}$ whose spectral properties approximate those of the true Koopman operator $\mathcal{K}$. This can be achieved by selecting a set of measurements and solving a corresponding least squares problem.

• SINDy [16]. The sparse identification of nonlinear dynamics (SINDy) algorithm seeks to approximate the dynamics function $\frac{dx}{dt}=f(x)$ with a generalized linear model whose basis is polynomials or triangle functions. By solving a sparse session problem, only a few bases are active in the dynamical system.

• Transformer [30] Transformer is a deep model good at learning long range dependencies benefit from the self-attention mechanism and has been successfully applied to dynamical system modeling [33].

• PLRNN [17] Piecewise-linear recurrent neural network (PLRNN) is a nonlinear state space model RNN composed of rectified-linear units. PLRNN is one of the state-of-the-art model specialized for dynamics reconstruction and system identification.

• dendPLRNN [18] Dendritic-PLRNN is an augmentation of PLRNN inspired by dendritic processing and is implemented based on spline basis expansion.

• LRU [11] Linear Recurrent Unit (LRU) is a linear RNN which can efficiently handle long sequence data and has strong performance on sequence tasks.

• DLinear [34] DLinear is a simple yet very effective baseline for time series forecasting which consists of just linear transformations along the time dimension.

A. Lorenz System

Lorenz system is notable for having chaotic solutions for certain parameter values and initial conditions, especially the butterfly-shaped Lorenz attractor. It’s a good example showing that chaotic systems can be completely deterministic and yet still be inherently impractical or even impossible to predict over longer periods. A Lorenz-63 system ¹ has the following form:

$$\left\{\begin{array}{l}\frac{dx}{dt}=\sigma (y-x)\\ \frac{dy}{dt}=x(\rho -z)-y\\ \frac{dz}{dt}=xy-\beta z\end{array}\right.$$ (12)

Here we use the classical setting of $\rho =28$, $\sigma =10$, $\beta =\frac{8}{3}$ with Gaussian noise $ϵ\sim \mathcal{N}\left(\mathbf{0},{0.01}^2\times \mathbf{I}\right)$. To test if the models can properly reconstruct the dynamics of the system, we report the mean square error at the first 64, 128, and 200 steps separately repeated 100 times with different dataset divisions and model initializations. Since the dimension of Lorenz system is not high (three dimensions), the hidden dimension is all set to 32. For PLRNN and dendPLRNN, the teacher force interval is set to 25 and the number of bases of dendPLRNN is 20 same as the original paper. The threshold of SINDy is set to 0.05. Since the Lorenz system doesn’t relate to change of dynamics with different states, we take the task-embedding component out of our BRICK model in this synthetic dataset.

Fig. 2 shows the predicted trajectories of different methods on the same initial condition with 1000 prediction steps. As shown in Fig. 2, most of the methods reconstruct the Lorenz-63 system reasonably except EDMD. Specifically, BRICK produces the most faithful reconstruction, showing a clear butterfly-like structure with well-defined lobes and smooth trajectories, matching closely with typical Lorenz attractor characteristics, revealing that it is possible to recover the system dynamics by combining Koopman operator theory with control input. DendPLRNN and PLRNN show similar overall structure but with some distortions - their trajectories are less smooth and the lobes appear slightly warped compared to the expected Lorenz pattern. The Transformer method generates dense trajectories but with substantial noise and irregularities in the attractor structure because it highly relies on precise observation. SINDy and EDMD show less faithful results than other deep models.

Fig. 2:

Comparison of prediction methods on the Lorenz-63 system over 1,000 time steps. Each method started from the same initial condition, with the blue line showing the ground truth trajectory and colored points indicating predictions from different methods.

Table II displays the mean square error (MSE) of different prediction steps for each method. Our BRICK method achieves 2^nd place in short-range prediction, 1^st place in mid-range prediction, and 2^nd place in long-range prediction, respectively. It seems that SINDy shows similar performance (3^rd place in short-range prediction, 2^nd place in mid-range prediction, and 1^st in long-range prediction) in terms of MSE. However, our method uncovers the attractor structure with significantly greater accuracy compared to SINDy.

TABLE II:

Mean squared error (MSE) comparison of different prediction methods on the Lorenz-63 system. Results are shown for predictions over 64, 128, and 200 time steps, reported as mean ± standard deviation.

	MSE
Model	64 steps	128 steps	200 steps
EDMD	2.247±0.150	1.363±0.117	1.138±0.107
SINDy	0.054±0.020	0.101±0.032	0.202±0.045
Transformer	0.090±0.030	0.321±0.057	0.438±0.066
PLRNN	0.056±0.008	0.367±0.019	0.666±0.026
dendPLRNN	0.021±0.005	0.190±0.014	0.430±0.021
Ours (BRICK)	0.042±0.006	0.082±0.010	0.388±0.020

B. Performance Evaluation on fMRI Dataset

1). Data Description:

(1) The Autism Brain Imaging Data Exchange (ABIDE) [35] provides anonymous resting-state functional magnetic resonance imaging (rs-fMRI) data collected from 17 international sites. Our analysis used brain networks and BOLD signal from 1009 subjects, including 516 Autism spectrum disorder (ASD) patients (51.14% positive cases). Brain regions were defined using the Craddock 200 atlas [36]. (2) Human Connectome Project-Aging (HCP-A) dataset [37] includes data from 717 subjects, encompassing both fMRI (4,846 time series) and Diffusion Weighted Imaging (DWI) (717) scans. This rich collection facilitates in-depth analyses of both functional and structural connectivity. The HCP-A dataset includes data from four brain tasks associated with memory: VISMOTOR, CARIT, FACENAME, and resting state. Each fMRI scan consists of 125 time points. In the following experiments on HCP-A, these tasks are treated as distinct categories in a four-class classification problem. We partition each into 116 regions using AAL atlas [38]. Thus, structural connectivity is a 116 × 116 matrix where each element is quantified by the number of nerve fibers linking two brain regions [39].

2). Evaluation Metrics:

The metrics of this experiment are as follows : (1) Mean Square Error (MSE) of 10 steps predicted BOLD signal. (2) The approximated Kullback-Leibler divergence $\left(D_{\mathbb{K}\mathbb{L}}\right)$ between true distribution $p_{\theta }(\mathit{x})$ and posterior distribution $q_{\psi }\left(\mathit{x}\mid {\mathit{g}}_0,\mathit{u}\right)$ (See Appendix 6.2 in [18]). KL divergence evaluates how well the model captures the underlying dynamics of brain states, accounting for natural variability in neural activity patterns. (3) The root mean square error (RMSE) between the reconstructed functional connectivities and the ground truth. The resulting RMSE measures preservation of network-level relationships between brain regions, which is crucial for understanding brain organization and function.

3). Implementations:

For experiments, we use an AdamW [40] optimizer with CosineAnnealing Schedule [41]. We use grid search to select the hyperparameters. The number of transformer layers is searched within [1, 2, 4] and found 2 layers provided the best performance. We test learning rate in [1e-4, 5e-4, 1e-3, 5e-3] with with 1e-3 providing fastest stable convergence. The batch size is 512. All models are trained for 1000 epochs, and the epoch with the highest performance on the validation set is used for performance comparison on the test set.

4). Comparison with Baselines:

Table III shows a comprehensive evaluation for the effectiveness of brain dynamics modeling approaches. First notice that the performance on ABIDE is worse than that of HCP-A because the fMRI images in ABIDE are collected from multiple centres while HCP-A data are collected in a highly standardized way. Traditional methods like EDMD and SINDy, while mathematically elegant in their approach to approximating dynamical systems through linear mappings and sparse identification, show limited success in capturing the complex, nonlinear nature of brain dynamics. One should notice that SINDy shows the weakest performance because it’s intractable to use polynomial features for such high dimensional systems (e.g., with the highest order of 3, the number of features scales up to the level of 10⁶). Moreover, SINDy is sensitive to the noise of states, making it less accurate.

TABLE III:

Performance Comparison Across Methods and Datasets (IC means Initial Condition module, C means Control module)

Method	ABIDE			HCP-A
	MSE	KL	RMSE	MSE	KL	RMSE
EDMD	0.841 ± 0.020	85.976 ± 0.383	0.319 ± 0.003	0.630 ± 0.015	53.715 ± 0.147	0.543 ± 0.001
SINDy	1.735 ± 0.036	152.172 ± 1.723	0.361 ± 0.003	1.097 ± 0.019	102.748 ± 0.917	0.574 ± 0.003
PLRNN	0.994 ± 0.015	65.330 ± 0.613	0.404 ± 0.012	0.673 ± 0.013	23.358 ± 0.346	0.202 ± 0.006
dendPLRNN	1.006 ± 0.008	63.329 ± 2.056	0.455 ± 0.016	0.654 ± 0.008	22.860 ± 0.315	0.198 ± 0.004
Transformer	0.477 ± 0.005	30.924 ± 0.538	0.193 ± 0.003	0.179 ± 0.002	5.170 ± 0.103	0.079 ± 0.001
Dlinear	0.271 ± 0.015	23.012 ± 0.338	0.119 ± 0.006	0.209 ± 0.008	6.721 ± 0.108	0.107 ± 0.004
LRU	0.335 ± 0.013	26.824 ± 0.448	0.133 ± 0.003	0.353 ± 0.012	7.090 ± 0.114	0.098 ± 0.005
BRICK(w/o C w/o IC)	0.428 ± 0.006	30.164 ± 0.609	0.145 ± 0.004	0.159 ± 0.009	5.278 ± 0.278	0.080 ± 0.004
BRICK(w/o IC)	0.428 ± 0.012	29.320 ± 0.846	0.147 ± 0.002	0.083 ± 0.004	4.298 ± 0.167	0.062 ± 0.003
BRICK(w/o C)	0.309 ± 0.008	25.788 ± 0.612	0.120 ± 0.005	0.153 ± 0.013	4.814 ± 0.232	0.072 ± 0.004
BRICK	0.168 ± 0.013	14.118 ± 0.996	0.081 ± 0.005	0.077 ± 0.001	4.637 ± 0.054	0.073 ± 0.001

The neural network-based approaches (PLRNN and dend-PLRNN) demonstrate improved performance, likely due to their ability to model nonlinear relationships based on a nonlinear state space model. However, these methods still face challenges in capturing long-range temporal dependencies in brain signals. Compared to their superior performance in Lorenz system, they struggle in the complex brain dynamics identification tasks. One possible reason is the limited number of parameters, leading to under-fitting. The significant performance leap achieved by the Transformer architecture highlights the importance of attention mechanisms in modelling brain dynamics, as they can effectively capture both local and global temporal relationships in neural activity patterns. Though being simple, the linear methods, DLinear and LRU, also showed consistent comparable performance on both datasets.

In contrast, our BRICK method achieves best performance, particularly in its full setting, suggests that combining controlled dynamics modeling with spatio-temporal measurements is crucial for accurately capturing brain state transitions. The superior performance of BRICKis evident not only in the aggregate metrics (Table III) but also in its ability to sustain prediction accuracy over time (Fig. 3) and reconstruct intricate connectivity patterns (Fig. 4). The prediction performance by SINDy and EDMD declines rapidly with increasing time steps, indicating their limitations in handling high-dimensional, nonlinear dynamics noted in the original analysis.

Fig. 3:

Comparison of prediction errors across different methods on the ABIDE and HCP-A datasets. The plots show Mean Squared Error (MSE) over increasing prediction steps, with lower values indicating better performance. Each color represents a different prediction method.

Fig. 4:

Comparison of reconstructed functional connectivity matrices across different methods. The top row shows the ground truth connectivity pattern alongside results from EDMD, DLinear and PLRNN methods, while the bottom row displays reconstructions from our proposed method, Transformer, LRU and dendPLRNN. The color scale ranges from −0.6 (dark blue) to 1.0 (dark red), representing correlation strengths.

The connectivity matrices provide additional insight into why neural network-based approaches perform better - they’re able to capture the underlying structure of brain networks, though with varying degrees of accuracy. BRICK success in both temporal prediction stability and connectivity reconstruction suggests that its architecture effectively combines the benefits of controlled dynamics modeling with the ability to capture complex spatio-temporal patterns in brain activity.

5). Ablation Study:

Table III also shows the performance of BRICKin different ablation settings by taking out the control module (denoted by w/o C) and initial condition (denoted by w/o IC). The base dynamics prototype (BRICK w/o C & w/o IC) establishes a solid foundation with its Koopman operator approach and fixed control module across subjects and tasks, showing initial improvements over existing counterpart methods and similar performance with Transformer. On top of this, the control module (BRICK w/o IC) demonstrates particular strength in task-related scenarios, as evidenced by its dramatic improvement of KL divergence from 5.278 to 4.298 on the HCP-A dataset, while showing minimal impact on resting-state ABIDE data. Since the setup of initial conditions provides sophisticated spatio-temporal measurements and functional connectivity modeling, our full method without the control module (BRICK w/o C) still yields promising results for resting-state data, where BRICK w/o C has reduced ABIDE’s KL divergence from 30.164 to 25.788, while having a slightly negative impact on task-based HCP-A data.

The contribution of each learning module in BRICK can be explained by examining the fundamental nature of each dataset and the modules’ design principles. The control module’s limited contribution on ABIDE’s resting-state data stems from the spontaneous, unstructured nature of brain activity during rest, where there are no clear task-related state transitions to regulate or control, making the task-based control module and classifier component less effective. Conversely, the initial condition module’s reduced effectiveness on HCP-A’s task-based data is because task-driven brain states are already highly structured and constrained by task requirements, making additional modeling of functional connectivity and spatial relationships less crucial, and potentially introducing unnecessary complexity through stochastic initial condition modelling. This complementary pattern underscores a key principle in brain dynamics modeling: the importance of aligning modeling components with the underlying characteristics of the neural activity being studied, regardless of whether it is a more structured, task-driven nature of HCP-A data or the more variable, spontaneous nature of resting-state ABIDE data.

Finally, the full BRICK setting achieves the best overall performance by leveraging this complementarity - the control module’s ability to handle task-specific variations combines synergistically with the initial condition module’s capacity to capture underlying brain state representations, resulting in optimal performance on ABIDE dataset and the second-best performance of HCP-A. This pattern suggests that comprehensive brain dynamics modeling requires both task-specific control mechanisms and sophisticated state representation capabilities, with different components becoming more or less crucial depending on the nature of the neural activity being analyzed.

6). Parameter Sensitivity Analysis:

We also conduct experiments on the hyperparameter sensitivity on the latent dimension and the loss balancing parameter $\beta$.

As shown in Fig. 6 (a)-(b), we evaluated model performance by varying the latent dimension multiplier from 1× to 3× the original dimension. For the ABIDE dataset, increasing latent dimensions initially improves performance across all metrics. In contrast, the HCP-A dataset demonstrates more varied responses to increased latent dimensions. All metrics exhibit a U-shaped curve with optimal performance at 2×. These findings indicate that different brain dynamics tasks may require specific optimal dimensionality settings, with task-related data (HCP-A) showing greater sensitivity to dimensional tuning than resting-state data (ABIDE). Overall, the 2× multiplier represents a balanced configuration that achieves strong performance across both datasets and all evaluation metrics.

Fig. 6:

Hyperparameter sensitivity evaluation. (a) The effect of latent dimension on ABIDE dataset. (b) The effect of latent dimension on HCP-A dataset. (c) The effect of loss balance parameter $\beta$ on ABIDE dataset. (d) The effect of loss balance parameter $\beta$ on HCP-A dataset. Each experiment is repeated with 10 different random seeds.

The parameter sensitivity analysis for the balance factor $\beta$ reveals robust performance across a wide range of values for both datasets. As shown in Fig. 6 (c)-(d), we evaluated the impact of $\beta$ on model performance by varying its value from 0.1 to 0.9. All metrics demonstrate remarkable stability across the entire range of $\beta$ values. This stability suggests that our model architecture effectively balances the classification and reconstruction objectives without requiring precise tuning of the $\beta$ hyperparameter.

C. Interpretation

In this part, we aim to interpret the learned system, especially the control module. We first randomly draw 20 samples from different tasks separately. Next we sum up and normalize the absolute value of each row of the control matrix, which measures the intensity of the response generated by the region given the same input $\mathit{u}$. Then we map and visualize the top weighted region of different tasks.

As shown in Fig. 5, we can identify top-10 critical brain regions associated with each task, aligning well with their respective cognitive functions. For instance, the CARIT task demonstrates strong activation in the temporal region and prefrontal cortex, indicating integrated processing between auditory processing, semantic memory, and executive control functions. The FACENAME task shows significant activation in the prefrontal cortex and temporal lobe, reflecting its involvement in memory encoding, retrieval, and facial recognition processes. The VISMOTOR task predominantly activates the motor regions and visual regions, emphasizing its role in visuomotor integration. The consistent activation patterns across subjects demonstrate both the specificity of regional recruitment for different cognitive demands and the integration of distributed networks for complex task execution. These findings support the functional specialization of brain networks for specific tasks and provide a foundation for further investigation into the interactions between these regions.

Fig. 5:

Critical brain regions (top-10) linked to three brain tasks: CARIT (left), FACENAME (middle), and VISMOTOR (right).

V. Discussion

In this section, we aim to discuss both the potential limitations of the current method as well as the potential applications.

One of the major considerations of our approach is the fitness of selected/learnt measurements. The quality of linearization depends critically on choosing appropriate observable functions, which is often problem-specific and non-trivial. Here we use the biologically inspired measurements and can only validate them empirically with no strong mathematical guarantee. Another potential limitation is that finite-dimensional approximations may not “close” (i.e., remain within the Koopman invariant subspace). With measurements that can not form a Koopman invariant subspace, the prediction error will grow fast in long-term prediction. For both of the limitations, we use the proposed control module to regulate to avoid quick divergence. However, this method has no rigorous guarantee to eliminate the limitations.

We further discuss the potential applications of BRICK. First, the current task-specific control module could be adapted to characterize disorder-specific deviations in brain dynamics, potentially serving as diagnostic biomarkers for conditions like Alzheimer’s disease, schizophrenia, or autism spectrum disorder. Second, our framework could be extended to create “disease fingerprints” by analyzing how the Koopman operator and control matrices differ between healthy controls and patient populations. Specifically, we can learn two Koopman matrices ${\mathbf{K}}_{nc}$, ${\mathbf{K}}_{patient}$ from healthy controls and patient populations to study their properties like eigenvalue distribution.

VI. Conclusion

This work presents BRICK (BRain dynamics Identification using Control and Koopman), a novel deep learning framework for understanding brain functional dynamics through the lens of controlled dynamical systems. Our approach addresses several fundamental challenges in modeling complex brain dynamics: First, we leverage Koopman operator theory to transform non-linear brain dynamics into a linear evolution in a higher-dimensional measurement space, providing both mathematical tractability and the ability to capture complex nonlinear behaviors. This linearization enables systematic analysis while preserving the rich dynamics of neural processes. Second, we introduce a sophisticated measurement generation scheme that combines functional connectivity patterns with spatial information through permutation-equivariant neural networks. This approach generates biologically meaningful measurements that form an effective Koopman invariant subspace, capturing both local and global patterns in brain activity. Third, we develop a task-related control module that adapts to individual differences while maintaining consistent underlying dynamics across subjects. This innovation allows the model to account for both shared neural mechanisms and individual variations in brain function.

Our comprehensive experiments on both synthetic data (Lorenz system) and real fMRI datasets (ABIDE and HCP-A) demonstrate BRICK’s superior performance compared to existing methods. The ablation studies reveal the complementary nature of our model’s components, with the control module proving particularly effective for task-related dynamics and the initial condition module excelling in capturing resting-state patterns. This suggests that comprehensive brain dynamics modeling requires both task-specific control mechanisms and sophisticated state representation capabilities.

Appendix

A. Implementation Details about BRICK

1). Neural Koopman Operator:

During training, we do not directly optimize the Koopman operator $\mathit{K}$. Instead, we utilize the eigendecomposition to enhance the training efficiency. Specifically, we first denote ${\tilde{\mathit{u}}}_0={\mathit{g}}_0$ and ${\tilde{\mathit{u}}}_t=C{\mathit{u}}_t$ for simplicity. Then we can notice that ${\mathit{g}}_{t+1}={\sum }_{i=0}^{t+1}{\mathit{K}}^i{\tilde{\mathit{u}}}_{t+1-i}$. By eigendecomposition $\mathit{K}=\mathit{P}\mathbf{\Lambda }{\mathit{P}}^{-1}$, where $\mathit{P}$ and $\mathbf{\Lambda }$ are complex. By left multiplying ${\mathit{P}}^{-1}$ on both sides, we can get ${\overline{\mathit{g}}}_{t+1}={\sum }_{i=0}^{t+1}{\mathbf{\Lambda }}^i\odot {\overline{\mathit{u}}}_{t+1-i}$, where ${\overline{\mathit{g}}}_t={\mathit{P}}^{-1}{\mathit{g}}_t$ and ${\overline{\mathit{u}}}_t={\mathit{P}}^{-1}{\tilde{\mathit{u}}}_t$. For this form of dynamical system, we can use a mature algorithm termed Parallel Scan^[1] to reduce the running time complexity to $\mathcal{O}(\log t)$. Finally, when it goes back the original state space by Eq. 2 ${\widehat{\mathit{x}}}_{t+1}={\mathit{W}}_{\mathit{x}}{\mathit{g}}_{t+1}$, we can reformulate it as ${\widehat{\mathit{x}}}_{t+1}={\overline{\mathit{W}}}_{\mathit{x}}{\overline{\mathit{g}}}_{t+1}$, where ${\overline{\mathit{W}}}_{\mathit{x}}={\mathit{W}}_{\mathit{x}}\mathit{P}$.

2). Control Module:

We first encode observations into embeddings using transformer encoders, which effectively capture temporal dependencies in the BOLD signals. These embeddings are then processed through a batch normalization layer and fed into both: (1) A classification head that predicts task states. (2) A control gate generation function using tanh activation. This implements the diagonal control matrix $\mathit{C}$ described in the paper. We chose this approach for parameter efficiency and to prevent overfitting, as mentioned in Sec. III-C.

B. Loss duing Training

We plot the training loss and validation loss for the first 500 epochs on Fig. 7, yielding a smooth optimization dynamics of BRICK.

Fig. 7:

Training loss and validation loss during training. The losses are processed by $\log (\cdot )$ function. The first 500 epochs are plotted for better visualization.