Fine-tuning Factor Augmented Neural Lasso for Heterogeneous Environments

1. Problem Formulation

The paper addresses the challenge of transfer learning in high-dimensional nonparametric regression settings where the target domain has limited sample size ($n_Q$) compared to the source domain ($n_P$), and where $p \gg n_Q$ (number of features exceeds target samples).

The problem involves two specific types of distribution shifts:

1. Covariate Shift: The marginal distribution of covariates differs between domains ($P \neq Q$), specifically manifested through different factor loading matrices ($B_P \neq B_Q$) in a latent factor model.
2. Posterior Shift: The conditional distribution of the response differs ($g_P \neq g_Q$), meaning the functional relationship between inputs and outputs changes.

Key Challenge: Existing methods struggle to simultaneously handle high-dimensional dependent covariates, nonparametric function estimation, variable selection, and these dual distribution shifts while avoiding "negative transfer" (where source data harms target performance).

Model Assumptions:

• Factor Structure: Covariates follow a low-rank factor model: $x = Bf + u$, where $f$ are latent factors and $u$ are idiosyncratic components.
• Regression Structure:
  • Source: $y_P = g_P(x_{J_P}) + \epsilon_P$ (Sparse nonparametric).
  • Target: $y_Q = g_Q(f_Q, u_{J_Q}) + \epsilon_Q$ (Factor-augmented sparse nonparametric).
• Transfer Mechanism: The target function is decomposed as a residual fine-tuning of the source function:
  $$g_Q(f_Q, u_{J_Q}) = h\left(f_Q, u_J, g_P(x_{J_P})\right)$$
  Here, $h$ is a low-complexity residual function, and $g_P$ is a "frozen" pre-trained source function.

2. Methodology: Fine-tuning FAN-Lasso

The authors propose the Fine-tuning Factor Augmented Neural Lasso (FAN-Lasso), a framework integrating three components:

A. Diversified Projection for Factor Estimation (Handling Covariate Shift)

To estimate latent factors $f_Q$ in the target domain when $n_Q$ is small:

• Diversified Projection Matrix: Instead of relying solely on target eigenvectors (which are noisy), the method constructs a diversified projection matrix $W$ using top eigenvectors of a pooled covariance matrix ($\hat{\Sigma}_A$) of source and target data.
• Transfer Around Boundary (TAB): A model-selection strategy is employed. If the difference between target and pooled covariance is below a threshold $\delta$, the pooled estimator is used; otherwise, the target-only estimator is used. This ensures robustness against significant domain shifts.
• Surrogate Factors: Latent factors are estimated as $\hat{f} = p^{-1} W^\top x$.

B. Residual Fine-Tuning Decomposition (Handling Posterior Shift)

The core innovation is treating the target learning task as learning a residual function $h$ rather than the full function $g_Q$.

• Frozen Source: The pre-trained source model $\hat{g}_P$ (trained via FAST-NN on source data) is treated as a fixed input feature.
• Target Network: A new neural network $h$ is trained on target data using inputs: $\{ \hat{f}_Q, \text{selected idiosyncratic features}, \hat{g}_P(x_Q) \}$.
• Benefit: This reduces the complexity of the target learning task from estimating a complex $g_Q$ to estimating a simpler residual $h$, leveraging the "knowledge" already captured in $g_P$.

C. Sparse Nonparametric Variable Selection

• FAST-NN Architecture: The method utilizes Factor Augmented Sparse Throughput Neural Networks.
• L0 Relaxation: A continuous relaxation of the $L_0$ penalty (clipped $L_1$ function $\psi_\tau$) is applied to the weights connecting raw covariates to the network. This enables automatic variable selection (identifying relevant $J_P$ and $J_Q$) within the deep learning framework, preventing the curse of dimensionality.

3. Key Contributions

1. Theoretical Framework for High-Dimensional Transfer: The paper provides the first minimax-optimal theoretical analysis of fine-tuning in high-dimensional nonparametric settings with variable selection.
2. Minimax Optimal Rates: The authors derive the excess risk bounds for the FAN-Lasso estimator. The convergence rate is shown to be:
   $$\left( \frac{\log(n_P + n_Q)}{n_P + n_Q} \right)^{\frac{2\gamma_P}{2\gamma_P+1}} + \left( \frac{\log n_Q}{n_Q} \right)^{\frac{2\gamma}{2\gamma+1}} + \frac{\log p}{n_Q}$$
   Where $\gamma_P$ is the complexity of the source function and $\gamma$ is the complexity of the residual function.
3. Statistical Acceleration Conditions: The theory proves that fine-tuning yields statistical acceleration over single-task learning if:
   • Source sample size is large ($n_P \gg n_Q$).
   • The source function is more complex than the residual ($\gamma_P < \gamma$).
   • Under these conditions, the rate depends on the joint sample size for the complex part and only the target sample size for the residual.
4. Robustness to Negative Transfer: The estimator is proven to be "oracle-robust." If the source data offers no advantage (e.g., $n_P$ is small or the residual is as complex as the target), the method automatically degrades to the optimal rate achievable without source data, without requiring prior knowledge of the relationship between domains.
5. Handling Distribution Shifts: The framework explicitly addresses both covariate shifts (via diversified projection) and posterior shifts (via residual decomposition), a combination rarely addressed in unified theoretical frameworks.

4. Results

Theoretical Results

• Upper Bounds: Theorem 3 establishes that the FAN-Lasso achieves the minimax optimal rate up to logarithmic factors.
• Lower Bounds: Theorem 4 provides a matching minimax lower bound, confirming the sharpness of the derived rates.
• Factor Transfer: Theorem 1 proves that the diversified projection matrix constructed from pooled data allows for accurate recovery of target latent factors even when $n_Q$ is very small, provided $n_P$ is large and the loading matrices are sufficiently similar.

Empirical Results

• Simulation Studies:
  • Covariate Shift: The transfer factor estimator ($\hat{W}_{TL}$) significantly outperforms target-only and source-only estimators in recovering the true loading space, especially when $n_Q$ is very small (e.g., $n_Q=7$).
  • Posterior Shift: FAN-Lasso consistently achieves the lowest Root Mean Squared Error (RMSE) compared to baselines (Vanilla-NN, FAST-NN, Fine-tuning Vanilla-NN). It converges to the performance of an "Oracle" estimator (which knows the true residual structure) as $n_Q$ increases.
• Real Data Application:
  • Applied to the Communities and Crime dataset (predicting violent crime rates).
  • Scenario: Transfer from high-density (urban) communities (Source, $n \approx 1154$) to low-density (rural) communities (Target, $n \approx 839$).
  • Outcome: FAN-Lasso achieved the lowest test RMSE ($0.1285$), outperforming the fine-tuned Vanilla-NN ($0.1373$) and source-only baselines, validating its effectiveness in real-world heterogeneous environments.

5. Significance

This paper makes a significant contribution to the theoretical understanding of Transfer Learning and Deep Learning in high-dimensional statistics.

• Bridging Theory and Practice: It moves beyond linear or parametric assumptions to provide rigorous guarantees for deep neural networks in transfer learning.
• Parameter Efficiency: It offers a theoretical justification for parameter-efficient fine-tuning (PEFT) methods, showing that learning a low-complexity residual is statistically optimal when source data is abundant.
• Robustness: The automatic safeguard against negative transfer is crucial for practical applications where the similarity between source and target domains is unknown.
• Scalability: By combining factor models with sparse neural networks, the method provides a scalable solution for high-dimensional data where traditional nonparametric methods fail due to the curse of dimensionality.

In summary, the Fine-tuning FAN-Lasso provides a unified, theoretically grounded, and empirically robust framework for adapting pre-trained models to new, data-scarce, high-dimensional environments with complex distribution shifts.