[论文解读] Learning Efficient Algorithms with Hierarchical Attentive Memory

Marcin Andrychowicz, Karol Kurach|arXiv (Cornell University)|Feb 9, 2016

Topic Modeling参考文献 22被引用 28

一句话总结

该论文提出了一种新型神经记忆架构——分层注意记忆（Hierarchical Attentive Memory, HAM），通过二叉树结构实现 $Θ(\log n)$ 的高效内存访问，使LSTM能够仅从纯输入输出对中学习复杂算法（如排序和搜索）。该模型在训练时未见过的更长序列上仍能泛化，并可模拟经典数据结构（如栈和队列）。

ABSTRACT

In this paper, we propose and investigate a novel memory architecture for neural networks called Hierarchical Attentive Memory (HAM). It is based on a binary tree with leaves corresponding to memory cells. This allows HAM to perform memory access in O(log n) complexity, which is a significant improvement over the standard attention mechanism that requires O(n) operations, where n is the size of the memory. We show that an LSTM network augmented with HAM can learn algorithms for problems like merging, sorting or binary searching from pure input-output examples. In particular, it learns to sort n numbers in time O(n log n) and generalizes well to input sequences much longer than the ones seen during the training. We also show that HAM can be trained to act like classic data structures: a stack, a FIFO queue and a priority queue.

研究动机与目标

• 为解决标准注意力机制在神经网络中因内存访问需 $Θ(n)$ 操作而导致的可扩展性问题，限制其在长序列任务中的应用。
• 设计一种随内存规模高效扩展的内存架构，同时支持从纯输入输出示例中学习算法行为。
• 使神经网络能够学习并泛化复杂算法（如排序和二分查找），其适用长度超过训练时所见序列长度。
• 证明该内存结构可模拟经典数据结构（如栈、FIFO队列和优先队列）。

提出的方法

• HAM 架构使用二叉树结构，其中叶节点表示内存单元，内部节点通过可微分的连接操作聚合信息。
• 内存访问通过从根节点到叶节点的树路径遍历完成，路径上的注意力概率通过可微分的路由机制计算。
• 模型采用一种软注意力变体（DHAM），通过路径概率对所有叶节点的期望激活进行计算，支持端到端反向传播。
• LSTM 控制器通过对树结构进行注意力来读取和写入内存单元，更新信息沿树向上传播以保持一致性。
• 该架构通过分层结构内置的区间操作偏好，支持算法任务的学习。
• 系统通过监督学习在输入输出对上进行训练，无需执行轨迹或奖励塑形。

实验结果

研究问题

• RQ1 神经网络能否仅从纯输入输出示例中学习对数字序列进行排序，并泛化到训练时未见的更长序列？
• RQ2 分层注意力机制是否相比标准注意力机制，能实现更快且更样本高效的算法任务学习？
• RQ3 HAM 内存结构能否模拟经典数据结构（如栈、队列和优先队列）？
• RQ4 HAM 的 $Θ(\log n)$ 内存访问复杂度是否足以支持 $O(n\log n)$ 时间复杂度算法（如排序）的学习？
• RQ5 与随机（硬）版本相比，可微分（软）版本的 HAM 在训练稳定性和泛化能力方面表现如何？

主要发现

• HAM 增强的 LSTM 能够以 $\Theta(n\log n)$ 时间复杂度成功学习对 $n$ 个数字进行排序，并泛化到远长于训练时所见长度的序列。
• 该模型在长度达到训练时最大长度 10 倍的输入序列上仍表现出良好泛化能力，显示出强大的归纳偏差。
• HAM 机制实现了 $\Theta\left(\log n\right)$ 的内存访问复杂度，显著优于标准注意力机制所需的 $\Theta(n)$ 操作。
• HAM 内存可模拟栈、FIFO 队列和优先队列，证明其作为通用内存抽象的多功能性。
• 可微分的软版本（DHAM）训练比随机版本更稳定，但在更大内存规模上的泛化能力较差。
• 分层结构提供了有利于学习基于区间的操作（如归并和二分查找）的归纳偏差，而这些操作是高效算法的核心。