#!/usr/bin/env python3 # -*- coding: utf-8 -*- # Copyright 2019 Shigeki Karita # Apache 2.0 (http://www.apache.org/licenses/LICENSE-2.0) """Multi-Head Attention layer definition.""" import math import numpy import torch from torch import nn from espnet2.asr.encoder.Spike_driven.Q_trick import MultiSpike class Q_MultiHeadedAttention_HierDecay(nn.Module): """Implementation of HD-RepSSA_S Args: n_head (int): The number of heads. n_feat (int): The number of features. dropout_rate (float): Dropout rate. layer_id (int): Layer ID for decay calculation. """ def __init__(self, n_head, n_feat, dropout_rate, layer_id): """Construct an MultiHeadedAttention object.""" super(Q_MultiHeadedAttention_HierDecay_v2, self).__init__() assert n_feat % n_head == 0 # We assume d_v always equals d_k self.d_k = n_feat // n_head self.h = n_head self.linear_q = nn.Linear(n_feat, n_feat, bias=False) self.linear_k = nn.Linear(n_feat, n_feat, bias=False) self.linear_v = nn.Linear(n_feat, n_feat, bias=False) self.v_sn = MultiSpike(n_feat) self.output_sn = MultiSpike(n_feat) self.linear_out = nn.Linear(n_feat, n_feat) self.attn = None self.dropout = nn.Dropout(p=dropout_rate) # Hierarchical decay calculation layer_decay = 1 - 2 ** (-5 - layer_id) decay = torch.log(torch.tensor(layer_decay).repeat(n_head)) self.register_buffer("decay", decay) def forward_qkv(self, query, key, value, iiter): """Transform query, key and value. Args: query (torch.Tensor): Query tensor (#batch, time1, size). key (torch.Tensor): Key tensor (#batch, time2, size). value (torch.Tensor): Value tensor (#batch, time2, size). Returns: torch.Tensor: Transformed query tensor (#batch, n_head, time1, d_k). torch.Tensor: Transformed key tensor (#batch, n_head, time2, d_k). torch.Tensor: Transformed value tensor (#batch, n_head, time2, d_k). """ n_batch = query.size(0) q = self.linear_q(query).view(n_batch, -1, self.h, self.d_k) k = self.linear_k(key).view(n_batch, -1, self.h, self.d_k) v = self.v_sn(self.linear_v(value)).view(n_batch, -1, self.h, self.d_k) q = q.transpose(1, 2) # (batch, head, time1, d_k) k = k.transpose(1, 2) # (batch, head, time2, d_k) v = v.transpose(1, 2) # (batch, head, time2, d_k) return q, k, v def forward_attention(self, value, scores, mask, inner_mask, iiter): """Compute attention context vector. Args: value (torch.Tensor): Transformed value (#batch, n_head, time2, d_k). scores (torch.Tensor): Attention score (#batch, n_head, time1, time2). mask (torch.Tensor): Mask (#batch, 1, time2) or (#batch, time1, time2). inner_mask (torch.Tensor): Inner mask for hierarchical decay. Returns: torch.Tensor: Transformed value (#batch, time1, d_model) weighted by the attention score (#batch, time1, time2). """ n_batch = value.size(0) scores += inner_mask if mask is not None: mask = mask.unsqueeze(1).eq(0) # (batch, 1, *, time2) min_value = torch.finfo(scores.dtype).min scores = scores.masked_fill(mask, min_value) self.attn = torch.softmax(scores, dim=-1).masked_fill( mask, 0.0 ) # (batch, head, time1, time2) else: self.attn = torch.softmax(scores, dim=-1) # (batch, head, time1, time2) p_attn = self.dropout(self.attn) x = torch.matmul(p_attn, value) # (batch, head, time1, d_k) x = ( x.transpose(1, 2).contiguous().view(n_batch, -1, self.h * self.d_k) ) # (batch, time1, d_model) x = self.output_sn(x) return self.linear_out(x) # (batch, time1, d_model) def forward(self, query, key, value, mask, iiter): """Compute scaled dot product attention. Args: query (torch.Tensor): Query tensor (#batch, time1, size). key (torch.Tensor): Key tensor (#batch, time2, size). value (torch.Tensor): Value tensor (#batch, time2, size). mask (torch.Tensor): Mask tensor (#batch, 1, time2) or (#batch, time1, time2). Returns: torch.Tensor: Output tensor (#batch, time1, d_model). """ slen = query.shape[1] if query.shape[1]==key.shape[1] else mask.shape[1] index = torch.arange(slen).to(self.decay) inner_mask = torch.abs(index.view(slen,1) - index.view(1, slen)) inner_mask = inner_mask * self.decay[:, None, None] q, k, v = self.forward_qkv(query, key, value, iiter) scores = torch.matmul(q, k.transpose(-2, -1))/ math.sqrt(self.d_k) return self.forward_attention(v, scores, mask, inner_mask, iiter) class Q_MultiHeadedAttention_HierDecay_woSoftMax(Q_MultiHeadedAttention): """Implementation of HD-RepSSA_S Args: n_head (int): The number of heads. n_feat (int): The number of features. dropout_rate (float): Dropout rate. layer_id (int): Layer ID for decay calculation. """ def __init__(self, n_head, n_feat, dropout_rate, layer_id): """Construct an MultiHeadedAttention object.""" super().__init__(n_head, n_feat, dropout_rate) assert n_feat % n_head == 0 # We assume d_v always equals d_k self.d_k = n_feat // n_head self.h = n_head self.linear_q = nn.Linear(n_feat, n_feat, bias=False) self.linear_k = nn.Linear(n_feat, n_feat, bias=False) self.linear_v = nn.Linear(n_feat, n_feat, bias=False) self.q_sn = MultiSpike(n_feat) self.k_sn = MultiSpike(n_feat) self.v_sn = MultiSpike(n_feat) self.output_sn = MultiSpike(n_feat) self.linear_out = nn.Linear(n_feat, n_feat) self.attn = None self.dropout = nn.Dropout(p=dropout_rate) layer_decay = 1 - 2 ** (-5 - layer_id) decay = torch.log(torch.tensor(layer_decay).repeat(n_head)) self.register_buffer("decay", decay) self.ln = torch.nn.LayerNorm(self.d_k) def forward_qkv(self, query, key, value, iiter): """Transform query, key and value. Args: query (torch.Tensor): Query tensor (#batch, time1, size). key (torch.Tensor): Key tensor (#batch, time2, size). value (torch.Tensor): Value tensor (#batch, time2, size). Returns: torch.Tensor: Transformed query tensor (#batch, n_head, time1, d_k). torch.Tensor: Transformed key tensor (#batch, n_head, time2, d_k). torch.Tensor: Transformed value tensor (#batch, n_head, time2, d_k). """ n_batch = query.size(0) q = self.q_sn(self.linear_q(query)).view(n_batch, -1, self.h, self.d_k) k = self.k_sn(self.linear_k(key)).view(n_batch, -1, self.h, self.d_k) v = self.v_sn(self.linear_v(value)).view(n_batch, -1, self.h, self.d_k) q = q.transpose(1, 2) # (batch, head, time1, d_k) k = k.transpose(1, 2) # (batch, head, time2, d_k) v = v.transpose(1, 2) # (batch, head, time2, d_k) return q, k, v def forward_attention(self, value, scores, mask, inner_mask, iiter): """Compute attention context vector. Args: value (torch.Tensor): Transformed value (#batch, n_head, time2, d_k). scores (torch.Tensor): Attention score (#batch, n_head, time1, time2). mask (torch.Tensor): Mask (#batch, 1, time2) or (#batch, time1, time2). Returns: torch.Tensor: Transformed value (#batch, time1, d_model) weighted by the attention score (#batch, time1, time2). """ n_batch = value.size(0) scores = scores / scores.detach().abs().sum(dim=-1, keepdim=True).clamp(min=1, max=5e4) if mask is not None: mask = mask.unsqueeze(1).eq(0) # (batch, 1, *, time2) min_value = torch.finfo(scores.dtype).min scores = scores.masked_fill(mask, min_value) self.attn = scores.masked_fill( mask, 0.0 ) # (batch, head, time1, time2) else: self.attn = scores # (batch, head, time1, time2) self.attn = inner_mask * self.attn p_attn = self.dropout(self.attn) x = self.ln(torch.matmul(p_attn, value)) # (batch, head, time1, d_k) x = ( x.transpose(1, 2).contiguous().view(n_batch, -1, self.h * self.d_k) ) # (batch, time1, d_model) x = self.output_sn(x) return self.linear_out(x) # (batch, time1, d_model) def forward(self, query, key, value, mask, iiter): """Compute scaled dot product attention. Args: query (torch.Tensor): Query tensor (#batch, time1, size). key (torch.Tensor): Key tensor (#batch, time2, size). value (torch.Tensor): Value tensor (#batch, time2, size). mask (torch.Tensor): Mask tensor (#batch, 1, time2) or (#batch, time1, time2). Returns: torch.Tensor: Output tensor (#batch, time1, d_model). """ slen = query.shape[1] if query.shape[1]==key.shape[1] else mask.shape[1] index = torch.arange(slen).to(self.decay) inner_mask = torch.abs(index.view(slen,1) - index.view(1, slen)) inner_mask = torch.exp(inner_mask * self.decay[:, None, None]) q, k, v = self.forward_qkv(query, key, value, iiter) scores = torch.matmul(q, k.transpose(-2, -1)) return self.forward_attention(v, scores, mask, inner_mask, iiter) class Q_MultiHeadedAttention(nn.Module): """Multi-Head Attention layer. Args: n_head (int): The number of heads. n_feat (int): The number of features. dropout_rate (float): Dropout rate. """ def __init__(self, n_head, n_feat, dropout_rate): """Construct an MultiHeadedAttention object.""" super(Q_MultiHeadedAttention, self).__init__() assert n_feat % n_head == 0 # We assume d_v always equals d_k self.d_k = n_feat // n_head self.h = n_head self.linear_q = nn.Linear(n_feat, n_feat, bias=False) self.linear_k = nn.Linear(n_feat, n_feat, bias=False) self.linear_v = nn.Linear(n_feat, n_feat, bias=False) self.v_sn = MultiSpike(n_feat) self.output_sn = MultiSpike(n_feat) self.linear_out = nn.Linear(n_feat, n_feat) self.attn = None self.dropout = nn.Dropout(p=dropout_rate) def forward_qkv(self, query, key, value, iiter): """Transform query, key and value. Args: query (torch.Tensor): Query tensor (#batch, time1, size). key (torch.Tensor): Key tensor (#batch, time2, size). value (torch.Tensor): Value tensor (#batch, time2, size). Returns: torch.Tensor: Transformed query tensor (#batch, n_head, time1, d_k). torch.Tensor: Transformed key tensor (#batch, n_head, time2, d_k). torch.Tensor: Transformed value tensor (#batch, n_head, time2, d_k). """ n_batch = query.size(0) q = self.linear_q(query).view(n_batch, -1, self.h, self.d_k) k = self.linear_k(key).view(n_batch, -1, self.h, self.d_k) v = self.v_sn(self.linear_v(value)).view(n_batch, -1, self.h, self.d_k) q = q.transpose(1, 2) # (batch, head, time1, d_k) k = k.transpose(1, 2) # (batch, head, time2, d_k) v = v.transpose(1, 2) # (batch, head, time2, d_k) return q, k, v def forward_attention(self, value, scores, mask, iiter): """Compute attention context vector. Args: value (torch.Tensor): Transformed value (#batch, n_head, time2, d_k). scores (torch.Tensor): Attention score (#batch, n_head, time1, time2). mask (torch.Tensor): Mask (#batch, 1, time2) or (#batch, time1, time2). Returns: torch.Tensor: Transformed value (#batch, time1, d_model) weighted by the attention score (#batch, time1, time2). """ n_batch = value.size(0) if mask is not None: mask = mask.unsqueeze(1).eq(0) # (batch, 1, *, time2) min_value = torch.finfo(scores.dtype).min scores = scores.masked_fill(mask, min_value) self.attn = torch.softmax(scores, dim=-1).masked_fill( mask, 0.0 ) # (batch, head, time1, time2) else: self.attn = torch.softmax(scores, dim=-1) # (batch, head, time1, time2) p_attn = self.dropout(self.attn) x = torch.matmul(p_attn, value) # (batch, head, time1, d_k) x = ( x.transpose(1, 2).contiguous().view(n_batch, -1, self.h * self.d_k) ) # (batch, time1, d_model) x = self.output_sn(x) return self.linear_out(x) # (batch, time1, d_model) def forward(self, query, key, value, mask, iiter): """Compute scaled dot product attention. Args: query (torch.Tensor): Query tensor (#batch, time1, size). key (torch.Tensor): Key tensor (#batch, time2, size). value (torch.Tensor): Value tensor (#batch, time2, size). mask (torch.Tensor): Mask tensor (#batch, 1, time2) or (#batch, time1, time2). Returns: torch.Tensor: Output tensor (#batch, time1, d_model). """ q, k, v = self.forward_qkv(query, key, value, iiter) scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k) return self.forward_attention(v, scores, mask, iiter) class Q_MultiHeadedAttention_woSoftMax(Q_MultiHeadedAttention): """Multi-Head Attention layer without SoftMax. Args: n_head (int): The number of heads. n_feat (int): The number of features. dropout_rate (float): Dropout rate. """ def __init__(self, n_head, n_feat, dropout_rate): """Construct an MultiHeadedAttention object.""" super().__init__(n_head, n_feat, dropout_rate) assert n_feat % n_head == 0 # We assume d_v always equals d_k self.d_k = n_feat // n_head self.h = n_head self.linear_q = nn.Linear(n_feat, n_feat, bias=False) self.linear_k = nn.Linear(n_feat, n_feat, bias=False) self.linear_v = nn.Linear(n_feat, n_feat, bias=False) self.q_sn = MultiSpike(n_feat) self.k_sn = MultiSpike(n_feat) self.v_sn = MultiSpike(n_feat) self.output_sn = MultiSpike(n_feat) self.linear_out = nn.Linear(n_feat, n_feat) self.attn = None # self.dropout = nn.Dropout(p=dropout_rate) self.ln = torch.nn.LayerNorm(self.d_k) # self.scale = self.d_k ** -0.5 def forward_qkv(self, query, key, value, iiter): """Transform query, key and value. Args: query (torch.Tensor): Query tensor (#batch, time1, size). key (torch.Tensor): Key tensor (#batch, time2, size). value (torch.Tensor): Value tensor (#batch, time2, size). Returns: torch.Tensor: Transformed query tensor (#batch, n_head, time1, d_k). torch.Tensor: Transformed key tensor (#batch, n_head, time2, d_k). torch.Tensor: Transformed value tensor (#batch, n_head, time2, d_k). """ n_batch = query.size(0) k = self.linear_k(key).view(n_batch, -1, self.h, self.d_k) v = self.linear_v(value).view(n_batch, -1, self.h, self.d_k) q = self.q_sn(self.linear_q(query)).view(n_batch, -1, self.h, self.d_k) q = q.transpose(1, 2) # (batch, head, time1, d_k) k = k.transpose(1, 2) # (batch, head, time2, d_k) v = v.transpose(1, 2) # (batch, head, time2, d_k) return q, k, v def forward_attention(self, value, scores, mask, iiter): """Compute attention context vector. Args: value (torch.Tensor): Transformed value (#batch, n_head, time2, d_k). scores (torch.Tensor): Attention score (#batch, n_head, time1, time2). mask (torch.Tensor): Mask (#batch, 1, time2) or (#batch, time1, time2). Returns: torch.Tensor: Transformed value (#batch, time1, d_model) weighted by the attention score (#batch, time1, time2). """ n_batch = value.size(0) scores = scores / scores.detach().abs().sum(dim=-1, keepdim=True).clamp(min=1, max=5e4) if mask is not None: mask = mask.unsqueeze(1).eq(0) # (batch, 1, *, time2) min_value = torch.finfo(scores.dtype).min scores = scores.masked_fill(mask, min_value) self.attn = scores.masked_fill( mask, 0.0 ) # (batch, head, time1, time2) else: self.attn = scores # (batch, head, time1, time2) p_attn = self.dropout(self.attn) x = self.ln(torch.matmul(p_attn, value))# (batch, head, time1, d_k) # x = torch.matmul(p_attn, value) * self.scale# (batch, head, time1, d_k) x = ( x.transpose(1, 2).contiguous().view(n_batch, -1, self.h * self.d_k) ) # (batch, time1, d_model) x = self.output_sn(x) return self.linear_out(x) # (batch, time1, d_model) def forward(self, query, key, value, mask, iiter): """Compute scaled dot product attention. Args: query (torch.Tensor): Query tensor (#batch, time1, size). key (torch.Tensor): Key tensor (#batch, time2, size). value (torch.Tensor): Value tensor (#batch, time2, size). mask (torch.Tensor): Mask tensor (#batch, 1, time2) or (#batch, time1, time2). Returns: torch.Tensor: Output tensor (#batch, time1, d_model). """ q, k, v = self.forward_qkv(query, key, value, iiter) scores = torch.matmul(q, k.transpose(-2, -1)) return self.forward_attention(v, scores, mask, iiter)