[论文解读] Discrete-Valued Neural Communication

Dianbo Liu, Alex Lamb+5 more|arXiv (Cornell University)|Jul 6, 2021

Advanced Graph Neural Networks参考文献 49被引用 13

一句话总结

该论文提出离散值神经通信（DVNC），一种通过共享可学习码本强制结构化深度学习模型（如Transformer、图神经网络和模块化网络）中组件间通信的方法，以生成离散符号。通过将通信限制为离散值，DVNC在分布外（OOD）任务上提升了系统性泛化能力，增强了抗噪声鲁棒性，并降低了模型参数量，在无需架构修改的情况下，对多种架构均实现了稳定的性能提升。

ABSTRACT

Deep learning has advanced from fully connected architectures to structured models organized into components, e.g., the transformer composed of positional elements, modular architectures divided into slots, and graph neural nets made up of nodes. In structured models, an interesting question is how to conduct dynamic and possibly sparse communication among the separate components. Here, we explore the hypothesis that restricting the transmitted information among components to discrete representations is a beneficial bottleneck. The motivating intuition is human language in which communication occurs through discrete symbols. Even though individuals have different understandings of what a "cat" is based on their specific experiences, the shared discrete token makes it possible for communication among individuals to be unimpeded by individual differences in internal representation. To discretize the values of concepts dynamically communicated among specialist components, we extend the quantization mechanism from the Vector-Quantized Variational Autoencoder to multi-headed discretization with shared codebooks and use it for discrete-valued neural communication (DVNC). Our experiments show that DVNC substantially improves systematic generalization in a variety of architectures -- transformers, modular architectures, and graph neural networks. We also show that the DVNC is robust to the choice of hyperparameters, making the method very useful in practice. Moreover, we establish a theoretical justification of our discretization process, proving that it has the ability to increase noise robustness and reduce the underlying dimensionality of the model.

研究动机与目标

• 为解决结构化神经架构中组件间动态稀疏通信的挑战。
• 通过启用可重用的离散符号实现组件间通信，提升分布外（OOD）设置下的系统性泛化能力。
• 提供一种理论基础坚实、鲁棒性强的组件间通信方法，降低模型参数量并提高抗噪声能力。
• 探究离散共享表征是否可作为专家组件间的通用语言，受人类语言与大脑功能的启发。

提出的方法

• DVNC将向量量化变分自编码器（VQ-VAE）扩展至多头离散化，各组件共享码本。
• 通过K-means聚类在训练数据上学习大小为L × m的共享码本，将连续激活映射到最近的离散码本向量。
• 该方法使用直通梯度估计器，将任务损失通过离散码本反向传播，同时通过梯度下降更新码本，损失函数包括重构损失和任务承诺损失。
• 承诺损失确保激活值保持在其分配码本向量附近，而任务损失仅反向传播至发送组件。
• 组件间的通信被替换为离散标记索引，保留模型架构的同时强制实现类似符号的交换。
• 码本与模型其余部分端到端联合训练，无需指数移动平均更新，实现直接优化。

实验结果

研究问题

• RQ1 强制组件间离散通信是否能提升结构化神经网络的系统性泛化能力？
• RQ2 基于共享可学习码本的通信是否能增强对分布偏移和分布外泛化的能力？
• RQ3 离散通信如何影响模型敏感度与潜在维度，特别是覆盖数的对数？
• RQ4 DVNC的性能是否对超参数选择具有鲁棒性，尤其是码本大小和损失权重？
• RQ5 离散通信能否作为专用组件间的通用语言，提升模块化与可重用性？

主要发现

• DVNC在多种架构（包括GNN、Transformer和RIM）中显著提升了分布外泛化能力，在OOD基准测试中表现一致提升。
• 在Sort-of-Clevr任务中，DVNC在OOD设置下相较基线模型将平均倒数排名（MRR）提升了15%。
• 在2D Shapes和Three-body任务中，DVNC模型在OOD泛化上的准确率比非离散化基线高出20%-30%。
• 该方法对超参数选择具有鲁棒性，在不同码本大小和损失权重下均保持稳定性能。
• 理论分析证实，DVNC提升了模型敏感度并降低了覆盖数的对数，表明潜在维度更低。
• 训练时间可扩展：GNN在单张GPU上约3小时完成训练，RIM/Transformer约12小时每项任务，所有实验总计约800 GPU小时。