S UBSPACE N ETWORK : D EEP M ULTI -T ASK C ENSORED R EGRESSION FOR M ODELING N EURODEGENERATIVE D ISEASES
Anonymous authors
Paper under double-blind review
A BSTRACT
Over the past decade a wide spectrum of machine learning models have been developed to model the neurodegenerative diseases, associating biomarkers, especially non-intrusive neuroimaging markers, with key clinical scores measuring the cognitive status of patients.
Multi-task learning (MTL) has been extensively explored in these studies to address challenges associated to high dimensionality and small cohort size.
However, most existing MTL approaches are based on linear models and suffer from two major limitations: 1) they cannot explicitly consider upper/lower bounds in these clinical scores; 2) they lack the capability to capture complicated non-linear effects among the variables.
In this paper, we propose the Subspace Network , an efﬁcient deep modeling approach for non-linear multi-task censored regression.
Each layer of the subspace network performs a multi-task censored regression to improve upon the predictions from the last layer via sketching a low-dimensional subspace to perform knowledge transfer among learning tasks.
Under mild assumptions, for each layer the parametric subspace can be recovered using only one pass of training data.
Empirical results demonstrate that the proposed subspace network quickly picks up correct parameter subspaces, and outperforms state-of-the-arts in predicting neurodegenerative clinical scores using information in brain imaging.
I NTRODUCTION
Recent years have witnessed increasing interests on applying machine learning (ML) techniques to analyze biomedical data.
Such data-driven approaches deliver promising performance improvements in many challenging predictive problems.
For example, in the ﬁeld of neurodegenerative diseases such as Alzheimer’s disease and Parkinson’s disease, researchers have exploited ML algorithms to predict the cognitive functionality of the patients from the brain imaging scans, e.g., using the magnetic resonance imaging (MRI) as in Adeli-Mosabbeb et al.
; Zhang et al.
; Zhou et al. (2011b).
As a key ﬁnding, there are typically various types of prediction targets (e.g., cognitive scores), and they can be jointly learned using multi-task learning (MTL), e.g., Caruana (1998); Evgeniou & Pontil (2004); Zhang et al.
(2012), where the predictive information is shared and transferred among related models to reinforce their generalization performance.
Two challenges persist despite the progress of applying MTL in disease modeling problems.
First, it is important to notice that clinical targets, different from typically regression targets, are often naturally bounded.
For example, the result from Mini-Mental State Examination (MMSE) test, a key reference for deciding cognitive impairments, ranges from 0 to 30 (a healthy subject): a smaller score indicates a higher level of cognitive dysfunction (please refer to Tombaugh & McIntyre (1992)).
Other cognitive scores, such as Clinical Dementia Rating Scale (CDR) Hughes et al.
and
Alzheimer’s Disease Assessment Scale-Cog (ADAS- Cog) Rosen et al.
(1984), also have speciﬁc upper and lower bounds.
Most existing approaches, e.g., Zhang et al.
Zhou et al.
(2011b);
Poulin et al.
(2011), relied on linear regression without considering the range constraint, partially due to the fact that mainstream MTL models for regression, e.g., Jalali et al.
; Argyriou et al.
(2007); Zhang et al.
; Zhou et al.
(2011a), are developed using the least squares loss and cannot be directly extended to censored regressions.
As the second challenge, a majority of MTL research focused on linear models because of computational efﬁciency and theoretical guarantees.
However,
V (1) Subspace
Sketching
U (1) Parameter
Subspace
R eL
U
V (2)
U (2)
R eL
U } Layerx y x
R eL
U
V (3)
U (3) } Layerx
V (k)
U (k) } Layer k x linear models cannot capture the complicated non-linear relationship between features and clinical targets.
For example, Association et al.
(2013) showed the early onset of Alzheimer’s disease to be related to single-gene mutations on chromosomes 21 , 14 , and 1 , and the effects of such mutations on the cognitive impairment are hardly linear (please refer to Martins et al.
; Sweet et al.
Recent advances in multi-task deep neural networks Seltzer & Droppo (2013); Zhang et al.
Wu et al.
(2015) provide a promising direction, but their model complexity and high demands of training data prohibit their broader usages in clinical cohort studies.
To address the aforementioned challenges, we propose a novel, solid, and efﬁcient deep modeling approach for non-linear multi-task censored regression, called Subspace Network , highlighting the following multi-fold technical innovations:
• It efﬁciently builds up a deep network in a layer-by-layer feedforward fashion, and in each layer considers a censored regression problem.
The layer-wise training allows us to grow a deep model efﬁciently.
• It explores a low-rank s ubspace structure that captures task relatedness for better predictions.
A critical difference on subspace decoupling between previous studies such as Mardani et al.
Shen et al.
and our method lies in our assumption of a low-rank structure on the parameter space among tasks rather than the original feature space.
• By leveraging the recent advances in online subspace sensing Mardani et al.
; Shen et al.
, we show that the parametric subspace can be recovered for each layer with feeding only one pass of the training data, allowing for more efﬁcient layer-wise training.
Synthetic experiments verify the technical claims of the proposed subspace network, and it outperforms various state-of-the-arts methods in modeling neurodegenerative diseases on real datasets.
M ULTI - TASK CENSORED REGRESSION VIA PARAMETER SUBSPACE SKETCHING AND REFINEMENT
In censored regression, we are given a set of N observations D = { ( x i , y i ) }
Ni =1 of D dimensional feature vectors { x i ∈ R D } and T corresponding outcomes { y i ∈ R T + } , where each outcome y
i,t ∈ R + , t ∈ { 1 , · · · , T } , can be cognitive scores (e.g., MMSE and ADAS-Cog) or other biomarkers of interest such as proteomics 1 .
For each outcome, the censored regression assumes a nonlinear relationship between the features and the outcome through a rectiﬁed linear unit (ReLU) transformation, i.e., y i,t = ReLU (cid:0) W (cid:62) t x i + (cid:15) (cid:1)
where W t ∈ R
D is the coefﬁcient for input features, (cid:15) is i.i.d.
noise, and ReLU is deﬁned by ReLU( z ) =
max( z, 0) .
We can thus collectively represent the censored regression for multiple tasks by:
[Equation]
where W =
[ W 1 , . . , W T ] (cid:62) ∈ R T × D is the matrix including all the linear coefﬁcients.
We consider the regression problem for each outcome as a learning task, and one commonly used task relationship assumption is that the transformation matrix W ∈ R T × D belongs to a linear low-rank subspace U .
The subspace allows us to represent W as product of two matrices, W = UV , where columns of U ∈ R T × R =
[ U 1 , . . . , U T ] (cid:62) span the linear subspace U , and V ∈ R R
× D is the embedding coefﬁcient.
We note that the output y can be entry-wise decoupled, such that for each component y i,t = ReLU( U (cid:62)
t V x i + (cid:15) ) .
By assuming a Gaussian noise (cid:15) ∼ N (0 , σ 2 ) , we derive the following likelihood function:
[Equation]
where φ is the probabilistic density function of the standardized Gaussian N (0 , 1) and Q is the standard Gaussian tail.
The likelihood of ( x i , y i ) pair is thus given by:
[Equation]
The likelihood function allows us to estimate subspace U and coefﬁcient V from data D .
To enforce a low-rank subspace, one common approach is to impose a trace norm on UV , where trace norm of a matrix A is deﬁned by (cid:107) A (cid:107) ∗ =
(cid:80) j s j ( A ) and s j ( A ) is the j th singular value of A .
Since (cid:107) UV (cid:107) ∗ =
min U,V 1 2 ( (cid:107) U (cid:107) 2 F + (cid:107) V (cid:107) 2 F ) , e.g., see Srebro et al.
; Mardani et al.
, the objective function of multi-task censored regression problem is given by:
[Equation]
2. O NLINE ALGORITHM
We propose to solve the objective in (2) via the block coordinate descent approach which is reduced to iteratively updating the following two subproblems:
[Equation]
Deﬁne the instantaneous cost of the i -th datum:
[Equation]
and the online optimization form of (2) can be recast as an empirical cost minimization given below:
[Equation]
According to the analysis in Section 2.2, one pass of the training data can warrant the subspace learning problem.
We outline the solver for each subproblem as follows:
Problem (P:V) sketches parameters in the current space.
We solve (P:V) using gradient descent.
The parameter sketching couples all the subspace dimensions in V (not decoupled as in Shen et al.
(2016)), and thus we need to solve this collectively.
The update of V ( V + ) can be obtained by solving the online problem given below:
[Equation]
V + can be computed by the following gradient update:
[Equation]
Require:
Training data D
= { ( x i , y i ) }
Ni =1 , rank parameters λ and R ,
Ensure: parameter subspace U , parameter sketch V
Initialize U − at random for i = 1 , . . . , N do // 1.
Sketching parameters in the current subspace
[Equation]
Parallel subspace reﬁnement { U t } Tt = for t = 1 , . . . , T do
U + t = arg min
U t − log Pr(
[Equation]
(cid:107)
[Equation]
(cid:107)end for
Set U − =
U + , V − =
V + end for where the gradient is given by:
[Equation]
where z i,t = σ − 1 ( − (cid:0) U − t (cid:1) (cid:62) V x ) .
The Algorithm 3 for solving (P:V) is summarized in the Appendix.
Problem (P:U) reﬁnes the subspace U + based on sketching.
We solve (P:U) using stochastic gradient descent (SGD).
We note that the problem is decoupled for different subspace dimensions t = 1 , . . .
, T (i.e., rows of U ).
With careful parallel design, this procedure can be done very efﬁciently.
Given a training data point ( x i , y i ) , the problem related to the t -th subspace basis is given by: [Equation] i
[Equation]
We can revise subspace by the following gradient update:
[Equation]
where the gradient is given by:
[Equation]
where z i,t = σ − 1 ( − U (cid:62) t V + x i ) .
We summarize the procedure in Algorithm 1 and show in Section [Equation] that under mild assumptions this procedure will be able to capture the underlying subspace structure in the parameter space with just one pass of the data [Equation]
2. T HEORETICAL RESULTS
We establish both asymptotic and non-asymptotic convergence properties for Algorithm 1.
The proof scheme is inspired by a series of previous works: Mairal et al.
; Kasiviswanathan et al.
Shalev-Shwartz et al.
(2012); Mardani et al.
; Shen et al.
We brieﬂy present the proof sketch, and refer readers to the literature for more details.
At each iteration i = 1 , 2 , ..., N , we sample ( x i , y i ) , and let U i , V i denote the intermediate U and V , to be differentiated from U t , V t which are the t -th columns of U, V .
For the proof feasibility, we assume that { ( x i , y i ) } Ni =1 are sampled i.i.d., and the subspace sequence { U i } Ni =1 lies in a compact set. Asymptotic Case: To estimate U , the Stochastic Gradient Descent (SGD) iterations can be seen as minimizing the approximate cost 1 N (cid:80)
Ni =1 g (cid:48) ( x i , y i , U, V ) , where g (cid:48) is a tight quadratic surrogate for Algorithm
Require:
Training data D
= { ( x i , y i ) } , target network depth K .
Ensure: The deep subspace network f
Set f
[0] ( x ) =
y and solve f [0] using Algorithm 1.
for k = 1 , . . . , K − 1 do // 1.
Subspace sketching based on the current subspace using Algorithm 1 :
[Equation]
Expand the layer using the reﬁned subspace as our new network :
[Equation]
end for return f = f [ K ] g based on the second-order Taylor approximation around U N − 1 .
Furthermore, g can be shown to be smooth, by bounding its ﬁrst-order and second-order gradients w.r.t.
each U t (similar to Appendixof Shen et al.
Following Mairal et al.
; Mardani et al.
(2015), it can then be established that, as N → ∞ , the subspace sequence { U i } Ni =1 asymptotically converges to a stationary-point of the batch estimator, under a few mild conditions .
We can sequentially show: 1) (cid:80) Ni =1 g (cid:48) ( x i , y i , U i , V i ) asymptotically converges to (cid:80) Ni =1 g ( x i , y i , U i , V i ) , according to the quasi-martingale property in the almost sure sense, owing to the tightness of g (cid:48) ; 2) the ﬁrst point implies convergence of the associated gradient sequence, due to the regularity of g ; 3) g t ( x i , y i , U, V ) is bi-convex for the block variables U t and V .
Non-Asymptotic Case: When N is ﬁnite, Mardani et al.
(2013) asserts that the distance between successive subspace estimates will vanish as fast as o (1 /i ) :
(cid:107) U i − U i − 1 (cid:107) F ≤ Bi , for some constant B that is independent of i and N .
Following Shen et al.
(2016) to exploit the unsupervised formulation of regret analysis as in Kasiviswanathan et al.
(2012); Shalev-Shwartz et al.
(2012), we can similarly obtain a tight regret bound that will again vanish if N → ∞ .
S UBSPACE NETWORK VIA HIERARCHICAL SKETCHING AND REFINEMENT
The single layer model in (1) has very limited power to capture the highly nonlinear regression relationships, as the parameters are linearly linked to the subspace except for a ReLU operation.
However, the single-layer procedure in Algorithm 1 has provided a building block, based on which we can develop an efﬁcient algorithm to train a deep subspace network in a greedy fashion.
We thus propose a network expansion procedure to overcome such limitation.
After we obtain the parameter subspace U and sketch V for the single-layer case (1), we project the data points by ¯ x = ReLU( UV x ) .
A straightforward idea of the expansion is to use (¯ x, y ) as the new samples to train another layer.
Let f [ k − 1] denote the network structure we obtained before the k -th expansion starts, k = 1 , 2 , ..., K − 1 , the expansion can recursively stack more ReLU layers:
[Equation]
However, we observe that simply stacking layers by repeating (4) many times can cause substantial information loss and degrade the generalization performance, especially since our training is layerby-layer without “looking back” (i.e., top-down joint tuning).
Inspired by deep residual networks by
He et al.
(2016) that exploit “skip connections” to pass lower-level data and features to higher levels, we concatenate the original samples with the newly transformed, censored outputs after each time of expansion, i.e., reformulating ¯ x = [ReLU( UV x ); x ] (similar manners could be found in Zhou &
Feng (2017)).The new form after the expansion is given below:
[Equation]
We summarize the network expansion process in Algorithm 2.
The architecture of the resulting subspace network (SN) is illustrated in Figure 1.
Compared to the single layer model (1), SN [Figure] gradually reﬁnes the parameter subspaces by multiple stacked nonlinear projections.
It is expected to achieve superior performance due to the higher learning capacity, and can also be viewed as a gradient boosting method.
Meanwhile, the layer-wise low-rank subspace structural prior would improve generalization further compared to the naive multi-layer networks.
E XPERIMENT 4. S IMULATIONS ON S YNTHETIC D ATA
Subspace recovery in a single layer model.
We ﬁrst evaluate the recovery of the original subspace by using the proposed Algorithm 1 on synthetic data.
We generated X ∈ R N × d , U ∈ R T × r and V ∈ R r × d , all as i.i.d.
random Gaussian matrices.
The target matrix Y ∈ R N × T was then synthesized using (1).
We set N = 5 , 000 , d = 200 , T = 100 r = 10 , and (cid:15) as a N (0 , 3 2 ) random noise.
Figure 2a displays the plot of subspace difference between the ground-truth U and the learned subspace U i throughout the iterations, i.e., (cid:107) U − U
i (cid:107) / (cid:107) U (cid:107) w.r.t.
i .
This result veriﬁes that Algorithmis able to correctly ﬁnd and smoothly converge to the underlying low-rank subspace of the synthetic data.
The objective values throughout the online training process of Algorithm 1 are plotted in Figure 2b.
We further show the plot of iteration-wise subspace differences , deﬁned as (cid:107) U i −
U i − 1 (cid:107) F / (cid:107) U (cid:107) , in Figure 2c, which complies with the o (1 /t ) result in our non-asymptotic analysis.
On the other hand, after one-pass of the entire data, we would also like to verify whether the model weights can be reliably recovered by multiplying the subspace and projection matrix.
We randomly pick two tasks and plot recovered weights versus ground truth.
Correlation between recovered weights and true weights for all 100 tasks are also computed.
As can be checked from Figure 4 in the Appendix, we display the distribution of computed correlations and scatter plot for two chosen tasks, and see that by passing data only once, our algorithm is able to recover the weights highly accurately (the majority of the predicted weights value having correlations with ground truth of above 0.9).
Multi-layer subspace network We re-generated synthetic data by repeatedly applying (1) for three times, each time followed the same setting as the single-layer model.
A three-layer SN was then learned using Algorithm 2.
As one simple baseline, a multi-layer perceptron (MLP) is trained, whose three hidden layers have the same dimensions as the three ReLU layers of the SN.
Inspired by Xue et al.
; Sainath et al.
; Wang et al.
, we then applied low-rank matrix factorization to each layer of MLP, with the same desired rank r , creating the factorized MLP (f-MLP) baseline that has the identical architecture (including both ReLU hidden layers and linear bottleneck layers) to SN.
We further re-trained f-MLP on the same data from end to end, leading to another retrained factorized MLP (rf-MLP) baseline.
We adopt the same deﬁnition of subspace difference as in the single-layer case for evaluation.
Table 1 compares the U subspace differences of the learned weight matrices, for each layer of SN, f-MLP, and rf-MLP, where SN achieves mostly the lowest differences.
[Table]
Beneﬁts of Going Deep.
We re-generated synthetic data again in the same way as the ﬁrst experiment; yet different from the ﬁrst one, we now aim to show that a deep SN will boost performance over
[Table]
[Figure] v e [Figure]
[Figure] [Figure] v e r age [Figure]
a li
[Figure]
single-layer subspace recovery, even in this simplest synthetic setting.
We will also compare SN with state-of-art approaches single task models and multi-task models.
Least square is treated as a naive baseline, while ridge (LS + (cid:96) 2 ) and lasso (LS + (cid:96) 1 ) regressions are considered for shrinkage or variables selection purpose; Censor regression (Censor), also known as Tobit model, predicts bounded targets with higher accuracy, e.g., Berberidis et al.
Multi-task models with regularizations on trace norm (Multi-trace) and (cid:96) 2 , 1 norm (Multi- (cid:96) 2 , 1 ) has been demonstrated to be successful on simultaneous structured/sparse learning, e.g., Yang et al.
; Zhang et al.
[Figure] [Table] We performed 10-fold random-sampling validation on the same dataset, i.e., randomly splitting into training and validation data 10 times.
For each split, we ﬁt model on training data and evaluated performance on validation data.
Mean of normalized mean square error (MNMSE) across all tasks is obtained as the overall performance for each split.
For methods without hyper parameters (least square and censor regression), an average of MNMSE (AMNMSE) for 10 splits is regarded as the ﬁnal performance; for methods with tunable parameters, e.g., λ in lasso, we perform a grid search on λ values and choose the optimal one.
We consider different splitting size with training samples containing [40%, 50%, 60%, 70%, 80%] of all the samples.
We grow the number of layers in SN from 2 to 20, and plot the errors in Figure 3a.
SN will steadily improves its performance with more layers, until reaching a plateau (here at around 5 layers, because of the simple underlying data distribution).
The observation is quite consistent among different splits, and also in more experiments. Table 2 further compares SN (with 20 layers), the performance for all approaches, from which we can conﬁrm that: (1) the censored model signiﬁcantly reduces the bias compared to uncensored models; (2) the multi-task model captures task relatedness and outperforms single task models; (3) by combining the best of both worlds, SN outperforms all competitors.
Computation speed All experiments are run on the same machine (1 x Six-core Intel Xeon E5v3 [3.50GHz], 12 logic core, 128 GB RAM), without GPU acceleration.
The running time for completing a 10-fold validation for all splitting percentages for synthetic data are given in Table3.
SN improves generalization performance without signiﬁcant loss of training speed, which is another advantage over other
4. E XPERIMENTS ON R EAL DATA
The standard deviation of 10 trials is given in parenthesis.
[Table]
We evaluated SN in a real clinical setting with the goal to build models to predict important clinical scores representing the subject’s cognitive status and signaling the progression of Alzheimer’s disease (AD), from structural Magnetic Resonance Imaging (sMRI) data
The ADNI phase 1 cohort 3 is used.
In the experiments, we worked with 1.5 Tesla structural MRI collected from the subjects at baseline, and performed cortical reconstruction and volumetric segmentations with the FreeSurfer following the procotol in Jack et al.
For each MRI image, we extracted 138 features representing the cortical thickness and surface areas of region-of-interests (ROIs) using the Desikan-Killiany cortical atlas Desikan et al.
After preprocessing, we obtained a dataset containing 670 samples and 138 feature dimensions.
These imaging features were used to predict a set of 30 clinical scores including ADAS scores Rosen et al.
(1984) at baseline and future (6 months from baseline), baseline Logical Memory from Wechsler
Memory Scale IV Scale—Fourth (2009), Neurobattery scores (i.e. immediate recall total score and Rey Auditory Verbal Learning Test scores), and the Neuropsychiatric Inventory (NPI) Cummings (1997) at baseline and future. A 6-layer SN with rank 5 per layer is employed by default.
σ 2 Calibration In synthetic formulation we assume that noise variance σ 2 is the same across all tasks, which is unnecessarily true in real cases.
To deal with heterogeneous σ 2 among tasks, we add a calibration step in our optimization process, where we estimate task-speciﬁc ˆ σ 2 t using (cid:107) y − ˆ y (cid:107) 22 /N before ReLu transformation, as the input for next layer and repeatomh layer-wise.
We compare performance of both non-calibrated and calibrated methods in next subsection.
Performance We adopt the same metrics as the synthetic experiments to evaluate performance on real data.
Different from synthetic data where the low-rank structure is predeﬁned, for real data, we try different tuning parameters as well as different low-rank structures.
Table 4 compares the performances between σ 2 non-calibrated versus calibrated models.
We observe a clear improvement by assuming different σ 2 across tasks.
Table 5 shows the results for all comparison methods, with SN outperforming all else.
Figure 3b shows the SN performance growth with increasing the number of layers.
Figure 3c further reveals the SN performance using varying rank estimations, when the groundtruth rank is unavailable in real data.
As expected, the U-shape curve suggests that an overly low rank may not be informative enough to recover the original weight space, while a high rank structure cannot enforce as strong a structural prior. [Table]
C ONCLUSIONS AND FUTURE WORK
In this paper we proposed a Subspace Network , an efﬁcient deep modeling approach for nonlinear multi-task censored regression, where each layer of the subspace network performs a multitask censored regression to improve upon the predictions from the last layer via sketching a lowdimensional subspace to perform knowledge transfer among learning tasks.
We show that under mild assumptions, for each layer we can recover the parametric subspace using only one pass of training data.
We demonstrate empirically that the subspace network can quickly capture correct parameter subspaces, and outperforms state-of-the-arts in predicting neurodegenerative clinical scores from brain imaging.
Based on similar formulations, the proposed method can be easily extended to cases where the targets have nonzero bounds, or both lower and upper bounds.
Appendix
Algorithm
Require: Training data ( x i , y i ) , U − , step size η ,
Ensure: sketch V
Initialize V − at random // 1.
Perform gradient step and update the current solution of V .
for t = 1 , . . . , T do
Compute z i,t = σ − 1 ( − (cid:0) U − t (cid:1) (cid:62) V x i )
Compute the gradient for y t : ∇ g t ( x i , y i,t ; U − , V + ) =
  − y i,t − ( U − t ) (cid:62)
V x i σ
U − t x (cid:62) i y i,t ∈ (0 , ∞ ) φ ( z i,t ) σ [1 − Q ( z i,t )]
U − t x (cid:62) i y i,t =end for // 2.
Update the current sketch V −
V + = V − − η (cid:34) T (cid:88) t
= ∇ g t ( x, y t ; U − , V + ) + λV (cid:35)
Set V − [-−   
 −−-]= V +
[Figure]
[Figure]
i c t ed v a l ue (
[Figure]
[Figure]
[Figure]