Covariate Adjusted Functional Mixed Membership Models

Functional clustering generally assumes that each sample path is drawn from one of $K$ underlying cluster-specific sub-processes (James and Sugar, 2003; Chiou and Li, 2007; Jacques and Preda, 2014). In the GP framework, assuming that $f^{(1)},\dots,f^{(K)}$ are the $K$ underlying cluster-specific sub-processes with corresponding mean $\left(\mu^{(1)},\dots\mu^{(K)}\right)$ and covariance functions $\left(C^{(1)},\dots,C^{(K)}\right)$, the general form of a GP finite mixture model is typically expressed as:

where $\rho^{(k)}$ ($\sum_{k=1}^{K}\rho^{(k)}=1$) are the mixing proportions and $f_{i}$ are the sample paths for $i=1,\dots,N$. Let $\mathbf{x}_{i}=\left(X_{i1}\dots X_{iR}\right)^{\prime}\in\mathbb{R}^{R}:=\mathcal{X}$ be the covariates of interest associated with the $i^{th}$ observation. Extending the multivariate finite mixture of regressions model (McLachlan et al., 2019; Faria and Soromenho, 2010; Grün et al., 2007; Khalili and Chen, 2007; Hyun et al., 2023; Devijver, 2015) to a functional setting, the covariate adjustments would be encoded through the mean by defining a set of cluster-specific regression functions $\mu^{(k)}(\mathbf{x}_{i}):\mathcal{X}\times\mathcal{T}\rightarrow\mathbb{R}$. The mixture of GPs in (1) would then be generalized to include covariate information as follows:

where cluster specific GPs are now assumed to depend on mean functions $\mu^{(k)}(\mathbf{x})$, which are allowed to vary over the covariate space. Alternative representations of covariate adjusted mixtures have been proposed to allow the mixing proportions to depend on covariate information through generalized regression functions $\rho^{(k)}(\mathbf{x}_{i})$ (Jordan and Jacobs, 1994) or implicitly defining covariate adjusted partitions through a cohesion prior (Park and Dunson, 2010).

In both the finite mixture of regressions and mixture of experts settings, covariate adjustment does not change the overall interpretation of the finite mixture modeling framework, and the underlying assumption is that a function $f_{i}$ exhibits uncertain membership to only one of $K$ well-defined subprocesses $f^{(1:K)}$. Following Marco et al. (2024a), we contrast this representation of functional data with one where a sample path $f_{i}$ is allowed mixed membership to one or more well-defined sub-processes $f^{(1:K)}$. In the mixed membership framework, the underlying subprocesses $f^{(1:K)}$ are best interpreted as latent functional features, which may partially manifest in a specific path $f_{i}$.

Our representation of a functional mixed membership model hinges on the introduction of a new set of latent variables $\mathbf{z}_{i}=(Z_{i1},\ldots,Z_{iK})^{\prime}$, where $Z_{ik}\in[0,1]$ and $\sum_{k=1}^{K}Z_{ik}=1$. Given the latent mixing variables $\mathbf{z}_{i}$, each sample path is assumed to follow the following sampling structure:

Thus we can see that under the functional mixed membership model, each sample path is assumed to come from a convex combination of the underlying GPs, $f^{(k)}$. To avoid unwarranted restrictions on the implied covariance structure, the functional mixed membership model does not assume that the underlying GPs are mutually independent. Thus, we will let $C^{(k,j)}$ represent the cross-covariance function between the $k^{th}$ GP and the $j^{th}$ GP, for $1\leq k\neq j\leq K$. Letting $\boldsymbol{\mathcal{C}}$ be the collection of covariance and cross-covariance functions, we can specify the sampling model of the functional mixed membership model as

Leveraging the mixture of regressions framework in (2), the proposed covariate adjusted functional mixed membership model becomes:

In Section 2.2, we review the multivariate Karhunen-Loève (KL) construction (Ramsay and Silverman, 2005a; Happ and Greven, 2018), which allows us to have a joint approximation of the covariance structure of the $K$ underlying GPs. Using the joint approximation, we are able to concisely represent the $K$ GPs, facilitating inference on higher-dimensional functions such as surfaces. Using the KL decomposition, we are able to fully specify our proposed covariate adjusted functional mixed membership model (Equation 5) in Section 2.3.

The proposed modeling framework in Equation (5) depends on a set of mean and covariance functions, which are naturally represented infinite-dimensional parameters. Thus, in order to work in finite dimensions, we assume that the underlying GPs are smooth and lie in the $P$-dimensional subspace, $\boldsymbol{\mathcal{S}}\subset L^{2}(\mathcal{T})$, spanned by a set of linearly independent square-integrable basis functions, $\{b_{1},\dots,b_{P}\}$ ($b_{p}:\mathcal{T}\rightarrow\mathbb{R}$). Although the choice of basis functions is user-defined, the basis functions must be uniformly continuous in order to satisfy Lemma 2.2 of Marco et al. (2024a). In this manuscript, we will utilize B-splines for all case studies and simulation studies.

The assumption that $f^{(k)}\in\boldsymbol{\mathcal{S}}$ allows us to turn an infinite-dimensional problem into a finite-dimensional problem, making traditional inference tractable. Although tractable, accounting for covariance and cross-covariance functions separately leads to a model that needs $\mathcal{O}(K^{2}P^{2})$ parameters to represent the overall covariance structure. While the number of clusters, $K$, supported in applications is typically small, the number of basis functions, $P$, tends to follow the observational grid resolution, and may therefore be quite large. This is particularly true when one considers higher-dimensional functional data, such as surfaces, which necessitate a substantial number of basis functions, making it computationally expensive to fit models. Thus, we will use the multivariate KL decomposition (Ramsay and Silverman, 2005a; Happ and Greven, 2018) to reduce the number of parameters needed to estimate the covariance surface to $\mathcal{O}(KPM)$, where $M$ is the number of eigenfunctions used to approximate the covariance structure. A more detailed derivation of a KL decomposition for functions in $\boldsymbol{\mathcal{S}}$ can be found in Marco et al. (2024a).

To achieve a concise joint representation of the $K$ latent GPs, we define a multivariate GP, which we denote by $\mathbf{f}(\mathbf{t})$, such that

where $\mathbf{t}=\left(t^{(1)},t^{(2)},\dots,t^{(K)}\right)$ and $t^{(1)},t^{(2)},\dots,t^{(K)}\in\mathcal{T}$. Since $\boldsymbol{\mathcal{S}}\subset L^{2}$ is a Hilbert space, with the inner product defined as the $L^{2}$ inner product, we can see that $\mathbf{f}\in\boldsymbol{\mathcal{H}}:=\boldsymbol{\mathcal{S}}\times\boldsymbol{\mathcal{S}}\times\dots\times\boldsymbol{\mathcal{S}}:=\boldsymbol{\mathcal{S}}^{K}$, where $\boldsymbol{\mathcal{H}}$ is defined as the direct sum of the Hilbert spaces $\boldsymbol{\mathcal{S}}$, making $\boldsymbol{\mathcal{H}}$ also a Hilbert space. Since $\boldsymbol{\mathcal{H}}$ is a Hilbert space and the covariance operator of $\mathbf{f}$, denoted $\boldsymbol{\mathcal{K}}$, is a positive, bounded, self-adjoint, and compact operator, we know that there exists a complete set of eigenfunctions $\boldsymbol{\Psi}_{1},\dots,\boldsymbol{\Psi}_{KP}\in\boldsymbol{\mathcal{H}}$ and the corresponding eigenvalues $\lambda_{1}\geq\dots\geq\lambda_{KP}\geq 0$ such that $\boldsymbol{\mathcal{K}}\boldsymbol{\Psi}_{p}=\lambda_{p}\boldsymbol{\Psi}_{p},$ for $p=1,\dots,KP$. Using the eigenpairs of the covariance operator, we can rewrite $f^{(k)}$ as

where $\chi_{m}\sim\mathcal{N}(0,1)$ and $\Psi_{m}^{(k)}(t)$ is the $k^{th}$ element of $\boldsymbol{\Psi}_{m}(\mathbf{t})$. Since $\Psi_{m}^{(k)}(t)\in\boldsymbol{\mathcal{S}}$, we know that there exists $\phi_{m}\in\mathbb{R}^{P}$ such that $\sqrt{\lambda_{m}}\Psi_{m}^{(k)}(t)=\boldsymbol{\phi}_{km}^{\prime}\mathbf{B}(t)$, where $\mathbf{B}^{\prime}(t):=[b_{1}(t),b_{2}(t),\dots,b_{P}(t)]$. Similarly, since $\mu^{(k)}(\mathbf{x},t)\in\boldsymbol{\mathcal{S}}$, we can introduce a mapping $\mathbf{g}:\mathbb{R}^{R}\rightarrow\mathbb{R}^{P}$, such that $\mu^{(k)}(\mathbf{x},t)=\mathbf{g}_{k}(\mathbf{x})^{\prime}\mathbf{B}(t)$. Therefore, we arrive at the general form of our decomposition:

Using this decomposition, our covariance and mean structures can be recovered such that $C^{(k,k^{\prime})}(s,t)=\mathbf{B}^{\prime}(s)\left(\sum_{m=1}^{KP}\boldsymbol{\phi}_{km}\boldsymbol{\phi}^{\prime}_{k^{\prime}m}\right)\mathbf{B}(t)$ and $\mu^{(k)}(\mathbf{x},t)=\mathbf{g}_{k}(\mathbf{x})^{\prime}\mathbf{B}(t)$, for $1\leq k,k^{\prime}\leq K$ and $s,t\in\mathcal{T}$. To reduce the dimensionality of the problem, we will only use the first $M$ eigenpairs to approximate the $K$ stochastic processes. Although traditional functional principal component analysis (FPCA) will choose the number of eigenfunctions based on the proportion of variance explained, the same strategy cannot be employed in functional mixed membership models because the allocation parameters of the model are typically not known. Therefore, we suggest selecting a large value of $M$, and then proceeding with stable estimation through prior regularization.

In this section, we fully specify the covariate adjusted functional mixed membership model utilizing a truncated version of the KL decomposition, specified in Equation 6. We start by first specifying how the covariates of interest influence the mean function. Under the standard function-on-scalar regression framework (Faraway, 1997; Brumback and Rice, 1998), we have $\mu^{(k)}(\mathbf{x}_{i},t)=\beta_{k0}+\sum_{r=1}^{R}X_{ir}\beta_{kr}(t)$ for $k=1,\dots,K$. Since we assumed that $\mu^{(k)}(\mathbf{x}_{i},t)\in\boldsymbol{\mathcal{S}}$ for $k=1,\dots,K$, we know that $\beta_{k0},\dots,\beta_{kR}\in\boldsymbol{\mathcal{S}}$. Therefore, there exist $\boldsymbol{\nu}_{k}\in\mathbb{R}^{P}$ and $\boldsymbol{\eta}_{k}\in\mathbb{R}^{P\times R}$ such that

Under a standardized set of covariates, $\boldsymbol{\nu}_{k}$ specifies the population-level mean of the $k^{th}$ feature and $\boldsymbol{\eta}_{k}$ encodes the covariate dependence of the $k^{th}$ feature. Thus, assuming the mean structure specified in Equation 7 and using a truncated version of the KL decomposition specified in Equation 6, we can specify the sampling model of our covariate adjusted functional mixed membership model.

Let $\{\mathbf{Y}_{i}(.)\}_{i=1}^{N}$ be the observed sample paths that we want to model to the $K$ features, $f^{(k)}$, conditionally on the covariates of interest, $\mathbf{x}_{i}$. Since the observed sample paths are observed at only a finite number of points, we will let $\mathbf{t}_{i}=[t_{i1},\dots,t_{in_{i}}]^{\prime}$ denote the time points at which the $i^{th}$ function was observed. Without loss of generality, we will define the sampling distribution over the finite-dimensional marginals of $\mathbf{Y}_{i}(\mathbf{t}_{i})$. Using the general form of our proposed model defined in Equation 5, we have

where $\mathbf{S}(\mathbf{t}_{i})=[\mathbf{B}(t_{1})\cdots\mathbf{B}(t_{n_{i}})]\in\mathbb{R}^{P\times n_{i}}$ and $\boldsymbol{\Theta}$ is the collection of the model parameters. As defined in Section 2, $Z_{ik}$ are variables that lie on the unit simplex, such that $Z_{ik}\in(0,1)$ and $\sum_{k=1}^{K}Z_{ik}=1$. From this characterization, we can see that each observation is modeled as a convex combination of realizations from the $K$ features with additional Gaussian noise, represented by $\sigma^{2}$. If we integrate out the $\chi_{im}$ variables, for $i=1,\dots,N$ and $m=1,\dots,M$, we arrive at the following likelihood:

where $\boldsymbol{\Theta}_{-\chi}$ is the collection of the model parameters excluding the $\chi_{im}$ parameters ($i=1,\dots,N$ and $m=1,\dots,M$) and $\mathbf{V}(\mathbf{t}_{i},\mathbf{z}_{i})=\sum_{k=1}^{K}\sum_{k^{\prime}=1}^{K}Z_{ik}Z_{ik^{\prime}}\left\{\mathbf{S}^{\prime}(\mathbf{t}_{i})\sum_{m=1}^{M}\left(\boldsymbol{\phi}_{km}\boldsymbol{\phi}^{\prime}_{k^{\prime}m}\right)\mathbf{S}(\mathbf{t}_{i})\right\}$. Equation 9 illustrates that the proposed covariate adjusted functional mixed membership model can be expressed as an additive model; the mean structure is a convex combination of the feature-specific means, while the covariance can be written as a weighted sum of covariance functions and cross-covariance functions.

To have an adequately expressive and scalable model, we approximate the covariance surface of the $K$ features using $M$ scaled pseudo-eigenfunctions. In this framework, orthogonality will not be imposed on the $\boldsymbol{\phi}_{km}^{\prime}\mathbf{B}(t)$ parameters, making them pseudo-eigenfunctions instead of the eigenfunctions described in Section 2.2. From a modeling prospective, this allows us to sample on an unconstrained space, facilitating better Markov chain mixing and easier sampling schemes. Although direct inference on the eigenfunctions is no longer available, a formal analysis can still be conducted by reconstructing the posterior samples of the covariance surface and calculating eigenfunctions from the posterior samples. To avoid overfitting of the covariance functions, we follow Marco et al. (2024a) by using the multiplicative gamma process shrinkage prior proposed by Bhattacharya and Dunson (2011) to achieve adaptive regularized estimation of the covariance structure. Therefore, letting $\phi_{kpm}$ be the $p^{th}$ element of $\boldsymbol{\phi}_{km}$, we have the following:

where $1\leq k\leq K$, $1\leq p\leq P$, $1\leq m\leq M$, and $2\leq j\leq M$. By letting $\alpha_{2}>\beta_{2}$, we can show that $\mathbb{E}(\tilde{\tau}_{mk})>\mathbb{E}(\tilde{\tau}_{m^{\prime}k})$ for $1\leq m<m^{\prime}\leq M$, leading to the prior on $\phi_{kpm}$ having stochastically decreasing variance as $m$ increases. This will have a regularizing effect on the posterior draws of $\phi_{kpm}$, making $\phi_{kpm}$ more likely to be close to zero as $m$ increases.

In functional data analysis, we often desire smooth mean functions to protect against overfit models. Therefore, we use a first-order random walk penalty proposed by Lang and Brezger (2004) on the $\boldsymbol{\nu}_{k}$ and $\boldsymbol{\nu}_{k}$ parameters to promote adaptive smoothing of the mean function of the features in our model. Therefore, we have that $P(\boldsymbol{\nu}_{k}\mid\tau_{\boldsymbol{\nu}_{k}})\propto\exp\left(-\frac{\tau_{\boldsymbol{\nu}_{k}}}{2}\sum_{p=1}^{P-1}\left(\nu_{pk}-{\nu}_{(p+1)k}\right)^{2}\right),$ for $k=1,\dots,K$, where $\tau_{\boldsymbol{\nu}_{k}}\sim\Gamma(\alpha_{\boldsymbol{\nu}},\beta_{\boldsymbol{\nu}})$ and $\nu_{pk}$ is the $p^{th}$ element of $\boldsymbol{\nu}_{k}$. Similarly, we have that $P(\{\eta_{prk}\}_{p=1}^{P}\mid\tau_{\boldsymbol{\eta}_{rk}})\propto\exp\left(-\frac{\tau_{\boldsymbol{\eta}_{rk}}}{2}\sum_{p=1}^{P-1}\left(\eta_{prk}-{\eta}_{(p+1)rk}\right)^{2}\right),$ for $k=1,\dots,K$ and $r=1,\dots,R$, where $\tau_{\boldsymbol{\eta}_{rk}}\sim\Gamma(\alpha_{\boldsymbol{\eta}},\beta_{\boldsymbol{\eta}})$ and $\eta_{prk}$ is the $p^{th}$ row and $r^{th}$ column of $\boldsymbol{\eta}_{k}$. Following previous mixed membership models (Heller et al., 2008; Marco et al., 2024a, b), we assume that $\mathbf{z}_{i}\mid\boldsymbol{\pi},\alpha_{3}\sim_{iid}Dir(\alpha_{3}\boldsymbol{\pi})$, $\boldsymbol{\pi}\sim Dir(\mathbf{c}_{\pi})$, and $\alpha_{3}\sim Exp(b)$. Lastly, we assume that $\sigma^{2}\sim IG(\alpha_{0},\beta_{0})$. The posterior distributions of the parameters in our model, as well as a sampling scheme with tempered transitions, can be found in Section 2 of the Supplementary Materials. A covariate adjusted model where the covariance is also dependent on the covariates of interest can be found in Section 4 of the Supplementary Materials.

Mixed membership models face multiple identifiability problems due to their increased flexibility over traditional clustering models (Chen et al., 2022; Marco et al., 2024a, b). Like traditional clustering models, mixed membership models also face the common label switching problem, where an equivalent model can be formulated by permuting the labels or allocation parameters. Although this is one source of non-identifiability, relabelling algorithms can be formulated from the work of Stephens (2000). More complex identifiability problems arise since the allocation parameters are now continuous variables on the unit simplex, rather than binary variables like in clustering. Specifically, an equivalent model can be constructed by rotating and/or changing the volume of the convex polytope constructed by the allocation parameters through linear transformations of the allocation parameters (Chen et al., 2022). Therefore, geometric assumptions must be made on the allocation parameters to ensure that the mixed membership models are identifiable.

One way to ensure that we have an identifiable model up to a permutation of the labels is to assume that the separability condition holds (Pettit, 1990; Donoho and Stodden, 2003; Arora et al., 2012; Azar et al., 2001; Chen et al., 2022). The separability condition assumes that at least one observation belongs completely to each of the $K$ features. These observations that belong only to one feature act as “anchors”, ensuring that the allocation parameters are identifiable up to a permutation of the labels. Although the separability condition is conceptually simple and relatively easy to implement, it makes strong geometric assumptions on the data-generating process for mixed membership models with three or more features. Weaker geometric assumptions, known as the sufficiently scattered condition (Chen et al., 2022), that ensure an identifiable model in mixed membership models with 3 or more features are discussed in Chen et al. (2022). Although the geometric assumptions are relatively weak, implementing these constraints is nontrivial. When extending multivariate mixed membership models to functional data and introducing a covariate-dependent mean structure, ensuring an identifiable mean and covariance structure requires further assumptions.

Lemma 2.1 states a set of sufficient conditions that lead to an identifiable mean and covariance structure up to a permutation of the labels. The proof of Lemma 2.1 can be found in Section 1 of the Supporting Materials. The first assumption is similar to the assumption needed in the linear regression setting to ensure identifiability; however, the intercept term is not included in the design matrix, as the $\boldsymbol{\nu}_{k}$ parameters already account for the intercept. From the second assumption we can see that we assume geometric assumptions on the allocation parameters, provide a lower bound for the number of functional observations, and we assume that the matrix $\mathbf{C}$ is full column rank. The matrix $\mathbf{C}$ is constructed using the allocation parameters, and although the condition on $\mathbf{C}$ may seem restrictive, it often holds in practice. We note that this condition may not hold if there are not enough unique allocation parameter values or if most of your allocation parameters lie on a lower-dimensional subspace of the unit $(K-1)$-simplex. Although the individual $\boldsymbol{\phi}_{km}$ parameters are not identifiable in our model, an eigen analysis can still be performed by constructing posterior draws of the covariance structure and calculating the eigenvalues and eigenfunctions of the posterior draws. Section 3 provides empirical evidence that the mean and covariance structure converge to the truth as more observations are accrued.

Function-on-scalar regression is a common method in FDA that allows the mean structure of a continuous stochastic process to be dependent on scalar covariates. In function-on-scalar regression, we often assume that the response is a GP, and that the covariates of interest are vector-valued. A comprehensive review of the broader area of functional regression can be found in Ramsay and Silverman (2005b) and Morris (2015). While there have been many advancements and generalizations (Krafty et al., 2008; Reiss et al., 2010; Goldsmith et al., 2015; Kowal and Bourgeois, 2020) of function-on-scalar regression since the initial papers of Faraway (1997) and Brumback and Rice (1998), the general form of function-on-scalar regression can be expressed as follows:

where $Y(t)$ is the response function evaluated at $t$, $\beta_{r}(\cdot)$ is the functional coefficient representing the effect that the $r^{th}$ covariate ($X_{r}$) has on the mean structure, and $\epsilon$ is a mean-zero Gaussian process with covariance function $\mathcal{C}$. The function $\mu:\mathcal{T}\rightarrow\mathbb{R}$ in Equation (10) represents the mean of the GP when all of the covariates, $X_{ir}$ are set to zero. Unlike the traditional setting for multiple linear regression in finite-dimensional vector spaces, function-on-scalar regression requires the estimation of the infinite-dimensional functions $\mu$ and $\beta_{1},\dots,\beta_{R}$ from a finite number of observed sample paths at a finite number of points ($\mathbf{Y}_{i}(\mathbf{t}_{i})$ for $i=1,\dots,N$ and $\mathbf{t}_{i}=[t_{i1},\dots,t_{in_{i}}]^{\prime}$).

To make inference tractable, we assume that the data lie in the span of a finite set of basis functions, which will allow us to expand $\mu$ and $\beta_{1},\dots,\beta_{R}$ as a finite sum of the basis functions. The set of basis functions can be specified using data-driven basis functions or by specifying the basis functions a-priori. If the basis functions are specified a-priori, common choices of basis functions are B-splines and wavelets due to their flexibility, as well as Fourier series for periodic functions. Alternatively, if the use of data-driven basis functions is desired, a common choice is to use the eigenfunctions of the covariance operator as basis functions. In order to estimate the eigenfunctions, functional principal component analysis (Shang, 2014) is often performed and the obtained estimates of the eigenfunctions are used. Functional principal component analysis faces a similar problem in that the objective is to estimate eigenfunctions using only a finite number of sample paths observed at a finite number of points. To solve this problem, Rice and Silverman (1991) proposes using a splines basis to estimate smooth eigenfunctions, while Yao et al. (2005) proposes using local linear smoothers to estimate the smooth eigenfunctions. Therefore, even using data-driven basis functions require smoothing assumptions, suggesting similar results between data-driven basis functions and a reasonable set of a-priori specified basis functions paired with a penalty to prevent overfitting.

Specifying the basis functions a-priori ($b_{1}(t),\dots,b_{P}(t)$), and letting $\mathbf{Y}_{i}(\mathbf{t}_{i})$ be the observed sample paths at points $\mathbf{t}_{i}=[t_{i1},\dots,t_{in_{i}}]^{\prime}$ ($i=1,\dots,N$), we can simplify Equation (10) to get

where $\tilde{\boldsymbol{\nu}}\in\mathbb{R}^{P}$, $\tilde{\boldsymbol{\eta}}\in\mathbb{R}^{P\times R}$, and $\mathbf{S}(\mathbf{t}_{i})=[\mathbf{B}(t_{1})\cdots\mathbf{B}(t_{n_{i}})]\in\mathbb{R}^{P\times n_{i}}$ are the set of basis function evaluated at the time points of interest. As specified in the previous sections, $\mathbf{B}^{\prime}(t):=[b_{1}(t),b_{2}(t),\dots,b_{P}(t)]$. Equation (11) shows that the function $\mu(.)$ evaluated at the points $\mathbf{t}_{i}$ can be represented by $\mathbf{S}^{\prime}(\mathbf{t}_{i})\tilde{\boldsymbol{\nu}}$, and similarly the functional coefficients $\beta_{1}(\cdot),\dots,\beta_{R}(\cdot)$ can be represented by $\mathbf{S}^{\prime}(\mathbf{t}_{i})\tilde{\boldsymbol{\eta}}$. Therefore, we are left to estimate $\tilde{\boldsymbol{\nu}}$, $\tilde{\boldsymbol{\eta}}$, and the parameters associated with the covariance function of $\epsilon(\cdot)$, denoted $\mathcal{C}$.

The covariance function $\mathcal{C}$ represents the within-function covariance structure of the data. In the simplest case, we often assume that $\boldsymbol{\epsilon}_{i}(\mathbf{t}_{i})\sim\mathcal{N}(\mathbf{0},\tilde{\sigma}^{2}\mathbf{I}_{n_{i}})$, which means that we only need $\tilde{\sigma}^{2}$ to specify $\mathcal{C}$. In more complex models (Faraway, 1997; Krafty et al., 2008), one may make less restrictive assumptions and assume that the covariance $\boldsymbol{\epsilon}_{i}(\mathbf{t}_{i})\sim\mathcal{N}(\mathbf{0}_{n_{i}},\tilde{V}(\mathbf{t}_{i})+\tilde{\sigma}^{2}\mathbf{I}_{n_{i}})$, where $\tilde{V}(\cdot)$ is a low-dimensional approximation of a smooth covariance surface using a truncated eigendecomposition. Although functional regression usually assumes that the functions are independent, functional regression models have been proposed to model between-function variation, or cases where observations can be correlated (Morris and Carroll, 2006; Staicu et al., 2010).

Assuming a relatively general covariance structure as in Krafty et al. (2008), the function-on-scalar model assumes the following distributional assumptions on our sample paths:

From Equation (12) it is apparent that the proposed covariate adjusted mixed membership model specified in Equation (9) is closely related to function-on-scalar regression. The key difference between the two models is that the covariate adjusted mixed membership model does not assume a common mean and covariance structure across all observations conditionally on the covariates of interest. Instead, the covariate adjusted mixed membership model allows each observation to be modeled as a convex combination of $K$ underlying features. Each feature is allowed to have different mean and covariance structures, meaning that the covariates are not assumed to have the same affect on all observations. By allowing this type of heterogeneity in our model, we can conduct a more granular analysis and identify subgroups that interact differently with the covariates of interest. Alternatively, if there are subgroups of the population that interact differently with the covariates of interest, then the results from a function-on-scalar regression model will be confounded, as the effects will likely be averaged out in the analysis. The finite mixture of experts model and the mixture of regressions model offer a reasonably granular level of insight, permitting inference at the sub-population level. However, they do not provide the individual-level inference achievable with the proposed covariate adjusted functional mixed membership model.

In the first simulation study, we explore the empirical convergence properties of the proposed covariate adjusted functional mixed membership models. Specifically, we generate data from a covariate adjusted functional mixed membership model and see how well the proposed framework can recover the true mean, covariance, and allocation structures. To evaluate how well we can recover the mean, covariance, and cross-covariance functions, we calculate the relative mean integrated square error (R-MISE), which is defined as $\text{R-MISE}=\frac{\int\{f(t)-\hat{f}(t)\}^{2}\text{d}t}{\int f(t)^{2}\text{d}t}\times 100\%$ or $\text{R-MISE}=\frac{\int_{t}\int_{x}\{f(t,x)-\hat{f}(t,x)\}^{2}\text{d}x\text{d}t}{\int_{t}\int_{x}f(t,x)^{2}\text{d}x\text{d}t}\times 100\%$ in the case of a covariate adjusted model. In this simulation study, $\hat{f}(t)$ will be the posterior median obtained from our posterior samples. To measure how well we recover the allocation structure, $Z_{ik}$, we calculated the root-mean-square error (RMSE). In addition to studying the empirical convergence properties of correctly specified models, we also included a scenario where we fit a covariate adjusted functional mixed membership model, when the generating truth had no covariate dependence. Conversely, we also studied the scenario where we fit a functional mixed membership model with no covariates, when the generating truth was generated from a covariate adjusted functional mixed membership model with one covariate.

Figure 1 contains a visualization of the performance metrics from each of the five scenarios considered in this simulation study. We can see that when the model is specified correctly, it does a good job in recovering the mean structure, even with relatively few observations. On the other hand, a relatively large number of observations are needed to recover the covariance structure when we have a correctly specified model. This simulation study also shows that we pay a penalty in terms of statistical efficiency when we over-specify a model; however, the over-specified model still shows signs of convergence to the true mean, covariance, and allocation structures. Conversely, an under-specified model shows no signs of converging to the true mean and covariance structure as we get more observations. Additional details on how the simulation was conducted, as well as more detailed visualizations of the simulation results, can be found in Section 3.1 of the Supplementary Materials.

When considering the use of a mixed membership model to gain insight into the heterogeneity found in a dataset, a practitioner must specify the number of features to use in the mixed membership model. Correctly specifying the number of features is crucial, as an over-specification will lead to challenges in interpretability of the results, and an under-specification will lead to a model that is not flexible enough to model the heterogeneity in the model. Although information criteria such as BIC and heuristics such as the elbow-method have been shown to be informative in choosing the number of parameters in mixed membership models (Marco et al., 2024a, b), it is unclear whether these results extend to covariate adjusted models. In this case, we consider the performance of AIC (Akaike, 1974), BIC (Schwarz, 1978), DIC (Spiegelhalter et al., 2002; Celeux et al., 2006), and the elbow-method in choosing the number of features in a covariate adjusted functional mixed membership model with one covariate. The information criteria used in this section are specified in Section 3.2 of the Supplementary Materials. As specified, the optimal model should have the smallest AIC, the largest BIC, and the smallest DIC.

To evaluate the performance of the information criteria and the elbow method, 50 datasets were generated from a covariate adjusted functional mixed membership model with three features and one covariate. Each dataset was analyzed using four covariate adjusted functional mixed membership models; each with a single covariate and a varying number of features ($K=2,3,4,5$). Figure 2 shows the performance of the three information criteria. Similarly to Marco et al. (2024a) and Marco et al. (2024b), we can see that BIC had the best performance, picking the three-feature model all 50 times. From the simulation study, we found that AIC and DIC tended to choose more complex models, with AIC and DIC choosing the correctly specified model 68% and 26% of the time, respectively. Looking at the log-likelihood plot, we can see that there is a distinct elbow at $K=3$, meaning that using heuristic methods such as the elbow method tends to be informative in choosing the number of features in covariate adjusted functional mixed membership models. Thus, we recommend using BIC and the elbow-method in conjunction to choose the number of features, as even though BIC correctly identified the model every time in the simulation, there were times when the BIC between the three-feature model and the four-feature model were similar. Additional details on how this simulation study was conducted can be found in Section 3.2 of the Supplementary Materials.

In the first part of this case study, we are primarily interested in modeling the developmental changes in alpha frequencies. As stated in Haegens et al. (2014), shifts in PAF can confound the alpha power measures, demonstrating the need for a model that controls for the age of the child. In this subsection, we fit a covariate adjusted functional mixed membership model using the log transformation of age as the covariate. Using conditional predictive ordinates and comparing the pseudomarginal likelihoods, we found that a log transformation of age provided a better fitting model than using a model with age untransformed. Although this model does not use information on the diagnostic group, we are also interested in characterizing the differences in the observed alpha oscillations between individuals across diagnostic groups.

From Figure 4, we can see that the first feature consists mainly of aperiodic neural activity patterns, which are commonly referred to as a $1/f$ trend or pink noise. On the other hand, the second feature can be interpreted as a distinct alpha peak, which is considered periodic neural activity. For children that heavily load on the second feature, we can see that as they age, the alpha peak increases in magnitude and the PAF shifts to a higher frequency, which has been observed in many other studies (Haegens et al., 2014; Rodríguez-Martínez et al., 2017; Scheffler et al., 2019). From a clinical perspective, it is also valuable to compare children’s alpha power conditional on age, since we know that there are developmental changes in alpha oscillations as children age. From Figure 4, we can also see that on average children with ASD have less attenuated alpha peaks compared to their age-adjusted TD counterpart.

One key feature of this model is that it allows us to directly compare individuals in the TD group and ASD group on the same spectrum. However, to do this, we had to assume that the alpha oscillations of children with ASD and TD children can be represented as a continuous mixture of the same two features shown in Figure 4. Just as developmental shifts in PAF can confound the measures of alpha power (Haegens et al., 2014), the assumption that the alpha oscillations of children with ASD and TD children can be represented by the same two features can also confound the results found in this section if the assumption is shown to be incorrect. In Section 4.2, we relax this assumption and allow the functional features to differ according to the diagnostic group.

Previous studies have shown that the emergence of alpha peaks and the developmental shifts in frequency are atypical in children with ASD (Dickinson et al., 2018; Scheffler et al., 2019; Marco et al., 2024a, b). Therefore, it may not be realistic to assume that the alpha oscillations of children with ASD and TD children can be represented by the same two features. In this section, we will fit a covariate adjusted functional mixed membership model using the log transformation of age, clinical diagnosis, and an interaction between the log transformation of age and clinical diagnosis as covariates of interest. By including the interaction between the log transformation of age and diagnostic group, we allow for differences in the developmental changes of alpha oscillations between diagnostic groups.

Compared to the model used in Section 4.1, this model is more flexible; it allows different features for each diagnostic group and different developmental trajectories for each diagnostic group. Although this model is more expressive, the comparison of individuals across diagnostic groups is no longer straightforward. The model in Section 4.1 did not use any diagnostic group information, which allowed us to easily compare TD children with ASD children as the distribution of the developmental trajectory depended only on the allocation parameters. On the other hand, comparing individuals across diagnostic groups is not as simple as looking at the inferred allocation parameters, as the underlying features are group-specific. Therefore, two individuals with the same allocation parameters that are in diagnostic groups will likely have completely different developmental trajectories. Moreover, in the model used in Section 4.1, we knew that at least one observation belonged to each feature, making the features interpretable as the extreme cases observed. However, for the more flexible model used in this subsection, there will not be at least one observation belonging to each feature for each diagnostic group. We can see that in the TD group, no individual has a loading of more than 50% on the first feature, meaning that the first TD-specific feature cannot be interpreted as the most extreme case observed in the TD group. Therefore, direct interpretation of the group-specific features is less useful, and instead we will focus on looking at the estimated trajectories of individuals conditional on empirical percentiles of the estimated allocation structure. Visualizations of the mean structure for the two features can be found in the Section 3.3 Supplementary Materials.

From Figure 5, we can see the estimated developmental trajectories of TD children and children with ASD, conditional on different empirical percentiles of the estimated allocation structure. Looking at the estimated developmental trajectories of children with a group-specific median allocation structure, we can see that both TD children and children diagnosed with ASD have a discernible alpha peak and that the expected PAF shifts to higher frequencies as the child develops. However, compared to TD children, ASD children have more of a $1/f$ trend, and the shift in PAF is not as large in magnitude. Similar differences between TD children and children diagnosed with ASD can be seen when we look at the $10^{th}$ percentile of observations that belong to the first group-specific feature. However, when looking at the group-specific $90^{th}$ percentile, we can see that children with ASD do not have a discernible alpha peak at 25 months old, and a clear alpha peak never really appears as they develop. Alternatively, TD children in the $90^{th}$ percentile still have a discernible alpha peak that becomes more prominent and shifts to higher frequencies as they develop. However, compared to TD individuals in the $10^{th}$ or $50^{th}$ percentile, there is more aperiodic signal ($1/f$ trend) present in the alpha power of TD children in the $90^{th}$ percentile.

As discussed in Section 2.5, the proposed covariate adjusted mixed membership model can be thought of as a generalization of function-on-scalar regression. Figure 6 contains the estimated developmental trajectories obtained by fitting a function-on-scalar regression model using the package “refund” package in R (Goldsmith et al., 2016). The results of the function-on-scalar regression coincide with the estimated developmental trajectory of the alpha oscillations obtained from our covariate adjusted mixed membership model, conditional on the group-specific mean allocation (group-specific mean allocations are represented by triangles in the bottom panel of Figure 5). As discussed in Section 2.5, function-on-scalar regression assumes a common mean for all individuals, controlling for age and diagnostic group. This is a relatively strong assumption in this case, as ASD is relatively broadly defined; with diagnosed individuals having varying degrees of symptoms. Indeed, from the function-on-scalar results, we would only be able to make population-level inference; giving no insight into the heterogeneity of alpha oscillation developmental trajectories for children with ASD (Figure 5). Overall, the added flexibility of covariate adjusted mixed membership models allows scientists to have greater insight into the developmental changes of alpha oscillations. Technical details on how the models for Section 4.1 and Section 4.2 were fit can be found in Section 3.3 of the Supplementary Materials. Section 3.3 of the Supplementary Materials also discusses model comparison using conditional predictive ordinates (Pettit, 1990; Chen et al., 2012; Lewis et al., 2014) and pseudomarginal likelihood.