src/Spike_driven/Spike_driven_modules/Q_attention.py

#!/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)