Estimating Learning Effects: A Short-Time Fourier Transform Regression Model for MEG Source Localization

Abstract

Magnetoencephalography (MEG) has a high temporal resolution well-suited for studying perceptual learning. However, to identify where learning happens in the brain, one needs to apply source localization techniques to project MEG sensor data into brain space. Previous source localization methods, such as the short-time Fourier transform (STFT) method by Gramfort et al. [6] produced intriguing results, but they were not designed to incorporate trial-by-trial learning effects. Here we modify the approach in [6] to produce an STFT-based source localization method (STFT-R) that includes an additional regression of the STFT components on covariates such as the behavioral learning curve. We also exploit a hierarchical $L_{21}$ penalty to induce structured sparsity of STFT components and to emphasize signals from regions of interest (ROIs) that are selected according to prior knowledge. In reconstructing the ROI source signals from simulated data, STFT-R achieved smaller errors than a two-step method using the popular minimum-norm estimate (MNE), and in a real-world human learning experiment, STFT-R yielded more interpretable results about what time-frequency components of the ROI signals were correlated with learning.

Appendix

1.1 Appendix 1

Short-Time Fourier Transform (STFT). Our approach builds on the STFT implemented by Gramfort et al. in [6]. Given a time series $\varvec{U} = \{U(t), t = 1,\cdots ,T\}$, a time step $\tau _0$ and a window size $T_0$, we define the STFT as

$$\begin{aligned} \varPhi (\{U(t)\},\tau ,\omega _h) = \sum _{t=1}^T U(t) K(t-\tau ) e^{(-i \omega _h)} \end{aligned}$$

(6)

for $\omega _h = 2\pi h/T_0, h = 0,1,\cdots , T_0/2$ and $\tau = \tau _0, 2\tau _0, \cdots n_0 \tau _0$, where $K(t-\tau )$ is a window function centered at $\tau$, and $n_0 = T/\tau _0$. We concatenate STFT components at different time points and frequencies into a single vector in $\varvec{V} \in \mathbb {C}^{s}$, where $s = (T_0/2+1) \times n_0$. Following notations in [6], we also call the $K(t-\tau ) e^{(-i \omega _h)}$ terms STFT dictionary functions, and use a matrix’s Hermitian transpose $\varvec{\varPhi ^H}$ to denote them, i.e. $(\varvec{U}^T)_{1\times T} = ({\varvec{V}}^T)_{1\times s} (\varvec{\varPhi ^H})_{s \times T}$.

1.2 Appendix 2

The Karush-Kuhn-Tucker Conditions. Here we derive the Karush-Kuhn-Tucker (KKT) conditions for the hierarchical $L_{21}$ problem. Since the term $f(\varvec{z}) = \frac{1}{2} \sum _{r=1}^q ||\varvec{M} ^{(r)} - \varvec{G}( \sum _{k=1}^p X_k^{(r)} \varvec{Z}_k ) \varvec{\varPhi ^H}||_F^2$ is essentially a sum of squared error of a linear problem, we can re-write it as $f(\varvec{z}) = \frac{1}{2} || \varvec{b} - \varvec{A} \varvec{z} ||^2$, where $\varvec{z}$ again is a vector concatenated by entries in $\varvec{Z}$, $\varvec{b}$ is a vector concatenated by $\varvec{M}^{(1)}, \cdots , \varvec{M}^{(q)}$, and $\varvec{A}$ is a linear operator, such that $\varvec{A} \varvec{z}$ is the concatenated $\varvec{G}( \sum _{k=1}^p X_k^{(r)} \varvec{Z}_k ) \varvec{\varPhi ^H}, r = 1,\cdots , q$. Note that although $\varvec{z}$ is a complex vector, we can further reduce the problem into a real-valued problem by rearranging the real and imaginary parts of $\varvec{z}$ and $\varvec{A}$. Here for simplicity, we only derive the KKT conditions for the real case. Again we use $\{g_1, \cdots , g_h, \cdots , g_N \}$ to denote our ordered hierarchical group set, and $\lambda _h$ to denote the corresponding penalty for group $g_h$. We also define diagonal matrices $\varvec{D}_h$ such that

$$\begin{aligned} \varvec{D}_h(l,l) = \left\{ \begin{array}{l l} 1 &{}\text { if } l\in g_h \\ 0 &{} \text { otherwise } \end{array} \right. \forall h \end{aligned}$$

therefore, the non-zero elements of $\varvec{D}_h \varvec{z}$ is equal to $\varvec{z} |_{g_h}$. With the simplified notation, we re-cast the original problem into a standard formulation:

$$\begin{aligned} \min _{\varvec{z}} (\frac{1}{2}\Vert \varvec{b} - \varvec{A} \varvec{z} \Vert ^2_2 + \sum _h \lambda _h \Vert \varvec{D}_h \varvec{z} \Vert _2) \end{aligned}$$

(7)

To better describe the KKT conditions, we introduce some auxiliary variables, $\varvec{u} = \varvec{A} \varvec{z}, \varvec{v}_h = \varvec{D}_h \varvec{z}$. Then (7) is equivalent to

$$\begin{aligned} \min _{\varvec{z}, \varvec{u} , \varvec{v}_h}&(\frac{1}{2}\Vert \varvec{b} - \varvec{u} \Vert ^2_2 + \sum _h \lambda _h \Vert \varvec{v}_h \Vert _2) \\ \text {such that }&\varvec{u} = \varvec{A} \varvec{z}, \quad \varvec{v}_h= \varvec{D}_h \varvec{z}, \forall h \end{aligned}$$

The corresponding Lagrange function is

$$\begin{aligned} L(\varvec{z}, \varvec{u} , \varvec{v}_h,\varvec{\mu },\varvec{\xi }_h ) = \frac{1}{2}\Vert \varvec{b} - \varvec{u} \Vert ^2_2 + \sum _h \lambda _h\Vert \varvec{v}_h \Vert _2 + \varvec{\mu }^T ( \varvec{A} \varvec{z} - \varvec{u}) + \sum _{h} \varvec{\xi }_h^T ( \varvec{D}_h \varvec{z} - \varvec{v}_h ) \end{aligned}$$

where $\varvec{\mu }$ and $\varvec{\xi }_h$’s are Lagrange multipliers. At the optimum, the following KKT conditions hold

$$\begin{aligned} \frac{\partial {L}}{\partial {\varvec{u}}}&= \varvec{u} - \varvec{b} - \varvec{\mu } = 0 \end{aligned}$$

(8)

$$\begin{aligned} \frac{\partial {L}}{\partial {\varvec{z}}}&= \varvec{A}^T \varvec{\mu } + \sum _h \varvec{D}_h \varvec{\xi }_h = 0 \end{aligned}$$

(9)

$$\begin{aligned} \frac{\partial {L}}{\partial {\varvec{v}_h}}&= \lambda _h \partial { \Vert \varvec{v}_h \Vert _2} - \varvec{\xi }_h \ni 0, \forall h \end{aligned}$$

(10)

where $\partial { \Vert \cdot \Vert _2}$ is the subgradient of the $L_2$ norm. From (8) we have $\varvec{\mu } = \varvec{u} - \varvec{b}$, then (9) becomes $\varvec{A}^T (\varvec{u} - \varvec{b}) + \sum _h \varvec{D}_h \varvec{\xi }_h = 0$. Plugging $\varvec{u}= \varvec{Az}$ in, we can see that the first term $\varvec{A}^T ( \varvec{u} - \varvec{b}) = \varvec{A}^T (\varvec{A} \varvec{z} - \varvec{b})$ is the gradient of $f(\varvec{z}) = \frac{1}{2} \Vert \varvec{b} - \varvec{A} \varvec{z}\Vert _2^2$. For a solution $\varvec{z}_0$, once we plug in $\varvec{v}_h = \varvec{D}_h \varvec{z}_0$, the KKT conditions become

$$\begin{aligned}&\nabla f(\varvec{z})_{\varvec{z} = \varvec{z}_0} + \sum _h \varvec{D}_h \varvec{\xi }_h = 0 \end{aligned}$$

(11)

$$\begin{aligned}&\lambda _h \partial { \Vert \varvec{D}_h \varvec{z}_0 \Vert _2} - \varvec{\xi }_h \ni 0, \forall h \end{aligned}$$

(12)

In (12), we have the following according to the definition of subgradients

$$\begin{aligned}&\varvec{\xi }_h = \lambda _h \frac{ \varvec{D}_h \varvec{z}_0 }{\Vert \varvec{D}_h \varvec{z}_0 \Vert _2} \text { if } \Vert \varvec{D}_h \varvec{z}_0 \Vert _2 > 0 \\&\Vert \varvec{\xi }_h\Vert _2 \le \lambda _h \text { if } \Vert \varvec{D}_h \varvec{z}_0 \Vert _2 = 0 \end{aligned}$$

Therefore we can determine whether (11) and (12) hold by solving the following problem.

$$\begin{aligned} \min _{\varvec{\xi }_h}&\frac{1}{2} \Vert \nabla f(\varvec{z})_{\varvec{z} = \varvec{z}_0} + \sum _h \varvec{D}_h \varvec{\xi }_h\Vert _2^2 \\ \text {subject to }&\varvec{\xi }_h = \lambda _h \frac{ \varvec{D}_h \varvec{z}_0 }{\Vert \varvec{D}_h \varvec{z}_0 \Vert _2} \text { if } \Vert \varvec{D}_h \varvec{z}_0 \Vert _2 > 0 \\&\Vert \varvec{\xi }_h\Vert _2 \le \lambda _h \text { if } \Vert \varvec{D}_h \varvec{z}_0 \Vert _2 = 0 \end{aligned}$$

which is a standard group lasso problem with no overlap. We can use coordinate-descent to solve it. We define $\frac{1}{2} \Vert \nabla f(\varvec{z})_{\varvec{z} = \varvec{z}_0} + \sum _h \varvec{D}_h \varvec{\xi }_h\Vert _2^2$ at the optimum as a measure of violation of the KKT conditions.

Let $f_{J}$ be the function f constrained on a set J. Because the gradient of f is linear, if $\varvec{z_0}$ only has non-zero entries in J, then the entries of $\nabla f( \varvec{z})$ in J are equal to $\nabla f_{J} (\varvec{z} |_J)$ at $\varvec{z} = \varvec{z}_0$. In addition, $\varvec{\xi }_h$’s are separate for each group. Therefore if $\varvec{z}_0$ is an optimal solution to the problem constrained on J, then the KKT conditions are already met for entries in J (i.e. $\left( \nabla f(\varvec{z})_{\varvec{z} = \varvec{z}_0} + \sum _h \varvec{D}_h \varvec{\xi }_h\right) |_{J} = 0$); for $g_h \not \subset J$, we use ( $\frac{1}{2}\Vert \left( \nabla f(\varvec{z})_{\varvec{z} = \varvec{z}_0} + \sum _h \varvec{D}_h \varvec{\xi }_h\right) |_{g_h} \Vert ^2$) at the optimum as a measurement of how much the elements in group $g_h$ violate the KKT conditions, which is a criterion when we greedily add groups (see Algorithm 2).