2023-NIPS《Joint Feature and Differentiable k-NN Graph Learning using Dirichlet Energy》

核心思想

论文提出了一种基于Dirichlet能量的深度无监督特征选择（Feature Selection, FS）框架，同时进行特征选择和可微分的k近邻（k-NN）图学习。其核心思想包括：

1. 联合特征选择与图学习：

   • 传统无监督特征选择方法（如Laplacian Score）依赖预计算的固定图结构，当数据包含大量无关特征时，图质量较差，影响特征选择的性能。论文提出通过Dirichlet能量联合优化特征选择和k-NN图结构，动态学习反映数据内在结构的图和信息性特征。
   • Dirichlet能量衡量特征在图结构上的局部平滑性（smoothness），用于识别与数据内在结构一致的重要特征。

2. 可微分框架：

   • 针对传统k-NN图构建中的不可微分排序问题，论文引入Gumbel Softmax技术解决离散特征选择的不可微分性，并利用**最优传输（Optimal Transport, OT）**理论实现可微分的k-NN选择。
   • 所有模块（特征选择器和图学习器）均算法化设计，具有可解释性和可扩展性，可作为其他图神经网络的前端模块。

3. 模块化设计：

   • 唯一特征选择器（Unique Feature Selector, UFS）：通过Gumbel Softmax生成离散特征选择矩阵，并通过正交化避免重复选择特征。
   • 可微分k-NN图学习器（Differentiable k-NN Graph Learner, DGL）：基于最优传输理论学习自适应k-NN图，优化图结构以反映选定特征子空间的内在结构。

4. 无监督特性：

   • 模型无需任务特定标签，适合多种下游任务（如分类、聚类），通过最小化Dirichlet能量选择通用特征。

目标函数

论文的目标函数旨在通过最小化Dirichlet能量联合优化特征选择和k-NN图学习，具体形式如下：

1. 总体目标函数：
   目标是通过选择特征子集$\hat{\mathbf{X}}\in {\mathbb{R}}^{n\times m}$和学习相似性矩阵$\mathbf{S}\in {\mathbb{R}}^{n\times n}$，最小化以下目标函数：

   $$\underset{\mathbf{F},\mathbf{S}}{\min }{\mathcal{L}}_{\text{obj}}(\hat{\mathbf{X}})=\mathrm{tr}({\hat{\mathbf{X}}}^T\mathbf{L}\hat{\mathbf{X}})+\alpha \Vert \mathbf{S}{\Vert }_F^2,$$

   约束条件为：

   • $\hat{\mathbf{X}}=\mathbf{X}\mathbf{F}$，其中$\mathbf{X}\in {\mathbb{R}}^{n\times d}$是输入数据，F∈{0,1}d×m\boldsymbol{F} \in \{0,1\}^{d \times m}F∈{0,1}d×m是二值特征选择矩阵，${\mathbf{F}}^T\mathbf{F}={\mathbf{I}}_m$确保选择$m$个唯一特征。
   • $\mathbf{S}{1}_n=1_n$，$s_{i,j}\ge 0$，$s_{i,i}=0$，确保相似性矩阵表示有效的k-NN图。
   • $\mathbf{L}=\mathbf{D}-\frac{\mathbf{S}+{\mathbf{S}}^T}{2}$是拉普拉斯矩阵，其中$\mathbf{D}$是对角度矩阵，$d_{i,i}={\sum }_{j=1}^n{s}_{i,j}$。

   其中：

   • $\mathrm{tr}({\hat{\mathbf{X}}}^T\mathbf{L}\hat{\mathbf{X}})={\sum }_{i=1}^m{\mathcal{L}}_{\text{dir}}({\hat{\mathbf{x}}}_i)$是Dirichlet能量，衡量选定特征在图上的平滑性。
   • $\alpha \Vert \mathbf{S}{\Vert }_F^2$是Tikhonov正则化项，防止$\mathbf{S}$退化为平凡解（如单位矩阵）。

2. 特征选择子问题：
   为选择$m$个特征，优化以下子问题：

   min⁡Ftr⁡(X^TLX^)s.t.X^=XF,F∈{0,1}d×m,FTF=Im. \min_{\boldsymbol{F}} \operatorname{tr}(\hat{\boldsymbol{X}}^T \boldsymbol{L} \hat{\boldsymbol{X}}) \quad \text{s.t.} \quad \hat{\boldsymbol{X}} = \boldsymbol{X} \boldsymbol{F}, \boldsymbol{F} \in \{0,1\}^{d \times m}, \boldsymbol{F}^T \boldsymbol{F} = \boldsymbol{I}_m. Fmin​tr(X^TLX^)s.t.X^=XF,F∈{0,1}d×m,FTF=Im​.

   • 使用Gumbel Softmax生成近似离散的特征选择向量${\hat{\mathbf{f}}}_i$：

     $${\hat{\mathbf{f}}}_i=\mathrm{softmax}\left(\frac{\log {\mathbf{w}}_i+{\mathbf{g}}_i}{T}\right),g_{i,j}=-\log (-\log u_{i,j}),u_{i,j}\sim \text{Uniform}(0,1),$$

     其中${\mathbf{w}}_i$是可学习参数，$T$是温度参数，逐渐退火以逼近离散分布。

   • 为确保特征唯一性，通过算法1（UFS）对$\hat{\mathbf{F}}$进行正交化处理：

     $$\mathbf{A}={\hat{\mathbf{F}}}^T\hat{\mathbf{F}}+ϵ{\mathbf{I}}_m,\mathbf{A}=\mathbf{P}\mathbf{\Lambda }{\mathbf{P}}^T,\mathbf{A}=\mathbf{L}{\mathbf{L}}^T,\mathbf{F}=\left[\begin{array}{c}{\mathbf{\Lambda }}^{1/2}{\mathbf{P}}^T\\ 0\end{array}\right]{\mathbf{L}}^{-1}.$$

3. 图学习子问题：
   为学习k-NN图的相似性矩阵$\mathbf{S}$，优化以下子问题：

   $$\underset{\mathbf{S}}{\min }\mathrm{tr}({\hat{\mathbf{X}}}^T\mathbf{L}\hat{\mathbf{X}})+\alpha \Vert \mathbf{S}{\Vert }_F^2\text{s.t.}\mathbf{S}{1}_n=1_n,s_{i,j}\ge 0,s_{i,i}=0.$$

   • 每行$s^i$独立优化，转换为：

     $$\underset{s^i}{\min }\frac{1}{2}{\Vert s^i+\frac{{\mathbf{e}}^i}{2{\alpha }_i}\Vert }_2^2\text{s.t.}{s}^i{1}_n=1,s_{i,j}\ge 0,s_{i,i}=0,$$

     其中${\mathbf{e}}^i$是距离向量，$e_{i,j}=\Vert {\hat{\mathbf{x}}}^i-{\hat{\mathbf{x}}}^j{\Vert }_2^2$，${\alpha }_i$是样本特定的正则化参数。

   • 通过KKT条件解得：

     $$s_{i,j}={\left(\frac{1}{k}+\frac{1}{k}\frac{{\mathbf{e}}^i{\mathbf{\delta }}_i^{(k)}}{2{\alpha }_i}-\frac{e_{i,j}}{2{\alpha }_i}\right)}_{+},$$

     其中${\mathbf{\delta }}_i^{(k)}$是标识${\mathbf{e}}^i$中前$k$个最小值的选择向量。

   • 为确保$s^i$只有$k$个非零元素，选择${\alpha }_i=\frac{1}{2}(k{\mathbf{e}}^i{\mathbf{\xi }}_i^{(k+1)}-{\mathbf{e}}^i{\mathbf{\delta }}_i^{(k)})$，最终解为：

     $$s_{i,{\sigma }_j}=\frac{{\mathbf{e}}^i{\mathbf{\xi }}_i^{(k+1)}-e_{i,{\sigma }_j}}{k{\mathbf{e}}^i{\mathbf{\xi }}_i^{(k+1)}-{\mathbf{e}}^i{\mathbf{\delta }}_i^{(k)}}\cdot \text{Bool}(1\le j\le k),$$

     其中${\mathbf{\xi }}_i^{(k+1)}$标识${\mathbf{e}}^i$中第$k+1$个最小值。

4. 可微分k-NN选择：

   • 传统排序算法不可微，论文通过最优传输（OT）将k-NN选择转化为可微分问题：

     $$\underset{\Gamma }{\min }⟨\mathbf{C},\mathbf{\Gamma }⟩+\gamma \sum _{i,j}{\Gamma }_{i,j}\log {\Gamma }_{i,j}\text{s.t.}\mathbf{\Gamma }{1}_{k+2}=\mathbf{\mu },{\mathbf{\Gamma }}^T{1}_n=\mathbf{\nu },{\Gamma }_{i,j}\ge 0,$$

     其中$\mathbf{C}$是成本矩阵，$c_{i,j}=(e_i-j+1{)}^2$，$\gamma$是熵正则化参数，$\mathbf{\mu }$和$\mathbf{\nu }$是概率分布。

   • 通过迭代Bregman投影求解$\mathbf{\Gamma }$，得到${\mathbf{\delta }}_i^{(k)}$和${\mathbf{\xi }}_i^{(k+1)}$：

     $${\mathbf{\delta }}_p^{(k)}=n\sum _{i=1}^k{\Gamma }_i,{\mathbf{\xi }}_p^{(k+1)}=n{\Gamma }_{k+1}.$$

优化过程

优化过程通过联合优化特征选择矩阵$\mathbf{F}$和相似性矩阵$\mathbf{S}$实现，采用基于梯度的深度学习框架（PyTorch实现），具体步骤如下：

1. 初始化：

   • 初始化特征选择参数$\mathbf{W}$（用于Gumbel Softmax）和图学习参数。
   • 设置超参数：学习率η∈{10−4,10−3,10−2,10−1,100,101}\eta \in \{10^{-4}, 10^{-3}, 10^{-2}, 10^{-1}, 10^0, 10^1\}η∈{10−4,10−3,10−2,10−1,100,101}，k-NN参数k∈{1,2,3,4,5}k \in \{1, 2, 3, 4, 5\}k∈{1,2,3,4,5}，熵正则化参数γ∈{10−3,10−2,10−1}\gamma \in \{10^{-3}, 10^{-2}, 10^{-1}\}γ∈{10−3,10−2,10−1}，温度参数$T$（初始值高，逐渐退火）。

2. 特征选择模块（UFS）：

   • 生成特征选择向量：通过Gumbel Softmax计算${\hat{\mathbf{f}}}_i$，逼近离散分布。
   • 正交化处理：
     • 计算$\mathbf{A}={\hat{\mathbf{F}}}^T\hat{\mathbf{F}}+ϵ{\mathbf{I}}_m$，$ϵ$为小正数（如$10^{-6}$），确保$\mathbf{A}$正定。
     • 进行Cholesky分解$\mathbf{A}=\mathbf{L}{\mathbf{L}}^T$，计算$\mathbf{F}=\hat{\mathbf{F}}({\mathbf{L}}^{-1}{)}^T$，确保${\mathbf{F}}^T\mathbf{F}={\mathbf{I}}_m$，避免重复选择特征。
   • 计算选定特征：$\hat{\mathbf{X}}=\mathbf{X}\mathbf{F}$。

3. 图学习模块（DGL）：

   • 距离计算：基于选定特征$\hat{\mathbf{X}}$，计算距离矩阵$\mathbf{E}$，$e_{i,j}=\Vert {\hat{\mathbf{x}}}^i-{\hat{\mathbf{x}}}^j{\Vert }_2^2$。
   • 可微分k-NN选择：
     • 初始化OT参数：${\mathbf{\mu }}_i=\frac{1}{n}$，${\nu }_j=\frac{1}{n}$（$j\le k+1$）或$\frac{n-k-1}{n}$（$j=k+2$），$c_{i,j}=(e_i-j+1{)}^2$。
     • 通过迭代Bregman投影（200次迭代）求解熵正则化OT问题，得到传输计划$\mathbf{\Gamma }$。
     • 计算选择向量${\mathbf{\delta }}_i^{(k)}$和${\mathbf{\xi }}_i^{(k+1)}$。
   • 计算相似性矩阵：根据公式（9）计算$s_{i,{\sigma }_j}$，构建$\mathbf{S}$，并计算对称化拉普拉斯矩阵$\mathbf{L}=\mathbf{D}-\frac{\mathbf{S}+{\mathbf{S}}^T}{2}$。

4. 联合优化：

   • 使用Adam优化器（1000个epoch）最小化目标函数$\mathrm{tr}({\hat{\mathbf{X}}}^T\mathbf{L}\hat{\mathbf{X}})+\alpha \Vert \mathbf{S}{\Vert }_F^2$。
   • 通过反向传播更新$\mathbf{W}$（特征选择参数）和图学习参数，联合优化$\mathbf{F}$和$\mathbf{S}$。
   • 温度参数$T$按退火策略逐渐减小，增强$\hat{\mathbf{F}}$的离散性。

5. 收敛与输出：

   • 监控目标函数值，收敛后输出选定特征$\hat{\mathbf{X}}$和图结构$\mathbf{S}$。
   • 对于下游任务，使用$\hat{\mathbf{X}}$进行分类（随机森林）、聚类（k-means）或重构（自编码器）。

目标函数的局限性及改进建议

尽管论文提出的目标函数和优化框架在理论和实验上表现出色，但仍存在以下局限性，并可通过以下改进建议进一步优化：

1. 局限性：对超参数的敏感性：

   • 问题：目标函数涉及多个超参数（如学习率$\eta$、k-NN参数$k$、熵正则化参数$\gamma$、特征数量$m$）。虽然实验表明模型对$\eta$、$k$、$\gamma$的敏感性较低，但$m$（选择特征数量）对性能影响较大，需通过交叉验证确定，增加了计算成本。
   • 改进建议：
     • 引入自适应特征数量选择机制。例如，基于Dirichlet能量的变化率自动确定$m$，当新增特征导致Dirichlet能量降低幅度小于阈值时停止选择。
     • 使用贝叶斯优化或元学习方法自动搜索超参数空间，减少人工调参成本。

2. 局限性：噪声鲁棒性有限：

   • 问题：尽管模型通过联合优化特征和图结构提高了对噪声的鲁棒性，但在高噪声数据集（如Madelon）上，初始图结构可能仍受无关特征影响，导致次优解。
   • 改进建议：
     • 在目标函数中引入噪声鲁棒性正则化项，如基于鲁棒统计量（如中位数而非均值）计算距离矩阵$\mathbf{E}$，或在Dirichlet能量中加入异常值检测机制。
     • 采用多阶段训练策略：先用简单方法（如Laplacian Score）预选特征，初步过滤噪声，再进行联合优化。

3. 局限性：计算复杂度：

   • 问题：可微分k-NN选择依赖最优传输和Bregman投影（200次迭代），对大规模数据集（高$n$）计算成本较高。此外，特征选择矩阵$\mathbf{F}$的正交化涉及Cholesky分解，复杂度为$O(m^3)$。
   • 改进建议：
     • 使用近似最优传输算法（如Sinkhorn算法的快速变体）降低OT计算复杂度。
     • 对于正交化，探索更高效的矩阵分解方法（如QR分解）或近似正交化策略。
     • 引入并行化或分布式计算，特别是在GPU上优化Bregman投影的迭代过程。

4. 局限性：图结构的局部性：

   • 问题：k-NN图只考虑局部邻居，可能忽略全局结构信息，导致在复杂数据分布（如非流形结构）上性能下降。
   • 改进建议：
     • 结合全局图结构学习方法，如在目标函数中加入谱聚类正则化项，鼓励$\mathbf{S}$捕捉全局簇结构。
     • 引入多尺度k-NN图，通过优化多个$k$值（如$k=5,10,20$）的加权组合，平衡局部和全局信息。

5. 局限性：可解释性有限：

   • 问题：虽然模块设计具有算法化可解释性，但最终选定特征和图结构的解释仍依赖Dirichlet能量，缺乏直观的语义解释。
   • 改进建议：
     • 在特征选择后，增加可视化模块（如t-SNE或UMAP）分析选定特征的语义分布，提供直观解释。
     • 引入基于注意力机制的特征重要性评分，结合Dirichlet能量生成特征排名的解释性报告。

6. 局限性：对特征标准化假设的依赖：

   • 问题：目标函数假设特征具有零均值和单位方差（$1_n^T{\mathbf{x}}_i=0$，${\mathbf{x}}_i^T{\mathbf{x}}_i=1$），在实际数据中可能不成立，需额外预处理。
   • 改进建议：
     • 在目标函数中加入自适应标准化项，自动学习特征的均值和方差，减少预处理依赖。
     • 探索非标准化特征的Dirichlet能量变体，如基于协方差矩阵调整$\mathbf{L}$的定义。

实验结果与验证

论文在合成数据集（Blobs、Moons、Circles）和12个真实数据集（包括文本、生物、图像和人工数据）上验证了模型的有效性：

• 合成数据集：展示了模型在噪声干扰下仍能有效选择信息性特征和学习准确的k-NN图（见Figure 3）。
• 真实数据集：
  • 分类任务（Table S1）：在Madelon（0.90±0.01）、Yale（0.81±0.05）、GLIOMA（0.81±0.09）等数据集上优于其他方法（如CAE、UDFS）。
  • 聚类任务（Table S2）：在Jaffe（0.87±0.06）、GLIOMA（0.75±0.08）等数据集上表现优异。
  • 重构任务（Table S5）：在Jaffe（0.22±0.01）、PROSTATE（0.73±0.13）等数据集上RMSE较低。
• 消融研究（Table S4、S5）：证明UFS和DGL模块对性能的贡献，去除DGL后性能显著下降（如Madelon分类从0.90降至0.50）。
• 参数敏感性（Figure S3）：模型对$\eta$、$k$、$\gamma$鲁棒，但对$m$敏感，需谨慎选择。

总结

论文通过Dirichlet能量提出了一种新颖的深度无监督特征选择框架，联合优化特征选择和k-NN图学习，利用Gumbel Softmax和最优传输技术实现可微分优化。其目标函数清晰，优化过程高效，实验结果验证了其在多种数据集上的优越性能。然而，模型对超参数$m$的敏感性、计算复杂度以及对特征标准化的依赖是主要局限性。未来可通过自适应特征选择、噪声鲁棒性增强和高效算法优化进一步改进，使其更适合大规模、复杂数据场景。