From 969df7190d5c584f04e3ae69aaf15cf9d2a29b8b Mon Sep 17 00:00:00 2001
From: Varuna Jayasiri Length of the memory (for masks) Total length of the memory and compressed memory (for masks) The default configs can and will be over-ridden when we start the experiment. The default configurations can and will be overridden when we start the experiment. If it’s configured not to use memory If the configurations specify not to use memory Return memories and the memories that were compressed.
-Memories that were compressed is needed for the reconstruction loss computation. The compression operation is defined as
$f_c: \mathbb{R}^{nc \times d} \rightarrow \mathbb{R}^{n \times d}$.
The paper introduces multiple choices for $f_c$ and we have only implemented
-1D convolution which seems to give best results.
+1D convolution which seems to give the best results.
Each layer has a separate compression operation $f_c^{(i)}$ where
$i$ is the layer number. Since training compression with BPTT requires maintaining
-a very large computational graph (many time steps), paper proposes
+a very large computational graph (many time steps), the paper proposes
an auto-encoding loss and an attention reconstruction loss.
-The auto-encoding loss, decodes the original memories from the compressed memories,
-and calculate the loss.
+The auto-encoding loss decodes the original memories from the compressed memories
+and calculates the loss.
Attention reconstruction loss computes the multi-headed attention results
-on the compressed memory and on uncompressed memory and get a mean squared error
+on the compressed memory and on uncompressed memory and gets a mean squared error
between them.
We have implemented the latter here since it gives better results. This implementation uses pre-layer norm while the paper uses post-layer norm.
+ This implementation uses pre-layer normalization
+while the paper uses post-layer normalization.
Pre-layer norm does the layer norm before FFN[../feedforward.html) and
-self attention, and the pass through in the residual connection is not normalized.
+self-attention, and the pass-through in the residual connection is not normalized.
This is supposed to be more stable in standard transformer setups. Here’s the training code and a notebook for training a compressive transformer
-model on Tiny Shakespeare dataset. Here are the training code and a notebook for training a compressive transformer
+model on the Tiny Shakespeare dataset. Run through each transformer layer Add to the list of feature vectors Memory Compressed Memory Run through the transformer XL layer Finally, normalize the vectors Attention reconstruction loss recreates the self attention output with
-uncompressed memory and with compressed memory and calculate mean squared error
+ Attention reconstruction loss recreates the self-attention output with
+uncompressed memory and with compressed memory and calculates the mean squared error
between the two. It does this without positional encoding. When calculating and training the compression function $f_c$ with attention
-reconstruction loss all parameters but $f_c$ are frozen.
-This includes key value projections and bias/scaling after normalization.58 m_len = len(mem[0]) if mem else 0
@@ -311,7 +311,7 @@
#
Configurations
-87class Configs(NLPAutoRegressionConfigs):
136 if self.mem_len == 0 and self.c_mem_len == 0:
@@ -726,7 +726,7 @@ and $N_m$ is the maximum number of memories we maintain (
mem_len
).<
#
195 return CompressedMemory(mem, c_mem), mem_to_compress
Training compression operation
53from typing import Optional, List
-54
-55import torch
-56import torch.nn.functional as F
-57from torch import nn
-58
-59from labml_helpers.module import Module, TypedModuleList
-60from labml_nn.transformers.feed_forward import FeedForward
-61from labml_nn.transformers.mha import PrepareForMultiHeadAttention
-62from labml_nn.transformers.xl.relative_mha import RelativeMultiHeadAttention
-63from labml_nn.utils import clone_module_list
54from typing import Optional, List
+55
+56import torch
+57import torch.nn.functional as F
+58from torch import nn
+59
+60from labml_helpers.module import Module, TypedModuleList
+61from labml_nn.transformers.feed_forward import FeedForward
+62from labml_nn.transformers.mha import PrepareForMultiHeadAttention
+63from labml_nn.transformers.xl.relative_mha import RelativeMultiHeadAttention
+64from labml_nn.utils import clone_module_list
66class Conv1dCompression(Module):
67class Conv1dCompression(Module):
74 def __init__(self, compression_rate: int, d_model: int):
75 def __init__(self, compression_rate: int, d_model: int):
79 super().__init__()
-80 self.conv = nn.Conv1d(d_model, d_model, kernel_size=compression_rate, stride=compression_rate)
80 super().__init__()
+81 self.conv = nn.Conv1d(d_model, d_model, kernel_size=compression_rate, stride=compression_rate)
mem
has shape [seq_len, batch, d_model]
82 def forward(self, mem: torch.Tensor):
83 def forward(self, mem: torch.Tensor):
[batch, features, sequence]
89 mem = mem.permute(1, 2, 0)
90 mem = mem.permute(1, 2, 0)
[batch, features, sequence]Get compressed memory by running it through the convolution layer
91 c_mem = self.conv(mem)
92 c_mem = self.conv(mem)
[batch, features, sequence]Permute back to form
[seq_len, batch, d_model]
93 return c_mem.permute(2, 0, 1)
94 return c_mem.permute(2, 0, 1)
[batch, features, sequence]This is the implementation of a single compressive transformer layer
96class CompressiveTransformerLayer(Module):
97class CompressiveTransformerLayer(Module):
[batch, features, sequence]
102 def __init__(self, *,
-103 d_model: int,
-104 self_attn: RelativeMultiHeadAttention,
-105 feed_forward: FeedForward,
-106 dropout_prob: float,
-107 compress: Conv1dCompression):
103 def __init__(self, *,
+104 d_model: int,
+105 self_attn: RelativeMultiHeadAttention,
+106 feed_forward: FeedForward,
+107 dropout_prob: float,
+108 compress: Conv1dCompression):
[batch, features, sequence]
115 super().__init__()
-116 self.compress = compress
-117 self.size = d_model
-118 self.self_attn = self_attn
-119 self.feed_forward = feed_forward
-120 self.dropout = nn.Dropout(dropout_prob)
-121 self.norm_self_attn = nn.LayerNorm([d_model])
-122 self.norm_ff = nn.LayerNorm([d_model])
116 super().__init__()
+117 self.compress = compress
+118 self.size = d_model
+119 self.self_attn = self_attn
+120 self.feed_forward = feed_forward
+121 self.dropout = nn.Dropout(dropout_prob)
+122 self.norm_self_attn = nn.LayerNorm([d_model])
+123 self.norm_ff = nn.LayerNorm([d_model])
[batch, features, sequence]
124 def concat_memory(self, z: torch.Tensor, mem: Optional[torch.Tensor], c_mem: Optional[torch.Tensor]):
125 def concat_memory(self, z: torch.Tensor, mem: Optional[torch.Tensor], c_mem: Optional[torch.Tensor]):
[batch, features, sequence]If there is no memory just return the token embeddings
133 if mem is None:
-134 return z
134 if mem is None:
+135 return z
[batch, features, sequence]If there are compressed memory concatenate that with memory
137 if c_mem is not None:
-138 mem = torch.cat((c_mem, mem), dim=0)
138 if c_mem is not None:
+139 mem = torch.cat((c_mem, mem), dim=0)
[batch, features, sequence]Run the memory through the normalization layer
141 mem = self.norm_self_attn(mem)
142 mem = self.norm_self_attn(mem)
[batch, features, sequence]Concatenate normalized memory and normalized token embeddings
143 return torch.cat((mem, z), dim=0)
144 return torch.cat((mem, z), dim=0)
[batch, features, sequence]
145 def forward(self, *,
-146 x: torch.Tensor,
-147 mem: Optional[torch.Tensor],
-148 c_mem: Optional[torch.Tensor],
-149 mask: torch.Tensor):
146 def forward(self, *,
+147 x: torch.Tensor,
+148 mem: Optional[torch.Tensor],
+149 c_mem: Optional[torch.Tensor],
+150 mask: torch.Tensor):
[batch, features, sequence]Normalize the vectors before doing self attention
159 z = self.norm_self_attn(x)
160 z = self.norm_self_attn(x)
[batch, features, sequence]Normalize and concatenate memory and compressed memory
161 m_z = self.concat_memory(z, mem, c_mem)
162 m_z = self.concat_memory(z, mem, c_mem)
[batch, features, sequence]Attention
163 self_attn = self.self_attn(query=z, key=m_z, value=m_z, mask=mask)
164 self_attn = self.self_attn(query=z, key=m_z, value=m_z, mask=mask)
[batch, features, sequence]Add the attention results
165 x = x + self.dropout(self_attn)
166 x = x + self.dropout(self_attn)
[batch, features, sequence]Normalize for feed-forward
168 z = self.norm_ff(x)
169 z = self.norm_ff(x)
[batch, features, sequence]Pass through the feed-forward network
170 ff = self.feed_forward(z)
171 ff = self.feed_forward(z)
[batch, features, sequence]Add the feed-forward results back
172 x = x + self.dropout(ff)
173 x = x + self.dropout(ff)
[batch, features, sequence]
175 return x
176 return x
[batch, features, sequence]This consists of multiple compressive transformer layers
178class CompressiveTransformer(Module):
179class CompressiveTransformer(Module):
[batch, features, sequence]
185 def __init__(self, layer: CompressiveTransformerLayer, n_layers: int):
-186 super().__init__()
186 def __init__(self, layer: CompressiveTransformerLayer, n_layers: int):
+187 super().__init__()
[batch, features, sequence]Make copies of the transformer layer
188 self.layers = clone_module_list(layer, n_layers)
189 self.layers = clone_module_list(layer, n_layers)
[batch, features, sequence]Final normalization layer
190 self.norm = nn.LayerNorm([layer.size])
191 self.norm = nn.LayerNorm([layer.size])
[batch, features, sequence]
192 def forward(self, x: torch.Tensor, mem: List[torch.Tensor], c_mem: List[torch.Tensor], mask: torch.Tensor):
193 def forward(self, x: torch.Tensor, mem: List[torch.Tensor], c_mem: List[torch.Tensor], mask: torch.Tensor):
[batch, features, sequence]
203 new_mem = []
204 new_mem = []
205 for i, layer in enumerate(self.layers):
206 for i, layer in enumerate(self.layers):
207 new_mem.append(x.detach())
208 new_mem.append(x.detach())
209 m = mem[i] if mem else None
210 m = mem[i] if mem else None
211 cm = c_mem[i] if c_mem else None
212 cm = c_mem[i] if c_mem else None
213 x = layer(x=x, mem=m, c_mem=cm, mask=mask)
214 x = layer(x=x, mem=m, c_mem=cm, mask=mask)
215 return self.norm(x), new_mem
216 return self.norm(x), new_mem
Attention Reconstruction Loss
-
Since this loss can be computed independently of the cross-entropy-loss of the model you can have a separate optimizer that only updates $f_c$. However, we use the same optimizer to update $f_c$ so when calculating -attention reconstruction loss we detach all other parameters except $f_c$ +attention reconstruction loss, we detach all other parameters except $f_c$ from the gradient computation.
218class AttentionReconstructionLoss:
219class AttentionReconstructionLoss:
layers
is the list of Compressive Transformer layers
236 def __init__(self, layers: TypedModuleList[CompressiveTransformerLayer]):
237 def __init__(self, layers: TypedModuleList[CompressiveTransformerLayer]):
240 self.layers = layers
-241 self.loss_func = nn.MSELoss()
241 self.layers = layers
+242 self.loss_func = nn.MSELoss()
243 def prepare_for_attn(self, pmha: PrepareForMultiHeadAttention, x: torch.Tensor):
244 def prepare_for_attn(self, pmha: PrepareForMultiHeadAttention, x: torch.Tensor):
Shape of the input except embedding dimension; [seq_len, batch_size]
.
253 head_shape = x.shape[:-1]
254 head_shape = x.shape[:-1]
Detach projection weights and bias
256 weight = pmha.linear.weight.detach()
-257 bias = pmha.linear.bias.detach() if pmha.linear.bias is not None else None
257 weight = pmha.linear.weight.detach()
+258 bias = pmha.linear.bias.detach() if pmha.linear.bias is not None else None
Linear transform
259 x = F.linear(x, weight, bias)
260 x = F.linear(x, weight, bias)
Split last dimension into heads
262 x = x.view(*head_shape, pmha.heads, pmha.d_k)
263 x = x.view(*head_shape, pmha.heads, pmha.d_k)
Output has shape [seq_len, batch_size, heads, d_k]
or [batch_size, d_model]
265 return x
266 return x
267 def attn(self, layer: RelativeMultiHeadAttention, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor):
268 def attn(self, layer: RelativeMultiHeadAttention, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor):
Calculate query, key and value projections
274 query = self.prepare_for_attn(layer.query, query)
-275 key = self.prepare_for_attn(layer.key, key)
-276 value = self.prepare_for_attn(layer.value, value)
275 query = self.prepare_for_attn(layer.query, query)
+276 key = self.prepare_for_attn(layer.key, key)
+277 value = self.prepare_for_attn(layer.value, value)
[seq_len, seq_len, batch_size, heads]
.
280 scores = torch.einsum('ibhd,jbhd->ijbh', query, key)
281 scores = torch.einsum('ibhd,jbhd->ijbh', query, key)
[seq_len, seq_len, batch_size, heads]
.
Scale scores $\frac{Q K^\top}{\sqrt{d_k}}$
283 scores *= layer.scale
284 scores *= layer.scale
[seq_len, seq_len, batch_size, heads]
.
$\underset{seq}{softmax}\Bigg(\frac{Q K^\top}{\sqrt{d_k}}\Bigg)$
287 attn = layer.softmax(scores)
288 attn = layer.softmax(scores)
291 return torch.einsum("ijbh,jbhd->ibhd", attn, value)
292 return torch.einsum("ijbh,jbhd->ibhd", attn, value)
Perform layer normalization with shift and scale parameters detached.
293 def norm(self, ln: nn.LayerNorm, x: torch.Tensor):
294 def norm(self, ln: nn.LayerNorm, x: torch.Tensor):
Detach shift(bias
) and scaling(weight
) parameters
299 weight = ln.weight.detach() if ln.weight is not None else None
-300 bias = ln.bias.detach() if ln.bias is not None else None
300 weight = ln.weight.detach() if ln.weight is not None else None
+301 bias = ln.bias.detach() if ln.bias is not None else None
Layer normalization
303 return F.layer_norm(x, ln.normalized_shape, weight, bias, ln.eps)
304 return F.layer_norm(x, ln.normalized_shape, weight, bias, ln.eps)
This calculates the loss for a layer
305 def calc_loss(self, layer: CompressiveTransformerLayer, h: torch.Tensor, mem: torch.Tensor):
306 def calc_loss(self, layer: CompressiveTransformerLayer, h: torch.Tensor, mem: torch.Tensor):
Detach the token embeddings and memory.
311 h = h.detach()
-312 mem = mem.detach()
312 h = h.detach()
+313 mem = mem.detach()
316 c_mem = layer.compress(mem)
317 c_mem = layer.compress(mem)
Normalize the embeddings and memories
319 h = self.norm(layer.norm_self_attn, h)
-320 mem = self.norm(layer.norm_self_attn, mem)
-321 c_mem = self.norm(layer.norm_self_attn, c_mem)
320 h = self.norm(layer.norm_self_attn, h)
+321 mem = self.norm(layer.norm_self_attn, mem)
+322 c_mem = self.norm(layer.norm_self_attn, c_mem)
Calculate attention with uncompressed memory
+Calculate the attention with uncompressed memory
324 attn_mem = self.attn(layer.self_attn, h, mem, mem)
325 attn_mem = self.attn(layer.self_attn, h, mem, mem)
Calculate the attention with compressed memory
326 attn_cmem = self.attn(layer.self_attn, h, c_mem, c_mem)
327 attn_cmem = self.attn(layer.self_attn, h, c_mem, c_mem)
Calculate the mean square error
329 return self.loss_func(attn_cmem, attn_mem)
330 return self.loss_func(attn_cmem, attn_mem)
331 def __call__(self, h: List[torch.Tensor], mem: List[torch.Tensor]):
332 def __call__(self, h: List[torch.Tensor], mem: List[torch.Tensor]):
Calculate the losses for each layer
333 losses = [self.calc_loss(layer, h[n], mem[n]) for n, layer in enumerate(self.layers)]
334 losses = [self.calc_loss(layer, h[n], mem[n]) for n, layer in enumerate(self.layers)]
Sum of the losses
335 return sum(losses)
336 return sum(losses)