TikTok Linear Attention：结合律、前缀状态与信息取舍

题意与假设

这道练习题保留了面经中值得训练的核心约束，并将需要澄清的条件单独列出。现场拿到变体时，先和面试官确认这些条件。

• 训练核K(q,k)=φ(q)ᵀφ(k)，φ=ELU+1；归一化分母为φ(q)ᵀΣφ(k)，输出V加权均值。
• 本例非负特征且有限输入，causal时query/key等长；不声称此核等于exp(q·k/√d)。

解题路线

1. 普通显式attention构造n×m两两score；线性核可按结合律先算S=Σφ(k)vᵀ和z=Σφ(k)，每query输出φ(q)ᵀS/(φ(q)ᵀz)，不建完整矩阵。
2. causal时S/z只累计到当前i，更新一个key/value后求当前query；这给递归前缀状态，不需缓存全部历史K/V，状态尺寸由特征与value维决定。
3. 所谓linear是固定特征维下对序列长度线性，仍有d·dv矩阵乘和特征映射成本。若特征维随精度增长，不能忽略；多头状态也乘head数。
4. 核替换改变相似性与表达，固定状态把历史压缩，难保留所有任意精细token两两关系；不等于无条件保留softmax结果。某些随机特征近似有精度/方差取舍。
5. 归一化保证正权重下尺度合理，但所有feature极小可能下溢，分母零需错误或规定epsilon；eps会改变数值目标。相对位置、mask和局部约束加入后结合律是否还成立需重新推。
6. 测试用完整两两核矩阵算加权结果，与聚合状态对照；再测causal前缀，不用拿普通softmax当同一个oracle。真实任务效果与信息损失靠模型评估，而非只看速度。

手动走一遍

在causal位置i只能把0..i的key/value加入状态；若先累计全序列再逐query输出，会泄露未来。V相同常数时归一化输出仍该常数，是数值不变量。

Python 参考实现

下面的代码按上文假设实现，适合本地运行并追加自己的测试用例。

代码区可横向滚动；使用键盘时，先聚焦代码区再按方向键查看。

from math import isfinite


def linear_attention(q, k, v, causal=False):
    if not q or not k or not v or len(k) != len(v) or (causal and len(q) != len(k)):
        raise ValueError('invalid sequence shape')
    d, dv = len(q[0]), len(v[0])
    if not d or not dv or any(len(row) != d for row in q + k) or any(len(row) != dv for row in v):
        raise ValueError('invalid feature shape')
    if any(not isfinite(x) for row in q + k + v for x in row):
        raise ValueError('finite values required')
    # Positive feature map ELU(x)+1, not exact exp(q dot k).
    from math import exp
    def phi(row):
        return [x + 1 if x >= 0 else exp(x) for x in row]
    qs, ks = [phi(row) for row in q], [phi(row) for row in k]
    sums, products = [0.0] * d, [[0.0] * dv for _ in range(d)]
    def add(index):
        for a in range(d):
            sums[a] += ks[index][a]
            for b in range(dv):
                products[a][b] += ks[index][a] * v[index][b]
    if not causal:
        for j in range(len(k)):
            add(j)
    output = []
    for i, row in enumerate(qs):
        if causal:
            add(i)
        denominator = sum(a * b for a, b in zip(row, sums))
        if denominator <= 0 or not isfinite(denominator):
            raise ValueError('degenerate feature normalization')
        output.append([sum(row[a] * products[a][b] for a in range(d)) / denominator for b in range(dv)])
    return output

复杂度

非causal约O((n+m)d·dv)，causal等长O(nd·dv)，状态O(d·dv+d)加输出O(ndv)；固定维下线性于长度，但不是O(n)且与维数无关。

易错点与追问

容易失分的边界

• 结合律实现仍等于普通softmax。
• causal先累计未来。
• 忽略特征维成本。

继续深挖

• 位置偏置怎样保结合律？
• 近似核精度如何测？
• 固定状态对检索细节有什么限制？