使用 TensorFlow.js 实现实时 tSNE 可视化

文 / Google 苏黎世软件工程实习生 Nicola Pezzotti

近年来，t 分布随机邻域嵌入 (tSNE) 算法已成为高维数据探索性数据分析中最常用和最具洞察力的技术之一。tSNE 用于解释 TensorFlow Embedding Projector 和 TensorBoard 等工具中的深度神经网络输出，它的强大特性之一是，只需对参数进行最少量的调整即可显示不同规模的高维数据点聚类。尽管有这些优点，但受限于计算复杂性，tSNE 算法只能应用于相对较小的数据集。为解决这一问题，已相继开发出几个 tSNE 演化版本（主要关注数据点之间相似性计算的可扩展性），但迄今为止，这些版本还不足以在可视化 tSNE 大数据集嵌入的演化时提供真正的交互式体验。

在 “Linear tSNE Optimization for the Web” （https://nicola17.github.io/tfjs-tsne-demo）一文中，我们提出了一种高度依赖现代图形硬件的全新 tSNE 方法。考虑到新方法的线性复杂性，我们的方法生成嵌入的速度比同类技术更快，甚至可以通过 WebGL 利用 GPU 功能在客户端的网络浏览器中执行。在这两种因素的共同作用下，大型高维数据集的实时交互式可视化得以实现。

此外，在 TensorFlow.js 系列中，我们将这项工作以开放源代码库的形式发布，希望更广的研究社区多加利用：

https://github.com/tensorflow/tfjs-tsne

用我们的技术实时演化完整 MNIST 数据集的 tSNE 嵌入。数据集包含 60000 个手写数字的图像：https://nicola17.github.io/tfjs-tsne-demo/

tSNE 的目标是对相似数据点的小 “邻域” 进行聚类，同时降低数据的整体维数，使其更易于可视化。换句话说，tSNE 目标函数测量相似数据的这些邻域在二维或三维空间中的保存程度，并相应地将它们组成聚类。

在此前的研究中，tSNE 目标函数的最小化作为 N 体模拟问题来执行，其中，点被随机放置在嵌入空间中，并且在每个点上施加两种不同类型的力。吸引力使点更靠近高维空间中最相似的点，而排斥力则使它们远离嵌入中的所有邻近点。

吸引力只作用于一小部分点（即相似的邻近点），而排斥力对所有点对都起作用。因此，tSNE 需要大量的计算和多次的目标函数迭代，使得数据集大小局限为仅数百个数据点。为改进蛮力解决方案，采用了 Barnes-Hut 算法近似计算排斥力和目标函数的梯度。这使得计算能力扩展到数万个数据点，但是计算 C++ 实现中的 MNIST 嵌入需要超过 15 分钟。

在论文中，我们针对上述扩展问题提出了一种解决方案，使用在 WebGL 中生成的纹理来近似计算目标函数的梯度。这一技术在每次最小化迭代中使用三通道纹理绘制 “排斥场”，其中 3 个分量被视为颜色并在 RGB 通道中绘制。为每个点获取排斥场，以表示该点产生的水平和垂直排斥力，以及用于归一化的第三分量。直观地说，归一化期用于确保移位的幅度与高维空间中的相似性度量值相匹配。此外，纹理的分辨率会根据情况自动改变，以便保持绘制的像素数量恒定。

三个函数的渲染，用于近似计算单点产生的排斥效应。在上图中，排斥力表示蓝色区域中的点被向左/向下推动，而红色区域中的点被向右/向上推动，而白色区域中的点不动

然后，将每个点的作用添加到 GPU 上，从而产生与下面 GIF 中呈现的近似排斥场的图像相似的纹理。相对于更常用的点间相互作用计算方法，这种新型排斥场方法对 GPU 更加友好。这是因为在 GPU 中可以同时且非常快速地计算多个点的排斥力。另外，我们在 GPU 中实现了点间吸引力的计算。

此动画显示了 tSNE 嵌入（左上）和用于使用归一化期（右上）、水平移位（左下）和垂直移位（右下）来近似计算梯度的标量场的演化

我们还将临时实现中的嵌入更新修改为一系列在 TensorFlow.js（一种 JavaScript 库，用于在网络浏览器中执行张量计算）中计算的标准张量操作。我们的方法已在 TensorFlow.js 系列中以开放源代码库形式发布，通过该方法，我们可以完全基于 GPU 来计算 tSNE 嵌入的演化，同时计算复杂性大大降低。

通过这种实现，以前需要 15 分钟计算的内容（在 MNIST 数据集上），现在可以在网络浏览器中实现实时可视化。此外，更大数据集的实时可视化也随之成为可能，此功能在分析深度神经输出时特别有用。我们研究的主要局限在于，这种技术目前只适用于 2D 嵌入。然而，2D 可视化通常比 3D 可视化更受欢迎，因为它们需要更多的交互来有效理解聚类结果。

未来工作

我们认为，在浏览器中运行快速、交互式 tSNE 实现将增强数据分析系统开发者的能力。我们特别感兴趣的是探索如何利用我们的实现来解释深度神经网络。另外，我们的实现展示了如何在 GPU 计算中运用横向思维（使用 RGB 纹理近似计算梯度）来显著加快算法计算速度。

今后，我们将探索如何拓展这种梯度近似法的应用范围，使其不仅用于加快其他降维算法的速度，而且还可以使用 TensorFlow.js 在网络浏览器中实现其他 N 体模拟。