Associative Transformer¶

会议: CVPR 2025
arXiv: 2309.12862
代码: 暂无公开
领域: LLM效率
关键词: Transformer, explicit memory, Hopfield network, sparse representation, bottleneck attention

一句话总结¶

提出 Associative Transformer (AiT)，通过在 Transformer 中引入可学习的显式记忆模块和 Hopfield 网络进行 token 重建，以更少的参数实现优于 ViT 的分类和关系推理性能。

研究背景与动机¶

领域现状：Vision Transformer (ViT) 通过 self-attention 机制在视觉任务中取得了显著进展，但其 token 表示缺乏显式结构化记忆的支持，所有信息都隐式编码在注意力权重中。

现有痛点：标准 Transformer 的 attention 机制对所有 token 进行全局交互，计算复杂度为 \(O(N^2)\)，且没有机制来维护跨样本的持久化信息表示。模型容易在小数据集上过拟合，且在需要关系推理的任务中表现受限。

核心矛盾：Transformer 虽然有强大的表示能力，但缺乏类似人脑"全局工作空间"（Global Workspace Theory）的机制。现有方法要么完全依赖隐式表示，要么引入外部记忆但缺乏有效的检索机制。

本文目标 如何在 Transformer 中引入持久化显式记忆，使 token 能够通过竞争访问共享的记忆池，同时保持计算效率？

切入角度：借鉴认知科学中的全局工作空间理论和联想记忆（Hopfield Network），设计 bottleneck 机制让 token 竞争进入共享记忆空间。

核心 idea：引入 Global Workspace Layer，结合 low-rank 显式记忆、bottleneck attention 和 Hopfield 网络，让 Transformer 具备持久化、竞争性的联想记忆能力。

方法详解¶

整体框架¶

AiT 在标准 ViT 的基础上在每个 Transformer block 中新增一个 Global Workspace Layer (GWL)。输入为图像 patch token 序列，经过 self-attention 后进入 GWL 进行记忆交互和 token 重建，最终输出增强的 token 表示。

关键设计¶

Low-rank Explicit Memory
- 功能：维护一个可学习的记忆池 \(\gamma \in \mathbb{R}^{M \times D}\)，\(M\) 为记忆槽位数（32-128），\(D\) 为低维嵌入维度（32）
- 核心思路：记忆通过 EWMA 持续更新 \(\gamma^{t+1} = (1-\alpha)\gamma^t + \alpha \cdot \text{LN}(\text{Concat}(h_1,...,h_S)W^O)\)，\(\alpha=0.1\)
- 设计动机：低维设计使记忆池可扩展到 32.8K 个 token 而不增加过多计算开销，跨 batch 更新使其积累全局统计信息
Bottleneck Attention
- 功能：通过 top-k 选择机制强制 token 竞争进入记忆空间
- 核心思路：计算每个 token 对各记忆槽的注意力分数，仅保留 top-k 个分数最高的 token 与记忆交互
- 设计动机：竞争机制模拟了全局工作空间的"广播"过程，确保只有最相关的信息被写入共享记忆
- Balance Loss 包含两部分：累积注意力均衡和选择频率均衡
Hopfield Network Token Reconstruction
- 功能：使用连续 Hopfield 网络从记忆中检索和重建 token 表示
- 核心思路：Hopfield 能量函数 \(E(\Xi^t) = -\text{lse}(\beta, f_{LT}(\gamma^{t+1})\Xi^t) + \frac{1}{2}\Xi^t(\Xi^t)^T\)
- 设计动机：Hopfield 网络天然适合从记忆池中提取匹配模式，且 FLOPs 仅占总计算量的 0.84%

损失函数 / 训练策略¶

总损失：\(\ell = \ell_{\text{class}} + \sigma \cdot \sum \ell_{\text{bottleneck}_i}\)，\(\sigma = 10^{-2}\)
批量大小：512（CIFAR），128（Pet），64（relational reasoning）
记忆槽数 M：32（CIFAR/Triangle），128（Pet）
Bottleneck 容量：512（CIFAR/Pet），64（Triangle）

实验关键数据¶

主实验¶

数据集	AiT-Base (91M)	AiT-Medium (45.9M)	ViT-Base (85.7M)	ViT-Medium
CIFAR10	85.44%	84.59%	83.82%	82.41%
CIFAR100	60.78%	60.58%	57.92%	55.78%
Triangle	99.64%	99.57%	99.63%	99.62%
平均	81.95%	81.58%	80.46%	79.27%

AiT-Medium（45.9M 参数）超过 ViT-Base（85.7M），参数量仅一半。ImageNet100: AiT-Medium 36.72% vs ViT-Base 34.62%。

消融实验¶

配置	平均准确率	变化
Full AiT-Small	79.70%	—
w/o Bottleneck	72.75%	-6.95%
w/o Self-Attention	73.31%	-6.39%
w/o Memory (=ViT)	77.40%	-2.30%
w/o Hopfield	78.48%	-1.22%
w/o Balance Loss	78.68%	-1.02%
Reset Memory	79.12%	-0.58%

关键发现¶

Bottleneck attention 贡献最大（-6.95%），竞争性访问机制是核心设计
去掉记忆后退化为 ViT（-2.30%），记忆模块提供了额外容量
Hopfield 计算开销极低（\(8.02 \times 10^6\) FLOPs，<0.84%），但贡献 -1.22%
Oxford Pet 实验中 ViT-Base 在 50 epoch 后过拟合，AiT-Small 持续上升
Sort-of-CLEVR 关系推理：AiT-Base 80.03%（关系任务），优于标准 Transformer

亮点与洞察¶

认知科学启发的架构设计：将全局工作空间理论引入 Transformer 是巧妙的跨学科迁移
参数效率反直觉：更小的 AiT-Medium 超过更大的 ViT-Base，结构化记忆比单纯增加参数更有效
Hopfield 网络的轻量应用：仅 0.84% 计算开销就带来稳定增益
记忆的 EWMA 更新可迁移到在线学习、持续学习等场景

局限与展望¶

实验仅在小规模数据集上验证，缺乏 ImageNet-1K 完整评估和下游密集预测任务
记忆槽数 M 和 bottleneck 容量 k 需手动调整
未探索与 LoRA 等高效微调方法的结合

评分¶

新颖性: ⭐⭐⭐⭐ 认知科学启发设计有创意，但 external memory in Transformer 非全新
实验充分度: ⭐⭐⭐ 数据集规模偏小，缺乏大规模评估
写作质量: ⭐⭐⭐⭐ 动机清晰，方法描述详细
价值: ⭐⭐⭐⭐ 为 Transformer 中结构化记忆探索了有前景的方向