未做分块时自定义Mistral模型滑动窗口注意力工作原理问询

用户代码片段 #

分词器实现代码 #

%%time
tokenizer = Tokenizer(models.BPE(byte_fallback=True))
trainer = trainers.BpeTrainer(vocab_size=vocab_size, special_tokens=["<pad>", "<unk>"], 
                              min_frequency=1500, show_progress=True)

tokenizer.train_from_iterator(seqs, trainer=trainer)

# Wrap in PreTrainedTokenizerFast
wrapped_tokenizer = PreTrainedTokenizerFast(
    tokenizer_object=tokenizer,
    unk_token="<unk>",
    pad_token="<pad>",
    max_length=max_length,
    truncation_side="left",
    padding_side="right",
)

自定义Mistral模型配置代码 #

custom_config = MistralConfig(
    vocab_size=160000,
    hidden_size=768,
    intermediate_size=2048,
    num_attention_heads=24,
    num_hidden_layers=24,
    max_position_embeddings=2048,
    sliding_window=1024,
    rms_norm_eps=1e-5,
    use_cache=False,
    pad_token_id=tokenizer.vocab['<pad>'],
    unk_token_id=tokenizer.vocab['<unk>'],
    bos_token_id=None,
    eos_token_id=None,
)
model = MistralForCausalLM(custom_config)
print(model)

核心疑问 #

已阅读滑动窗口注意力相关论文，但仍有两点困惑：

• 若未在分词器中设置分块参数（如return_overflowing_tokens=True），启用滑动窗口注意力后会有什么表现？
• 第一个Transformer层中，滑动注意力是否以步长1滑动，最终生成max_position_embeddings-1的特征图？因未找到步长参数，对此存疑。

解答 #

1. 未分词器分块时的滑动窗口行为 #

即使不在分词器层面做序列分块，滑动窗口注意力会在模型内部自动生效：

• 序列中第i个token仅能关注到[i - sliding_window + 1, i]范围内的token，也就是最多往前回溯1024个token。
• 比如输入序列长度为2048时，前1024个token里，第k个token（k≤1024）能看到自身及之前的所有token；从第1025个token开始，每个token只能访问最近的1024个上下文，无法触及更早的内容。
• 整个序列会被模型完整处理，滑动窗口的限制是在注意力计算阶段自动触发的，和分词器是否分块无关。

2. 滑动窗口的步长问题 #

滑动窗口注意力没有单独的步长参数，它的滑动逻辑是固定逐token进行的（步长为1）：

• 每个token的注意力窗口都以自身为终点，向前覆盖sliding_window长度的范围，相当于窗口以步长1逐个token向右移动。
• 最终生成的特征图长度和输入序列长度一致（即2048，并非max_position_embeddings-1），每个位置的特征都基于对应窗口内的上下文计算得到。
• 这种设计是为了保证每个token都能利用到最近的有效上下文，同时把注意力计算的时间/空间复杂度从O(n²)降到O(n*window_size)，大幅降低资源消耗。

内容的提问来源于stack exchange，提问作者Jonathan