【MHA】之 Attention Mask (with back & forward trace) / Causal Mask (with back trace)

最新推荐文章于 2025-06-03 14:27:28 发布
最新推荐文章于 2025-06-03 14:27:28 发布
阅读量4.8k 收藏 7
点赞数 3
CC 4.0 BY-SA版权
分类专栏： Python TensorFlow MHA 文章标签： ubuntu linux 运维
版权声明：本文为博主原创文章，遵循 CC 4.0 BY-SA 版权协议，转载请附上原文出处链接和本声明。
本文链接：https://blog.csdn.net/u010637291/article/details/121438116
本文介绍如何在TensorFlow中实现不同的Attention Mask，包括因果Mask、带N帧回溯的因果Mask及自定义Mask等，适用于多头注意力机制。

摘要生成于 C知道 ，由 DeepSeek-R1 满血版支持， 前往体验 >

文章目录

在multihead attention 中可添加attention mask，对输入进行范围限定，如

  • 因果mask (causal mask)：即可限定只看当前点前面的数据，不可看该点之后的数据。从矩阵上看，causal mask类似一个倒三角，下半部分为1，上半部分为0；
  • 因果mask带n_backtrace：即可限定每一点尽可最多向前看n_backtrace帧。从矩阵上看，即在上面的倒三角中，再在最左侧截去一部分，使得其为宽度为n_backtrace的斜带1；
  • 前后向N帧：即在上述带有n_backtrace的causal mask上，再以同样方式，向前即向右扩展一个宽度为n_backtrace的斜带1；
  • 类似的，可根据自定义需求，自行设定mask。

ref：
MHA TFA 的 实现: https://github.com/tensorflow/addons/blob/v0.15.0/tensorflow_addons/layers/multihead_attention.py#L23-L298

1. Attention Mask or Causal Mask

可指定causal参数,来生成普通的attention mask 还是causal mask:

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

from tensorflow.keras.layers import Layer, Masking
import tensorflow as tf

class AttentionMask(Layer):
    """
	Computes attention mask.
	"""

    def __init__(self, causal, mask_value=-1e9):
        """
        Argument/s:
			causal - causal attention mask flag.
			mask_value - value used to mask components that aren't to be attended
				to (typically -1e9).
        """
        super(AttentionMask, self).__init__()
        self.causal = causal
        self.mask_value = mask_value
        if not isinstance(mask_value, float): raise ValueError("Mask value must be a float.")

    def call(self, inp):
        """
		Compute attention mask.

		Argument/s:
			inp - used to compute sequence mask.

		Returns:
			Attention mask.
		"""
        batch_size = tf.shape(inp)[0]
        max_seq_len = tf.shape(inp)[1]
        flat_seq_mask = Masking(mask_value=0.0).compute_mask(inp)
        seq_mask = self.merge_masks(tf.expand_dims(flat_seq_mask, axis=1), tf.expand_dims(flat_seq_mask, axis=2))
        ### HERE !!! ###
        causal_mask = self.lower_triangular_mask([1, max_seq_len, max_seq_len]) if self.causal else None
        ################
        logical_mask = self.merge_masks(causal_mask, seq_mask)
        unmasked = tf.zeros([batch_size, max_seq_len, max_seq_len])
        masked = tf.fill([batch_size, max_seq_len, max_seq_len], self.mask_value)
        att_mask = tf.where(logical_mask, unmasked, masked)
        seq_mask = tf.cast(seq_mask, tf.float32)
        return att_mask, seq_mask

    def lower_triangular_mask(self, shape):
        """
		Creates a lower-triangular boolean mask over the last 2 dimensions.

		Argument/s:
			shape - shape of mask.

		Returns:
			causal mask.
		"""
        row_index = tf.math.cumsum(
            tf.ones(shape=shape, dtype=tf.int32, name="row"), axis=-2)
        col_index = tf.math.cumsum(
            tf.ones(shape=shape, dtype=tf.int32, name="col"), axis=-1)
        return tf.math.greater_equal(row_index, col_index)

    def merge_masks(self, x, y):
        """
		Merges a sequence mask and a causal mask to make an attantion mask.

		Argument/s:
			x - mask.
			y - mask.

		Returns:
			Attention mask.
		"""
        if x is None: return y
        if y is None: return x
        return tf.math.logical_and(x, y)

测试：

if __name__ == '__main__':
    input = tf.ones([64, 526, 40])

    attention_mask = AttentionMask(causal=0)(input)
    causal_mask = AttentionMask(causal=1)(input)
    print('done')

实验结果为：

在这里插入图片描述

其中attention mask为：
在这里插入图片描述

causal mask为：
在这里插入图片描述

2. Causal Mask (with n_backtrce)

即带有n_backtrce的因果mask,继承上面的AttentionMask:

from tensorflow.keras.layers import Masking
import tensorflow as tf

from AttentionMask import AttentionMask

class AttentionMask_Causal_Backtrace(AttentionMask):
    """
	Computes attention mask appropriate for tf.keras.layers.MultiHeadAttention.
	"""

    def __init__(self, causal, n_backtrace=None):
        """
		Argument/s:
			causal - causal attention mask flag.
			n_backtrace - (int) number of backtrace
		"""
        super().__init__(causal)
        self.causal = causal
        self.n_backtrace = n_backtrace

    def call(self, inp):
        """
		Compute attention mask.

		Argument/s:
			inp - used to compute sequence mask.

		Returns:
			Attention mask.
		"""
        batch_size = tf.shape(inp)[0]
        max_seq_len = tf.shape(inp)[1]
        flat_seq_mask = Masking(mask_value=0.0).compute_mask(inp)
        seq_mask = self.merge_masks(tf.expand_dims(flat_seq_mask, axis=1), tf.expand_dims(flat_seq_mask, axis=2))
        ### HERE !!! ###
        causal_mask = self.lower_triangular_mask([batch_size, max_seq_len, max_seq_len]) if self.causal else None
        bt_mask = self.backtrace_mask([1, max_seq_len, max_seq_len]) \
            if self.causal and self.n_backtrace else None
        ################
        logical_mask = self.merge_masks(causal_mask, seq_mask)
        logical_mask = self.merge_masks(logical_mask, bt_mask)
        att_mask = tf.cast(logical_mask, tf.float32)
        att_mask = tf.reshape(att_mask, [batch_size, 1, max_seq_len, max_seq_len])
        return att_mask

    def backtrace_mask(self, shape):
        """
		Creates a lower-triangular boolean mask over the last 2 dimensions.

		Argument/s:
			shape - shape of mask.

		Returns:
			causal mask.
		"""
        row_index = tf.math.cumsum(
            tf.ones(shape=shape, dtype=tf.int32, name="row"), axis=-2)
        col_index = tf.math.cumsum(
            tf.ones(shape=shape, dtype=tf.int32, name="col"), axis=-1)
        return tf.math.less_equal(row_index, col_index + self.n_backtrace)

测试：

if __name__ == '__main__':
    input = tf.ones([64, 526, 40])

    causal_mask = AttentionMask_Causal_Backtrace(causal=1, n_backtrace=None)(input)
    causal_mask_backtrace = AttentionMask_Causal_Backtrace(causal=1, n_backtrace=50)(input)
    print('done')

实验结果：
在这里插入图片描述

其中causal_mask为：
在这里插入图片描述

在这里插入图片描述

causal_mask_backtrace为:

在这里插入图片描述

在这里插入图片描述

在这里插入图片描述

测试样例2：

causal_mask_backtrace = AttentionMask_Causal_Backtrace(causal=1, n_backtrace=5)(input)

在这里插入图片描述

3. Attention Mask with backstrace and forwardtrace

from tensorflow.keras.layers import Masking
import tensorflow as tf

from AttentionMask import AttentionMask

class AttentionMask_Backtrace_Forwardtrace(AttentionMask):
    """
	Computes attention mask appropriate for tf.keras.layers.MultiHeadAttention.
	"""

    def __init__(self, causal, n_backtrace=None, n_forwardtrace=None):
        """
		Argument/s:
			causal - causal attention mask flag.
			n_backtrace - (int) number of backtrace
		"""
        super().__init__(causal)
        self.causal = causal
        self.n_backtrace = n_backtrace
        self.n_forwardtrace = n_forwardtrace

    def call(self, inp):
        """
		Compute attention mask.

		Argument/s:
			inp - used to compute sequence mask.

		Returns:
			Attention mask.
		"""
        batch_size = tf.shape(inp)[0]
        max_seq_len = tf.shape(inp)[1]
        flat_seq_mask = Masking(mask_value=0.0).compute_mask(inp)
        seq_mask = self.merge_masks(tf.expand_dims(flat_seq_mask, axis=1), tf.expand_dims(flat_seq_mask, axis=2))
        ### HERE !!! ###
        bt_ft_mask = self.backtrace_forwardtrace_mask([1, max_seq_len, max_seq_len]) \
            if self.n_backtrace and self.n_forwardtrace else None
        ################
        logical_mask = self.merge_masks(bt_ft_mask, seq_mask)
        att_mask = tf.cast(logical_mask, tf.float32)
        att_mask = tf.reshape(att_mask, [batch_size, 1, max_seq_len, max_seq_len])
        return att_mask

    def backtrace_forwardtrace_mask(self, shape):
        """
		Creates a lower-triangular boolean mask over the last 2 dimensions.

		Argument/s:
			shape - shape of mask.

		Returns:
			causal mask.
		"""
        row_index = tf.math.cumsum(
            tf.ones(shape=shape, dtype=tf.int32, name="row"), axis=-2)
        col_index = tf.math.cumsum(
            tf.ones(shape=shape, dtype=tf.int32, name="col"), axis=-1)
        bt_mask = tf.math.less_equal(row_index, col_index + self.n_backtrace)
        ft_mask = tf.math.greater_equal(row_index + self.n_forwardtrace, col_index)
        bt_ft_mask = self.merge_masks(bt_mask, ft_mask)
        return bt_ft_mask

测试：

if __name__ == '__main__':
    input = tf.ones([64, 526, 40])

    bt_ft_mask = AttentionMask_Backtrace_Forwardtrace(causal=0, n_backtrace=2, n_forwardtrace=5)(input)
    print('done')

实验结果：

在这里插入图片描述

bt_ft_mask为：

在这里插入图片描述

4. Customized Mask

class AttentionMask_Customization(AttentionMask):
    """
	Computes attention mask appropriate for tf.keras.layers.MultiHeadAttention.
	"""

    def __init__(self, causal, trace=None):
        """
		Argument/s:
			causal - causal attention mask flag.
			n_backtrace - (int) number of backtrace
		"""
        super().__init__(causal)
        self.causal = causal
        self.trace = trace

    def call(self, inp):
        """
		Compute attention mask.

		Argument/s:
			inp - used to compute sequence mask.

		Returns:
			Attention mask.
		"""
        batch_size = tf.shape(inp)[0]
        max_seq_len = tf.shape(inp)[1]
        flat_seq_mask = Masking(mask_value=0.0).compute_mask(inp)
        seq_mask = self.merge_masks(tf.expand_dims(flat_seq_mask, axis=1), tf.expand_dims(flat_seq_mask, axis=2))
        ### HERE !!! ###
        customized_mask = self.customized_mask(batch_size, max_seq_len, self.trace)
        ################
        logical_mask = self.merge_masks(customized_mask, seq_mask)
        att_mask = tf.cast(logical_mask, tf.float32)
        att_mask = tf.reshape(att_mask, [batch_size, 1, max_seq_len, max_seq_len])
        return att_mask

    @tf.function
    def customized_mask(self, batchsize, max_length, trace):
        mask = tf.ones(shape=[batchsize, trace, trace], dtype=tf.int32, name="row")
        shape_pad = int(max_length - trace)
        mask = tf.pad(mask, paddings=[[0, 0], [shape_pad, 0], [shape_pad, 0]])
        mask = tf.cast(mask, dtype=bool)
        return mask

测试：

if __name__ == '__main__':
    input = tf.ones([64, 526, 40])

    customized_mask = AttentionMask_Customization(causal=1, trace=5)(input)
    print('done')

实验结果：
在这里插入图片描述

TensorFlow 构造 attention mask 或 causal mask
Love-Coding
07-14 1550
def upper_triangle_bias(D, dtype=tf.float32): """Create a upper triangle matrix for decoding bias.""" upper_triangle_DxD = 1 - tf.matrix_band_part( tf.ones([D, D], dtype=dtype), -1, 0) tensor_1xDxD = tf.expand_dims(upper_triangle_DxD *
超细粒度分析 XLNet 中神奇的 Attention Mask
海晨威
11-29 1326
BERT后时代，XLNet 在20个任务上超越BERT，引起不小的关注。最重要的改进是 XLNet 实现了不在输入中加[Mask] 标志，同样可以利用上下文信息，解决了BERT预训练和微调不一致的问题。 1、XLNet如何实现在不加 [Mask] 情况下利用上下文信息呢？ XLNet 通过 Permutation Language Modeling (PLM) 重排输入文本，使得当前词的下文也有可能出现在当前词的「上文」中，而实际的操作不是通过改变文本输入顺序来实现，而是依靠关键的 Attention M
手写mask|代码详解，TriangularCausalMask/ProbMask/LocalMask
最新发布
哪惧明天，风高路斜
06-03 1046
多种掩码图解及代码详细注释
Attention Mask
weixin_43889416的博客
03-29 3391
Mask大致分为两种 Padding Mask：在NLP任务中，由于句子长度不一致，经常会进行padding操作，在sequence中加入零向量。这部分padding不应该起作用，但是在Attention的计算中用到了softmax等操作，即便0向量也会参与计算(e^0=1)，因此需要手动将这部分信息mask才行。padding mask主要包含两种： key mask：在计算score之后，且softmax之前进行，将值设为很小的数字(如-e^12)，这样经过的softmax之后值几乎为0 query m
attention-mask
weixin_43178406的博客
08-14 1105
参考链接：https://blog.csdn.net/weixin_40901056/article/details/97514718 def mask(inputs, queries=None, keys=None, type=None): """Masks paddings on keys or queries to inputs inputs: 3d tensor. (N, T_q, T_k) queries: 3d tensor. (N, T_q, d) keys:
LLM - Make Causal Mask 构造因果关系掩码
BITDDD小栈
09-25 5677
LLM MakeCausalMask 构造因果关系掩码函数解析。
GPT - 因果掩码（Causal Mask）
qq_60245590的博客
04-10 1155
本节代码定义了一个函数，用于生成因果掩码（Causal Mask）。因果掩码通常用于自注意力机制中，以确保模型在解码时只能看到当前及之前的位置，而不能看到未来的信息。这种掩码在自然语言处理任务（如语言生成）中非常重要，因为它模拟了人类阅读或写作时的顺序性。
Attention-mask 在transformer模型框架中的作用
li1784506的博客
11-19 7924
注意力机制的解释性博客比较多质量良莠不齐，推荐大家观看李宏毅老师关于注意力机制的讲解视频以及本人觉得对注意力机制讲解比较透彻的一篇博客[为更好解读注意力机制中attention-mask 的作用，现将注意力机制的原理进行总结。假设两个输入经过Wq、Wk、Wv矩阵（可训练）线性变换后获得q1=(1,2),q2=(0,1)，k1=(1,0),k2=(0,1)，v1=(1,0),v2=(0,1)向量。注意力机制核心就是向量q与向量k点乘后获得相似性分数（一个标量)。同理 q2也与所有的k向量点乘获得。
Attention mask理解
热门推荐
qq_39783265的博客
06-27 1万+
Transformer的矩阵维度分析和Mask详解 Transformer 中self-attention以及mask操作的原理以及代码解析
Attention中的Mask: query mask, key mask, future mask
chenxinvhai89的专栏
11-25 1449
版权声明：本文为博主原创文章，遵循 CC 4.0 BY-SA 版权协议，转载请附上原文出处链接和本声明。 本文链接：https://blog.csdn.net/weixin_40901056/article/details/97514718 Attention简介 Attention是2015年被提出来的，在NLP领域大放光彩。Attention具有在繁多信息中自动focus到重点的能力，而且A...
Transformer模型：intra-attention mask实现
weixin_62472350的博客
07-14 1256
前言这是对Transformer模型的Word Embedding、Postion Embedding、Encoder self-attention mask内容的续篇。
Mask-guided Contrastive Attention Model for Person Re-Identification学习笔记
与贰的博客
07-28 3658
刚刚去火车站买了下火车票，两个小时就为了刷个下铺。。。。。。突然看自己的笔记本，这周学的东西还有很多没有整理，上周的东西也还有很多没有清尾，最可耻的是上上周还有一篇计划了很久的东西，看了一半挂在那里没有看。郭师叔给的论文可能才看了1/8，最近又和我将强化学习是个很火的东西，花皮的deep learning买回来这么久也还没有正式看，还有一直心心念念的slam十四讲。我又不喜欢熬夜，每天还要定打不饶...
Causal mask代码阅读
weixin_45649595的博客
12-23 1397
transformer代码详解
Unilm生成式之Attention Mask解读
liliang199的专栏
03-13 2156
生成式对话说明
attention_mask的用处
qq_26751379的博客
11-06 4032
attention_mask：避免用注意力机制的时候关注到填充符（翻译可能不准，大致这个意思） 参考链接：https://huggingface.co/transformers/model_doc/bert.html?highlight=tfbertmodel%20from_pretrained#tfbertmodel
生成模型的中Attention Mask说明
weixin_42223207的博客
09-25 7830
生成模型中的attention mask 学习理解
一步一步理解大模型：因果掩码
chattyfish的博客
04-18 8721
为了防止这种情况，我们使用了一个注意力子序列屏蔽（attention subsequence mask），将所有未来的令牌设置为零，有效地从注意力机制中屏蔽了它们。具体实现中，这种输入不是串行输入的，而且提前准备的，蒙住从第二个到结尾的所有词，蒙住从第三个到结尾的所有词。在Transformer模型中，自注意力（self-attention）是在整个令牌（token）序列上计算的，包括当前令牌之后的令牌。这使得模型在进行预测时只能关注过去和当前的令牌，并确保它仅基于每个时间步骤可用的信息进行预测。
BERT模型中的input_ids和attention_mask参数
python_plus的博客
02-20 6645
介绍了BERT模型中的input_ids和attention_mask参数的含义和使用方法
MHA mask代码实现
02-01
### MHA Multi-Head Attention Mask Implementation Code Example In the context of implementing masking within a multi-head attention mechanism (MHA), it is essential to understand how masks are applied during the computation process. Masks play an important role by allowing certain positions not to contribute to the output, which can be particularly useful in tasks like language modeling where future tokens should not influence current predictions. Below is a Python code snippet demonstrating how one might implement masked multi-head self-attention using PyTorch: ```python import torch import torch.nn as nn import math class MultiHeadAttention(nn.Module): def __init__(self, d_model, num_heads): super(MultiHeadAttention, self).__init__() assert d_model % num_heads == 0, "d_model must be divisible by num_heads" self.d_k = d_model // num_heads self.num_heads = num_heads self.q_linear = nn.Linear(d_model, d_model) self.v_linear = nn.Linear(d_model, d_model) self.k_linear = nn.Linear(d_model, d_model) self.out = nn.Linear(d_model, d_model) def forward(self, q, k, v, mask=None): batch_size = q.size(0) # Linear transformations k = self.k_linear(k).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) q = self.q_linear(q).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) v = self.v_linear(v).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) # Calculate compatibility scores and apply scaling factor scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k) # Apply mask if provided if mask is not None: mask = mask.unsqueeze(1) # For broadcasting across heads dimension scores = scores.masked_fill(mask == 0, float('-inf')) # Softmax along last dimension before applying dropout attn_weights = torch.softmax(scores, dim=-1) output = torch.matmul(attn_weights, v) # Concatenate outputs from different heads together again concat_output = output.transpose(1, 2).contiguous().view(batch_size, -1, self.num_heads * self.d_k) # Final linear transformation after concatenation step final_output = self.out(concat_output) return final_output, attn_weights ``` This implementation includes support for optional masking through `mask` parameter passed into the `forward()` method[^1]. When present, this mask will prevent specific elements from contributing towards computing attention weights via setting their corresponding score values very low (`float('-inf')`) so they effectively get ignored when calculating softmax probabilities later on. --related questions-- 1. How does masking work inside transformer models? 2. What modifications would need to occur in order to adapt this code sample specifically for GQA or MQA implementations instead of standard MHA? 3. Can you explain why we use `-inf` value while applying mask in attention mechanisms? 4. Is there any difference between padding masks versus subsequent masks used within transformers architecture? If yes, what distinguishes them apart? 5. Could you provide more details about ziya_finetune project mentioned earlier regarding inference quantization examples?
评论
添加红包

请填写红包祝福语或标题

个

红包个数最小为10个

元

红包金额最低5元

当前余额3.43元 前往充值 >
需支付：10.00元
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付元
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明：

1.余额是钱包充值的虚拟货币，按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载，可以购买VIP、付费专栏及课程。

余额充值