Fast Weight Systems

#

9from typing import Optional
10
11import torch
12from torch import nn
13
14from labml_helpers.module import Module
15from labml_nn.transformers.fast_weights import DPFP
16from labml_nn.transformers.feed_forward import FeedForward
17from labml_nn.transformers.mha import PrepareForMultiHeadAttention
18from labml_nn.utils import clone_module_list

#

21class FastWeightsAttention(Module):

#

22    def __init__(self, heads: int, d_model: int, dropout_prob: float, phi: DPFP):
23        super().__init__()

#

Number of features per head

26        self.d_k = d_model // heads

#

28        self.heads = heads

#

These transform the query multi-headed attention.

31        self.query = PrepareForMultiHeadAttention(d_model, heads, self.d_k, bias=False)

#

These transform the key and value for multi-headed attention.

33        self.key = PrepareForMultiHeadAttention(d_model, heads, self.d_k, bias=False)
34        self.value = PrepareForMultiHeadAttention(d_model, heads, self.d_k, bias=False)
35
36        self.gate = nn.Sequential(PrepareForMultiHeadAttention(d_model, heads, 1, bias=False),
37                                  nn.Sigmoid())
38
39        self.phi = phi

#

Output layer

42        self.output = nn.Linear(d_model, d_model)

#

Dropout

44        self.dropout = nn.Dropout(dropout_prob)

#

46    def forward(self, x: torch.Tensor, weights: Optional[torch.Tensor]):
47        query = self.phi(self.query(x))
48        key = self.phi(self.key(x))
49        value = self.value(x)
50
51        if weights is None:
52            weights = key.new_zeros((key.shape[0], key.shape[1], value.shape[2], key.shape[2]))
53
54        value_existing = torch.einsum('bhvk,bhk->bhv', weights, key)
55
56        beta = self.gate(x)
57
58        weights = weights + torch.einsum('bhv,bhk->bhvk', beta * (value - value_existing), key)
59
60        x = torch.einsum('bhvk,bhk->bhv', weights, query)

#

Concatenate multiple heads

63        x = x.reshape(x.shape[0], -1)

#

Output layer

66        return self.output(x), weights

#

69class FastWeightsAttentionTransformerLayer(Module):

#

70    def __init__(self, *,
71                 d_model: int,
72                 attn: FastWeightsAttention,
73                 feed_forward: FeedForward,
74                 dropout_prob: float):
75        super().__init__()

#

Transformer size $d_{m o d e l}$

77        self.size = d_model

#

79        self.attn = attn
80        self.feed_forward = feed_forward
81        self.dropout = nn.Dropout(dropout_prob)

#

Normalization layers

84        self.norm_self_attn = nn.LayerNorm([d_model])
85        self.norm_ff = nn.LayerNorm([d_model])

#

87    def forward(self, x: torch.Tensor, weights: Optional[torch.Tensor]):
88        attn, weights = self.attn(x, weights)

#

Add the self attention results

90        x = x + self.dropout(attn)

#

Normalize for feed-forward

93        z = self.norm_ff(x)

#

Pass through the feed-forward network

95        ff = self.feed_forward(z)

#

Add the feed-forward results back

97        x = x + self.dropout(ff)

#

100        return x, weights

#

103class FastWeightsAttentionTransformer(Module):

#

104    def __init__(self, layer: FastWeightsAttentionTransformerLayer, n_layers: int):
105        super().__init__()

#

Make copies of the transformer layer

107        self.layers = clone_module_list(layer, n_layers)

#

Final normalization layer

109        self.norm = nn.LayerNorm([layer.size])

#

111    def forward(self, x_seq: torch.Tensor):

#

Split the input to a list along the sequence axis

113        x_seq = torch.unbind(x_seq, dim=0)

#

List to store the outputs

115        res = []

#

For each input step

117        weights = [None for _ in range(len(self.layers))]
118
119        for x in x_seq:

#

Run through each layer

121            for i, layer in enumerate(self.layers):

#

Get layer output

123                x, weights[i] = layer(x, weights[i])
124
125            res.append(x)

#

Stack the output tensors

128        res = torch.stack(res)

#

Normalize the output

130        return self.norm(res)