论文阅读笔记——Multi-Token Attention

MTA 论文

在 Transformer 中计算注意力权重时,仅依赖单个 Q 和 K 的相似度,无法有效捕捉多标记组合信息。(对于 A、B 两个词,单标记注意力需要分别计算两个词的注意力分数,再通过后处理定位共同出现的位置或通过多层隐式堆叠,增加模型深度和容量)。MTA 显示建模多标记依赖,同时不牺牲全局交互和额外参数。(通过卷积运算让他能够看到邻近的Q、K 以及其他注意力头的信息)

在 Transformer 其他部分,如 FFN 的输入/输出加卷积,主要是为了捕捉词元表示 之间的局部依赖关系,不直接改变注意力机制本身如何计算相关性
MTA 的卷积直接作用在 Q K T / A QK^T/A QKT/A ,意味着卷积直接参与了决定哪些上下文位置应该被关注的过程,在处理词元间的关系强度

提出两种方式:pre-softmax convolution 和 post-softmax convolution,MTA 默认采用 Pre-softmax Q-K Convolution 和 Post-softmax Head Mixing Convolution。二者区别在于是在 softmax 之前还是之后进行。

Q-K convolution

a i j = S o f t m a x ( ∑ i ′ = 0 c q − 1 ∑ j ′ = − ⌊ c k / 2 ⌋ ⌈ c k / 2 ⌉ − 1 1 i ≥ j − j ′ θ i ′ , j ′ q i − i ′ k j − j ′ ⊤ / d ) ( 1 ) a_{ij}=\mathrm{Softmax}\left(\sum_{i^{\prime}=0}^{c_{q}-1}\sum_{j^{\prime}=-\lfloor c_{k}/2\rfloor}^{\lceil c_{k}/2\rceil-1}\mathbf{1}{i\geq j- j^{\prime}}\theta{i^{\prime},j^{\prime}}q_{i-i^{\prime}}k_{j-j^{\prime}}^{\top}/\sqrt{d}\right) \qquad \qquad(1) aij=Softmax i′=0∑cq−1j′=−⌊ck/2⌋∑⌈ck/2⌉−11i≥j−j′θi′,j′qi−i′kj−j′⊤/d (1)

在卷积中,为防止未来信息泄露,需要做 Masking。理想的 Masking 比较复杂(见式(1)),采用一种简化形式:用 0 Mask 掉未来的 Q K T QK^T QKT 值,做卷积,再用 − ∞ -\infty −∞ Mask 掉结果中非法位置,再做 Softmax。
A = S o f t m a x ( M a s k − ∞ ( C o n v 2 d θ ( M a s k 0 ( A ^ ) ) ) ) . A=\mathrm{Softmax}\left(\mathrm{Mask}{-\infty}\left(\mathrm{Conv}2\mathrm{d}\theta\left(\mathrm{Mask}_0(\hat{A})\right)\right)\right). A=Softmax(Mask−∞(Conv2dθ(Mask0(A^)))).

Head Mixing Convolution

允许不同注意力头之间共享信息,放大重要信号。将 M 个头分成 M / c h M/c_h M/ch 个组,每组 c h c_h ch 个头。在每组的头内左 1D 卷积。同样可以在 softmax 之前或之后进行。

Group Normalization with depth scaling

改善梯度流,对抗深层网络中残差连接可能带来的主导效应(让模型更关注注意力部分输出,而不是仅仅传递上一层信息)。

在每个头的输出上独立应用组归一化,并结合一个随层数变化的缩放因子。

核心矛盾:在「增强注意力精度」和「保持计算效率」之间尚未找到完美平衡,当前更适合对计算资源不敏感的高精度场景。

实验结果

1.找字母块任务 ,验证 MTA 能够解决 [多条件匹配] 问题。

MTA 错误率接近 0% ,而 Transformer 失败率超 50%

2.LLM,在 105B 词元数据上训练 880M 参数模型

  • MTA 仅在 1/4 的层 使用 Key-Query 卷积(核大小: c q = 6 , c k = 11 c_q=6,c_k=11 cq=6,ck=11)。
  • 所有层使用 Head 卷积(核大小 c h = 2 c_h=2 ch=2)。
相关推荐
数智工坊2 小时前
【Transfer CLIP论文阅读】跨模态大模型赋能!CLIP迁移学习实现超强泛化图像去噪
论文阅读·人工智能·迁移学习
AI周红伟2 小时前
三年狂赚1.75亿!卖课,才是中国AI最容易赚钱的生意
人工智能·深度学习·学习·机器学习·copilot·openclaw
东京老树根2 小时前
SAP学习笔记 - BTP CAP开发06 - Customize Criticality,Map Page,Value help
笔记·学习
AM越.3 小时前
助睿:!!零代码解决!!订单利润分流数据加工o(* ̄▽ ̄*)ブ
数据仓库·笔记·etl·助睿
沪漂阿龙4 小时前
从买菜做饭到大模型:一份真正看懂深度学习的硬核指南
人工智能·深度学习
魔都大虾4 小时前
旧时光里面有那些情话句子 什么比较热门
笔记
Java成神之路-5 小时前
【算法刷题笔记】全题型导航目录
笔记·算法
zhangrelay5 小时前
云课实践速通系列-基础篇汇总-必修-通识基础和专业基础-2026--工科--自动化、电气、机器人、测控等
linux·笔记·单片机·学习·ubuntu·机器人·自动化
05候补工程师5 小时前
【编译原理】自顶向下语法分析深度解析:从 LL(1) 文法判定、改写到预测分析表
经验分享·笔记·考研·自然语言处理
2zcode5 小时前
基于深度学习的肺部听诊音疾病智能诊断方法研究
人工智能·深度学习