class WeightSharePositionwiseFF(nn.Module): def __init__(self, d_model, d_inner, dropout, pre_lnorm=False): super(WeightSharePositionwiseFF, self).__init__() self.d_model = d_model self.d_inner = d_inner self.dropout = dropout self.ff1_net = nn.Conv1d(d_model, d_inner, kernel_size=1) self.drop1 = VariationalHidDropout(dropout=dropout) self.ff2_net = nn.Conv1d(d_inner, d_model, kernel_size=1) self.drop2 = VariationalHidDropout(dropout=dropout) self.pre_lnorm = pre_lnorm def wnorm(self): print("Weight normalization applied to PosFF") self.ff1_net, self.ff1_fn = weight_norm(module=self.ff1_net, names=['weight'], dim=0) self.ff2_net, self.ff2_fn = weight_norm(module=self.ff2_net, names=['weight'], dim=0) def reset(self, bsz, qlen): self.drop1.reset_mask(torch.zeros(bsz, self.d_inner, qlen)) self.drop2.reset_mask(torch.zeros(bsz, self.d_model, qlen)) if 'ff1_fn' in self.__dict__: self.ff1_fn.reset(self.ff1_net) if 'ff2_fn' in self.__dict__: self.ff2_fn.reset(self.ff2_net) def copy(self, func): self.ff1_net.weight.data = func.ff1_net.weight.data.clone() self.ff2_net.weight.data = func.ff2_net.weight.data.clone() self.ff1_net.bias.data = func.ff1_net.bias.data.clone() self.ff2_net.bias.data = func.ff2_net.bias.data.clone() self.drop1.mask = func.drop1.mask.clone() self.drop2.mask = func.drop2.mask.clone() def forward(self, inp, attn_out=None): assert inp.size(1) == self.d_model, "Feature dimension not match!!" if self.pre_lnorm: inp = F.layer_norm(inp.transpose(1, 2), (self.d_model, )).transpose(1, 2) relu_out1 = self.drop1(F.relu(self.ff1_net(inp))) out2 = self.drop2(self.ff2_net(relu_out1)) output = out2 + inp if not self.pre_lnorm: output = F.layer_norm(output.transpose(1, 2), (self.d_model, )).transpose(1, 2) return output
class WeightShareSelfAttention(nn.Module): # This is similar to the RelPartialLearnableMultiHeadAttn class in Transformer-XL def __init__(self, d_model, n_head, d_head, dropout, dropatt, pre_lnorm=False, local_size=None): super(WeightShareSelfAttention, self).__init__() self.d_model = d_model self.n_head = n_head self.d_head = d_head self.dropout = dropout self.scale = 1 / (d_head ** 0.5) self.qkv_net = nn.Conv1d(d_model, 3 * n_head * d_head, kernel_size=1, bias=False) self.r_net = nn.Conv1d(d_model, n_head * d_head, kernel_size=1, bias=False) self.r_w_bias = nn.Parameter(torch.rand(n_head, d_head).uniform_(-0.05, 0.05)) self.r_r_bias = nn.Parameter(torch.rand(n_head, d_head).uniform_(-0.05, 0.05)) self.o_net = nn.Conv1d(n_head * d_head, d_model, kernel_size=1) self.dropatt = VariationalAttnDropout(dropout=dropatt) self.drop = VariationalHidDropout(dropout=dropout) self.pre_lnorm = pre_lnorm self.local_size = local_size def wnorm(self): print("Weight normalization applied to SA") self.qkv_net, self.qkv_fn = weight_norm(module=self.qkv_net, names=['weight'], dim=0) self.r_net, self.r_fn = weight_norm(module=self.r_net, names=['weight'], dim=0) self.o_net, self.o_fn = weight_norm(module=self.o_net, names=['weight'], dim=0) def reset(self, bsz, qlen, klen): self.dropatt.reset_mask(torch.zeros(bsz, self.n_head, qlen, klen)) self.drop.reset_mask(torch.zeros(bsz, self.d_model, qlen)) if 'qkv_fn' in self.__dict__: self.qkv_fn.reset(self.qkv_net) if 'r_fn' in self.__dict__: self.r_fn.reset(self.r_net) if 'o_fn' in self.__dict__: self.o_fn.reset(self.o_net) def copy(self, func): # Destructive copy self.qkv_net.weight.data = func.qkv_net.weight.data.clone() self.r_net.weight.data = func.r_net.weight.data.clone() self.r_w_bias.data = func.r_w_bias.data.clone() self.r_r_bias.data = func.r_r_bias.data.clone() self.o_net.weight.data = func.o_net.weight.data.clone() self.o_net.bias.data = func.o_net.bias.data.clone() self.dropatt.mask = func.dropatt.mask.clone() self.drop.mask = func.drop.mask.clone() def _rel_shift(self, x): # x has dimension (bsz x n_head x qlen x klen) bsz, n_head, qlen, klen = x.size() x_padded = F.pad(x, (1,0)) x_padded = x_padded.view(bsz, n_head, klen+1, qlen) return x_padded[:,:,1:].view_as(x) def forward(self, z1ss, pos_emb, u1ss, mems=None): # Note: In this context, qlen means the length of the (small) subsequence; and mlen describes # the length of the padding. Their sum is klen. bsz, d_model, qlen = z1ss.size() r_w_bias, r_r_bias = self.r_w_bias, self.r_r_bias n_head, d_head = self.n_head, self.d_head rlen = pos_emb.size(2) if mems is None: mems = torch.tensor([]).view(0,0,0) mlen = mems.size(2) cat = torch.cat([mems, z1ss], dim=-1) if self.pre_lnorm: cat = F.layer_norm(cat.transpose(1,2), (d_model,)).transpose(1,2) w_heads = self.qkv_net(cat) # (N, 3C, L) r_head_k = self.r_net(pos_emb) # Input injection w_heads += u1ss w_head_q, w_head_k, w_head_v = torch.chunk(w_heads, 3, dim=1) w_head_q = w_head_q[:,:,-qlen:] klen = w_head_k.size(2) w_head_q = w_head_q.view(bsz, n_head, d_head, qlen) # bsz x n_head x d_head x qlen w_head_k = w_head_k.view(bsz, n_head, d_head, klen) # bsz x n_head x d_head x klen w_head_v = w_head_v.view(bsz, n_head, d_head, klen) # bsz x n_head x d_head x klen r_head_k = r_head_k.view(n_head, d_head, rlen) # n_head x d_head x rlen #### compute attention score rw_head_q = w_head_q + r_w_bias[:,:,None] # bsz x n_head x d_head x qlen AC = torch.einsum('bndi,bndj->bnij', rw_head_q, w_head_k) rr_head_q = w_head_q + r_r_bias[:,:,None] BD = torch.einsum('bndi,ndj->bnij', rr_head_q, r_head_k) BD = self._rel_shift(BD) # for the sake of relative positional embedding attn_score = AC + BD # bsz x n_head x qlen x klen attn_score.mul_(self.scale) #### compute attention probability # We apply a local mask, with local horizon size of mlen local_size = self.local_size or 1000 attn_mask = torch.triu(torch.ones(qlen, klen), diagonal=1+mlen).byte()[None,:,:] attn_mask += torch.tril(torch.ones(qlen, klen), diagonal=mlen-local_size).byte()[None,:,:] if attn_mask is not None and attn_mask.any().item(): attn_score = attn_score.float().masked_fill( attn_mask[None,:,:,:], -float('inf')).type_as(attn_score) attn_prob = F.softmax(attn_score, dim=-1) # bsz x n_head x qlen x klen #### compute attention vector attn_vec = torch.einsum('bnij,bndj->bndi', (attn_prob, w_head_v)) # [bsz x d x qlen] attn_vec = attn_vec.contiguous().view(bsz, n_head*d_head, attn_vec.size(-1)) ##### linear projection attn_out = self.o_net(attn_vec) attn_out = self.drop(attn_out) ##### residual connection + layer normolization (if applicable) if self.pre_lnorm: out = attn_out + z1ss else: out = F.layer_norm((attn_out + z1ss).transpose(1,2), (d_model,)).transpose(1,2) return out
class DEQTransformerLM(nn.Module): def __init__(self, n_token, n_layer, eval_n_layer, n_head, d_model, d_head, d_inner, dropout, dropatt, tie_weights=True, d_embed=None, div_val=1, tie_projs=[False], pre_lnorm=False, wnorm=False, tgt_len=None, mem_len=None, local_size=0, pretrain_steps=1, cutoffs=[], load=''): super().__init__() self.n_token = n_token d_embed = d_model if d_embed is None else d_embed self.d_embed = d_embed self.d_model = d_model self.n_head = n_head self.d_head = d_head self.word_emb = AdaptiveEmbedding(n_token, d_embed, d_model, cutoffs, div_val=div_val) self.iodrop = VariationalDropout() self.dropout = dropout self.drop = VariationalHidDropout(dropout=dropout) self.pretrain_steps = pretrain_steps self.tgt_len = tgt_len self.mem_len = mem_len self.local_size = local_size self.max_klen = tgt_len + mem_len self.n_layer = n_layer self.eval_n_layer = eval_n_layer self.inject_conv = nn.Conv1d(d_model, 3*d_model, kernel_size=1) self.pos_emb = PositionalEmbedding(self.d_model) self.func = RelPartialLearnableDecoderLayer(n_head, d_model, d_head, d_inner, dropout=dropout, dropatt=dropatt, pre_lnorm=pre_lnorm, local_size=local_size) self.func_copy = copy.deepcopy(self.func) if wnorm: # If wnorm is specified, we need to make sure to do the deepcopy first # because pytorch prevents non-user defined variables from being copied. # The deepcopy will only access and copy the `xxx.weight` instead # of messing with `xxx.weight_g` that wnorm actually optimizes. self.func.wnorm() for params in self.func_copy.parameters(): params.requires_grad_(False) # Turn off autograd for func_copy self.deq = TransformerDEQForward(self.func) self.deqback = TransformerDEQBackward(self.func, self.func_copy) # use adaptive softmax (including standard softmax) # (Note: To use sample softmax, refer to the Transformer-XL implementation) self.crit = ProjectedAdaptiveLogSoftmax(n_token, d_embed, d_model, cutoffs, div_val=div_val) if tie_weights: for i in range(len(self.crit.out_layers)): self.crit.out_layers[i].weight = self.word_emb.emb_layers[i].weight if tie_projs: for i, tie_proj in enumerate(tie_projs): if tie_proj and div_val == 1 and d_model != d_embed: self.crit.out_projs[i] = self.word_emb.emb_projs[0] elif tie_proj and div_val != 1: self.crit.out_projs[i] = self.word_emb.emb_projs[i] if len(load) > 0: params_dict = torch.load(load) self.load_weights(params_dict) print(f"Finished loading. d_embed={self.inject_conv.weight.data.size(1)}") def reset_length(self, tgt_len, mem_len): self.tgt_len = tgt_len self.mem_len = mem_len def load_weights(self, params_dict): self.load_state_dict(params_dict) def save_weights(self, name='pretrained_deq'): with open(f'{name}.pkl', 'wb') as f: print(f"Saving weights at {name}.pkl") torch.save(self.state_dict(), f) def init_mems(self): if self.mem_len <= 0: print("init_mems: Hmmmm... you shouldn't be here.") return None # mems is not None with torch.no_grad(): param = next(self.parameters()) mems = [torch.empty(0, dtype=param.dtype, device=param.device), torch.empty(0, dtype=param.dtype, device=param.device)] return mems # For z0 and u0 def _update_mems(self, z1s, us, z0, qlen, mlen): # does not deal with None if self.mem_len <= 0: print("_update_mems: Hmmmm... you shouldn't be here.") return None # mems is not None with torch.no_grad(): end_idx = mlen + qlen beg_idx = max(0, end_idx - self.mem_len) # Account for when mlen = 0 zs = torch.cat([z0, z1s], dim=2) new_z0 = zs[:,:,beg_idx:end_idx].detach().permute(2,0,1).contiguous() # seq_len x bsz x d_model new_u0 = us[:,:,beg_idx:end_idx].detach().permute(2,0,1).contiguous() return [new_z0, new_u0] def _forward(self, dec_inp, subseq_len, mems=None, f_thres=30, b_thres=40, train_step=-1): """ Apply the DEQ-Transformer language model on input word tokens :param dec_inp: Input words of shape (seq_len x bsz) and dtype torch.LongTensor :param subseq_len: The subsequence length with which we feed the segments of the data to DEQ :param mems: History madding and the transformed input corresponding to it; must be a tuple (z0, u0) where z0 has dimension (bsz x d_model x pad_len) and u0 has size (bsz x 3*d_model x pad_len) :param f_thres: Forward pass threshold :param b_thres: Backward pass threshold :param train_step: The number of training step that the current iteration is at :return: tuple(output sequence, new memory) """ # Assume dec_inp has shape (qlen x bsz) dec_inp = dec_inp.t() bsz, qlen = dec_inp.size() word_emb = self.word_emb(dec_inp) word_emb = self.iodrop(word_emb, self.dropout) u1s = self.inject_conv(word_emb.transpose(1,2)) # bsz x 3*d_model x qlen z0, u0 = mems d_model = self.d_model if z0 is not None and z0.nelement() > 0: assert z0.size(2) == u0.size(2), "Padding fixed points and padding embedding dimensions don't agree" else: z0 = torch.zeros(bsz, d_model, 0) u0 = torch.zeros(bsz, 3*d_model, 0) mlen = z0.size(2) klen = mlen + qlen # qlen is seq_len, mlen is pad_len pos_seq = torch.arange(klen-1, -1, -1.0) pos_emb = self.drop(self.pos_emb(pos_seq)) # bsz x d_model x (qlen + mlen) for positional embedding us = torch.cat([u0, u1s], dim=2) z1s = torch.zeros(bsz, d_model, qlen) # bsz x d_model x (qlen + mlen) for initial estimate of output if 0 <= train_step < self.pretrain_steps: # In pretraining mode with stacked (weight-tied) layers. NOT recommended for large models (as then # a stacking of, for example, 16 layers would be extremely inefficient). One can also train with # M layers and evaluate using N layers (which typically leads to worse performance). n_layer = self.n_layer if self.training or train_step > 0 else self.eval_n_layer torch.cuda.empty_cache() for i in range(n_layer): z1s = self.func(z1s, us, z0, pos_emb) else: # Compute the equilibrium via DEQ. When in training mode, we need to register the analytical backward # pass according to the Theorem 1 in the paper. z1s = self.deq(z1s, us, z0, pos_emb=pos_emb, subseq_len=subseq_len, threshold=f_thres, train_step=train_step) if self.training: z1s = self.deqback(z1s, us, z0, pos_emb=pos_emb, subseq_len=subseq_len, threshold=b_thres, train_step=train_step) core_out = self.iodrop(z1s, self.dropout) core_out = core_out.permute(2,0,1).contiguous() # qlen x bsz x d_model new_mems = self._update_mems(z1s, us, z0, mlen, qlen) return core_out, new_mems def forward(self, data, target, mems, train_step=-1, **kwargs): # nn.DataParallel does not allow size(0) tensors to be broadcasted. # So, have to initialize size(0) mems inside the model forward. # Moreover, have to return new_mems to allow nn.DataParallel to piece # them together. if not mems: mems = self.init_mems() else: for i in range(len(mems)): mems[i] = mems[i].permute(1,2,0).contiguous() # bsz x [-1] x seq_len qlen, bsz = data.size() mlen = 0 if mems[0].nelement() == 0 else mems[0].size(2) klen = mlen + qlen # Reset dropout in self.func, and copy everything (weights, dropout masks) to self.func_copy self.drop.reset_mask(torch.zeros(1, self.d_model, klen)) self.func.reset(bsz, qlen, klen) self.func_copy.copy(self.func) tgt_len = target.size(0) subseq_len = kwargs.get('subseq_len', 75) f_thres = kwargs.get('f_thres', 30) b_thres = kwargs.get('b_thres', 40) hidden, new_mems = self._forward(data, subseq_len=subseq_len, mems=mems, f_thres=f_thres, b_thres=b_thres, train_step=train_step) pred_hid = hidden[-tgt_len:] loss = self.crit(pred_hid.view(-1, pred_hid.size(-1)), target.view(-1)) loss = loss.view(tgt_len, -1) if new_mems is None: return [loss] else: return [loss] + new_mems