最近在看 MIT 6.S191 Lecture 9 讲 Context Parallelism 的时候,碰到了一个词:Causal Attention。
我对 Self-Attention 的公式并不陌生——之前那篇已经把 QKV 和 routing score table 拆得比较清楚了。但 Causal Attention 这个概念我一直没有认真想过。回头看,我写 Attention 公式、写 Multi-Head、写 FlashAttention 的时候,脑子里跑的都是"所有 token 都能看到所有 token"的 Full Self-Attention。
起因 # 最近在过 MIT 6.S191 和分布式训练相关的材料,各种 Norm 反复出现。LayerNorm 在写 Transformer Block 那篇时已经接触过,BatchNorm 在做 CV 的那几年天天调,RMSNorm 在看 LLM 代码时到处都是。单独看每个都不算陌生。
上一篇结尾我写:Multi-Head Attention 仍然只是一个 information routing / read mechanism,一个完整的 Transformer layer 还包含 FFN、residual connection、layer normalization 等组件。
上一篇走到了一个结论:去掉 recurrence 之后,Self-Attention 必须重新获得 positional information。position 可以加进 token 的 representation,两者相加不要求后续计算无损还原,只要下游能利用其中的信息就行。
上一篇结尾我留了一个问题:一组 \(W_Q, W_K, W_V\) 只定义了一个 matching / routing space,如果不同类型的信息需求需要不同的 routing 策略,一个 space 够吗?
这个问题让我停了下来。因为到 single-head self-attention 为止,每个 token 已经可以访问整个 sequence 了——任意两个位置之间的通路都是通的。如果表达能力不够,加大 \(d_k\) 不就行了?为什么还要搞出 Multi-Head?
上一篇结尾留了三个问题:谁和我相关?相关程度是多少?我该从对方那里读什么?
Self-attention 允许 sequence 内的位置之间直接建立通路,不再需要沿着 recurrent state 一步步搬运。但"允许直接交互"和"知道怎么交互"是两件事——architecture 把路打通了,到底是谁在决定走哪条路、读多少、读什么?
上一篇《从 RNN / LSTM / GRU 到早期 Attention》结尾我留了一句话:早期的 attention 并不是为了取代 RNN 而出现的,它只是夹在 encoder–decoder 中间解决 fixed-context bottleneck;但“根据当前需求去读取一组 representation”这种机制,并不一定要依附在 RNN 上。这篇接着讲那个“另一个故事”。
起因 # 最近在做 CS336(Language Modeling from Scratch)的 Transformer 实现,代码里反复出现这种写法: