AutoCorrelation Mechanism with the following two phases: (1) period-based dependencies discovery (2) time delay aggregation This block can replace the self-attention family mechanism seamlessly.
| 26 | |
| 27 | |
| 28 | class AutoCorrelation(nn.Module): |
| 29 | """ |
| 30 | AutoCorrelation Mechanism with the following two phases: |
| 31 | (1) period-based dependencies discovery |
| 32 | (2) time delay aggregation |
| 33 | This block can replace the self-attention family mechanism seamlessly. |
| 34 | """ |
| 35 | def __init__(self, mask_flag=True, factor=1, scale=None, attention_dropout=0.1, output_attention=False): |
| 36 | super(AutoCorrelation, self).__init__() |
| 37 | self.factor = factor |
| 38 | self.scale = scale |
| 39 | self.mask_flag = mask_flag |
| 40 | self.output_attention = output_attention |
| 41 | self.dropout = nn.Dropout(attention_dropout) |
| 42 | |
| 43 | def time_delay_agg_training(self, values, corr): |
| 44 | """ |
| 45 | SpeedUp version of Autocorrelation (a batch-normalization style design) |
| 46 | This is for the training phase. |
| 47 | """ |
| 48 | head = values.shape[1] |
| 49 | channel = values.shape[2] |
| 50 | length = values.shape[3] |
| 51 | # find top k |
| 52 | top_k = int(self.factor * math.log(length)) |
| 53 | mean_value = torch.mean(torch.mean(corr, dim=1), dim=1) |
| 54 | index = torch.topk(torch.mean(mean_value, dim=0), top_k, dim=-1)[1] |
| 55 | weights = torch.stack([mean_value[:, index[i]] for i in range(top_k)], dim=-1) |
| 56 | # update corr |
| 57 | tmp_corr = torch.softmax(weights, dim=-1) |
| 58 | # aggregation |
| 59 | tmp_values = values |
| 60 | delays_agg = torch.zeros_like(values).float() |
| 61 | for i in range(top_k): |
| 62 | pattern = torch.roll(tmp_values, -int(index[i]), -1) |
| 63 | delays_agg = delays_agg + pattern * \ |
| 64 | (tmp_corr[:, i].unsqueeze(1).unsqueeze(1).unsqueeze(1).repeat(1, head, channel, length)) |
| 65 | return delays_agg |
| 66 | |
| 67 | def time_delay_agg_inference(self, values, corr): |
| 68 | """ |
| 69 | SpeedUp version of Autocorrelation (a batch-normalization style design) |
| 70 | This is for the inference phase. |
| 71 | """ |
| 72 | batch = values.shape[0] |
| 73 | head = values.shape[1] |
| 74 | channel = values.shape[2] |
| 75 | length = values.shape[3] |
| 76 | # index init |
| 77 | init_index = torch.arange(length).unsqueeze(0).unsqueeze(0).unsqueeze(0)\ |
| 78 | .repeat(batch, head, channel, 1).to(values.device) |
| 79 | # find top k |
| 80 | top_k = int(self.factor * math.log(length)) |
| 81 | mean_value = torch.mean(torch.mean(corr, dim=1), dim=1) |
| 82 | weights, delay = torch.topk(mean_value, top_k, dim=-1) |
| 83 | # update corr |
| 84 | tmp_corr = torch.softmax(weights, dim=-1) |
| 85 | # aggregation |