11fromtypingimportList,Tuple,NamedTuple1213importtorch14importtorch.nnasnn1516fromlabmlimportexperiment,tracker,monit,logger17fromlabml.configsimportoption18fromlabml.loggerimportText19fromlabml\_helpers.metrics.simple\_stateimportSimpleStateModule20fromlabml\_helpers.moduleimportModule21fromlabml\_helpers.train\_validimportBatchIndex,hook\_model\_outputs22fromlabml\_nn.experiments.nlp\_autoregressionimportNLPAutoRegressionConfigs23fromlabml\_nn.transformers.compressiveimportCompressiveTransformer,AttentionReconstructionLoss,\24CompressiveTransformerLayer,Conv1dCompression

27classCompressedMemory(NamedTuple):28mem:List[torch.Tensor]29c\_mem:List[torch.Tensor]

32classAutoregressiveModel(Module):

37def\_\_init\_\_(self,n\_vocab:int,d\_model:int,transformer:CompressiveTransformer):38super().\_\_init\_\_()

40self.src\_embed=nn.Embedding(n\_vocab,d\_model)

42self.transformer=transformer

44self.generator=nn.Linear(d\_model,n\_vocab)

46self.mask\_x=None47self.mask\_mem=None

49defforward(self,x:torch.Tensor,mem:CompressedMemory):

51ifmemisnotNone:52mem,c\_mem=mem.mem,mem.c\_mem53else:54mem=[]55c\_mem=[]

58m\_len=len(mem[0])ifmemelse059ifc\_mem:60m\_len+=len(c\_mem[0])

63ifself.mask\_xisNoneorself.mask\_x.shape[0]\<len(x):64fromlabml\_nn.transformers.utilsimportsubsequent\_mask65self.mask\_x=subsequent\_mask(len(x)).to(x.device)

67ifself.mask\_memisNoneorself.mask\_mem.shape[1]\<m\_lenorself.mask\_mem.shape[0]\<len(x):68self.mask\_mem=self.mask\_x.new\_ones(len(x),m\_len,1)

71ifm\_len:72mask=torch.cat((self.mask\_mem[:len(x),:m\_len],self.mask\_x[:len(x),:len(x)]),dim=1)

74else:75mask=self.mask\_x[:len(x),:len(x)]

78x=self.src\_embed(x)

80res,mem=self.transformer(x,mem,c\_mem,mask)

82res=self.generator(res)

84returnres,mem

87classConfigs(NLPAutoRegressionConfigs):

94model:AutoregressiveModel

97d\_model:int=128

99heads:int=4

101dropout:float=0.0

103d\_ff:int=256

105n\_layers:int=6

107mem\_len:int=8

109memory=SimpleStateModule()

111attention\_reconstruction\_loss:AttentionReconstructionLoss

113compression\_rate:int=4

115c\_mem\_len:int=128

117definit(self):

119tracker.set\_scalar("accuracy.\*",True)120tracker.set\_scalar("loss.\*",True)

122tracker.set\_scalar("ar\_loss.\*",False)

124hook\_model\_outputs(self.mode,self.model,'model')

126self.state\_modules=[self.accuracy,self.memory]

[email protected]\_grad()129defmerge\_compress\_memory(self,mem:CompressedMemory,new\_mem:List[torch.Tensor])\130-\>Tuple[CompressedMemory,List[torch.Tensor]]:

136ifself.mem\_len==0andself.c\_mem\_len==0:137returnCompressedMemory([],[]),[]

140ifmemisnotNone:141mem,c\_mem=mem.mem,mem.c\_mem142else:143mem,c\_mem=[],[]

146ifmem:147mem=[torch.cat((m,x),dim=0)form,xinzip(mem,new\_mem)]148else:149mem=new\_mem

152iflen(mem[0])\>self.mem\_len:

156n\_c\_mem=(len(mem[0])-self.mem\_len+self.compression\_rate-1)//self.compression\_rate

158n\_old=n\_c\_mem\*self.compression\_rate

160mem\_to\_compress=[]

162uncompressed\_mem=[]

164forminmem:

166cm,m=torch.split(m,[n\_old,len(m)-n\_old])

168mem\_to\_compress.append(cm)

170uncompressed\_mem.append(m)

172mem=uncompressed\_mem

175new\_c\_mem=[]176fori,layerinenumerate(self.model.transformer.layers):177new\_c\_mem.append(layer.compress(mem\_to\_compress[i]))

180ifc\_mem:181c\_mem=[torch.cat((m,nm),dim=0)form,nminzip(c\_mem,new\_c\_mem)]

183else:184c\_mem=new\_c\_mem

187iflen(c\_mem[0])\>self.c\_mem\_len:188c\_mem=[m[-self.c\_mem\_len:]forminc\_mem]

190else:191mem\_to\_compress=[]

195returnCompressedMemory(mem,c\_mem),mem\_to\_compress

197defstep(self,batch:any,batch\_idx:BatchIndex):

203data,target=batch[0].to(self.device),batch[1].to(self.device)

206ifself.mode.is\_train:207tracker.add\_global\_step(data.shape[0]\*data.shape[1])

210withself.mode.update(is\_log\_activations=batch\_idx.is\_last):

212mem=self.memory.get()

214output,new\_mem=self.model(data,mem)

216mem,mem\_to\_compress=self.merge\_compress\_memory(mem,new\_mem)

218self.memory.set(mem)

221loss=self.loss\_func(output,target)222tracker.add("loss.",loss)

225ifmem\_to\_compress:

227ar\_loss=self.attention\_reconstruction\_loss(new\_mem,mem\_to\_compress)

229tracker.add("ar\_loss.",ar\_loss)

231loss=loss+ar\_loss

234self.accuracy(output,target)235self.accuracy.track()

238ifself.mode.is\_train:

240loss.backward()

242torch.nn.utils.clip\_grad\_norm\_(self.model.parameters(),max\_norm=self.grad\_norm\_clip)

244self.optimizer.step()

246ifbatch\_idx.is\_last:247tracker.add('model',self.model)

249self.optimizer.zero\_grad()

252tracker.save()

254defsample(self):

260prompt=self.prompt

262log=[(prompt,Text.subtle)]

264mem=CompressedMemory([],[])

266foriinmonit.iterate('Sample',25):

268data=self.text.text\_to\_i(prompt).unsqueeze(-1)

270data=data.to(self.device)

272output,new\_mem=self.model(data,mem)

274output=output.argmax(dim=-1).squeeze(1)

276prompt+=self.prompt\_separator+self.text.itos[output[-1]]

278prompt=prompt[-1:]

280log+=[(self.prompt\_separator+self.text.itos[output[-1]],Text.value)]

282mem,\_=self.merge\_compress\_memory(mem,new\_mem)

285logger.log(log)

288@option(Configs.model)289defautoregressive\_model(c:Configs):

293fromlabml\_nn.transformers.xlimportRelativeMultiHeadAttention294fromlabml\_nn.transformers.feed\_forwardimportFeedForward295m=AutoregressiveModel(c.n\_tokens,c.d\_model,CompressiveTransformer(296CompressiveTransformerLayer(d\_model=c.d\_model,297self\_attn=RelativeMultiHeadAttention(c.heads,c.d\_model,c.dropout),298feed\_forward=FeedForward(c.d\_model,c.d\_ff,c.dropout),299dropout\_prob=c.dropout,300compress=Conv1dCompression(c.compression\_rate,c.d\_model)),c.n\_layers))301returnm.to(c.device)

304@option(Configs.attention\_reconstruction\_loss)305defattention\_reconstruction\_loss(c:Configs):

309returnAttentionReconstructionLoss(c.model.transformer.layers)

312defmain():

317experiment.create(name="compressive\_transformer",comment='')

319conf=Configs()

321experiment.configs(conf,

323{'tokenizer':'character',324'text':'tiny\_shakespeare',325'optimizer.learning\_rate':2.5e-4,326'optimizer.optimizer':'AdamW',327'prompt':'It is',328'prompt\_separator':'',329330'train\_loader':'sequential\_train\_loader',331'valid\_loader':'sequential\_valid\_loader',332333'seq\_len':8,334'mem\_len':8,335'epochs':128,336'batch\_size':32,337'inner\_iterations':25,338'compression\_rate':2,339})

342experiment.add\_pytorch\_models({'model':conf.model})

345withexperiment.start():

347conf.run()

351if\_\_name\_\_=='\_\_main\_\_':352main()