[论文解读] Learning identifiable and interpretable latent models of high-dimensional neural activity using pi-VAE

Ding Zhou, Xue-Xin Wei|arXiv (Cornell University)|Nov 9, 2020

Neural dynamics and brain function参考文献 69被引用 49

一句话总结

pi-VAE 将可辨识的 VAE 概念与任务变量条件化相结合，在高维神经数据中学习可解释、可辨识的潜在结构，相较于标准 VAE 和调谐曲线模型，提升拟合度和潜在结构的可解释性。

ABSTRACT

The ability to record activities from hundreds of neurons simultaneously in the brain has placed an increasing demand for developing appropriate statistical techniques to analyze such data. Recently, deep generative models have been proposed to fit neural population responses. While these methods are flexible and expressive, the downside is that they can be difficult to interpret and identify. To address this problem, we propose a method that integrates key ingredients from latent models and traditional neural encoding models. Our method, pi-VAE, is inspired by recent progress on identifiable variational auto-encoder, which we adapt to make appropriate for neuroscience applications. Specifically, we propose to construct latent variable models of neural activity while simultaneously modeling the relation between the latent and task variables (non-neural variables, e.g. sensory, motor, and other externally observable states). The incorporation of task variables results in models that are not only more constrained, but also show qualitative improvements in interpretability and identifiability. We validate pi-VAE using synthetic data, and apply it to analyze neurophysiological datasets from rat hippocampus and macaque motor cortex. We demonstrate that pi-VAE not only fits the data better, but also provides unexpected novel insights into the structure of the neural codes.

研究动机与目标

• 激发对既灵活又可解释的高维神经数据模型的需求。
• 开发一个生成框架，联合建模潜在结构及其与任务变量的关系。
• 通过引入标签先验和泊松观测噪声实现可识别性和可解释性。
• 证明在神经数据集上，pi-VAE 相较于基线可获得更好的数据拟合和更易解释的潜在结构。

提出的方法

• 定义生成模型 p_theta(x,z|u) = p_f(x|z) p_{T,lambda}(z|u)。
• 将标签先验 p_{T,lambda}(z|u) 建模为条件独立的指数族分布，其自然参数 lambda(u) 由神经网络参数化。
• 将 p_f(x|z) 表示为以发射率 f(z) 实现的泊松观测，通过广义可逆流（GIN）实现，以处理高维输出。
• 将 GIN 扩展为将 m 维 z 映射到 n 维 x，确保 f 是单射。
• 采用带有 q(z|x,u) ∝ q_phi(z|x) p_{T,lambda}(z|u) 的可辨识的 VAE 风格推理方案，并使用高斯近似以便可训练。
• 在温和条件下证明可识别性，并用合成数据和电生理数据集进行验证。

实验结果

研究问题

• RQ1 pi-VAE 是否能够在利用任务变量作为标签的同时，恢复高维神经记录中的可解释且可辨识的潜在结构？
• RQ2 将标签先验引入是否相对于标准 VAE 或调谐曲线方法改善模型拟合度和潜在解耦？
• RQ3 在使用 pi-VAE 时，真实神经数据（如海马体 CA1 和 M1/PMd）会呈现出何种潜在几何结构？
• RQ4 在神经数据集上，相对于传统方法，pi-VAE 在解码和编码任务中的表现如何？

主要发现

• pi-VAE 在对保留数据的拟合方面优于 VAE 和调谐曲线模型，测量指标是测试数据的边际对数似然更高。
• pi-VAE 在解码任务相关变量（如到达方向）方面比调谐曲线模型更准确，尤其是在运动的早期阶段。
• 来自 pi-VAE 的潜在空间显示出可解释的几何结构，分离方向信息与时间动态（例如，在猴子到达数据中，前两个潜在维度编码方向，其他维度捕捉轨迹演变）。
• 在大鼠 CA1 数据上，pi-VAE 显示与轨道几何对齐的潜在流形，并呈现与位置细胞属性一致的方向分离，且时间结构与 θ 节律（约 10 Hz）相关。
• 来自 pi-VAE 的潜在结构比其他方法（UMAP、PfLDS、LFADS、分解 PCA）更解耦且在物理意义上更有意义。
• 标签先验和可辨识框架使后验推断 p(z|x,u) 可解释，直到可辨识变换为止，支持对神经码的科学洞见。