CITE THIS NOTEBOOK: Transformers: LMs, sequence classifiers and Vision by Nicolò Monti. Wolfram Community MAR 6 2023.
Since its introduction in 2017, Transformers has revolutionised almost every area of Machine Learning, where we have seen a sudden convergence of all architectures towards them.
Note: This Notebook assumes that the reader is already familiar with the topics covered.
Case-study: Parentheses matching
Case-study: Parentheses matching
To explore how we can make use of Transformers in the Wolfram Language, we will train an autoregressive model on parenthesis patterns.
We start by defining the net encoder and (also used as decoder), which will act as our “bridge” between the input/output modality and the model. In this case, we’re only interested in processing two tokens, “(“ and “)”.
In[]:=
tokenizerLM=NetEncoder[{"Characters",{"(",")"}}]
Out[]=
NetEncoder
In[]:=
detokenizerLM=NetDecoder[{"Class",{"(",")"}}]
Out[]=
NetDecoder
We now need to define some hyperparameters for our model, which are completely arbitrary and have been obtained empirically.
In[]:=
device="GPU";
In[]:=
tokensDimLM=2;predictorDropoutRate=0.64;sequenceLength=64;hiddenDim=64;headsDim=32;
We can now move on to the very first part of our model: The encoding. This block involves both embedding the tokens in the latent space and providing them with a positional encoding. We do this by simply learning an embedding for both the tokens and the position indices of the words and adding them up.
In[]:=
positionalEncoding=NetGraph[{"indices"->SequenceIndicesLayer[],"positionEmbedding"->EmbeddingLayer[hiddenDim,sequenceLength],"plus"->ThreadingLayer[Plus],"normalize"->BatchNormalizationLayer[Interleaving->True],"dropout"->DropoutLayer[predictorDropoutRate]},{NetPort["Input"]->"indices"->"positionEmbedding",{NetPort["Input"],"positionEmbedding"}->"plus"->"normalize"->"dropout"},"Input"->{"Varying",hiddenDim},"Output"->{"Varying",hiddenDim}]
Out[]=
NetGraph
In[]:=
embeddingPositionalEncoding[tokensDim0_Integer]:=Module[{tokensDim=tokensDim0},(NetGraph[{"indices"->SequenceIndicesLayer[],"wordEmbedding"->EmbeddingLayer[hiddenDim,tokensDim],"positionEmbedding"->EmbeddingLayer[hiddenDim,sequenceLength],"plus"->ThreadingLayer[Plus],"normalize"->BatchNormalizationLayer[Interleaving->True],"dropout"->DropoutLayer[predictorDropoutRate]},{NetPort["Input"]->"indices"->"positionEmbedding",NetPort["Input"]->"wordEmbedding",{"wordEmbedding","positionEmbedding"}->"plus"->"normalize"->"dropout"},"Input"->"Varying","Output"->{"Varying",hiddenDim}])]
Now we implement one of the most important parts of the architecture, the pointwiseBlock. The pointwiseBlock is applied to the embedding of each token individually, and is thus the part of the model that actually learns the task (along with the Attention Mechanisms).
In[]:=
pointwiseBlock=NetChain[{LinearLayer[hiddenDim],Ramp,DropoutLayer[predictorDropoutRate],BatchNormalizationLayer[Interleaving->True],LinearLayer[hiddenDim],Ramp,DropoutLayer[predictorDropoutRate],BatchNormalizationLayer[Interleaving->True],LinearLayer[hiddenDim],Ramp,DropoutLayer[predictorDropoutRate],BatchNormalizationLayer[Interleaving->True]},"Input"->hiddenDim,"Output"->hiddenDim]
Out[]=
NetChain
The selfAttentionCausalBlock is concerned with making sure that each embedding can “communicate” with each other, allowing the model to consider each piece of input as it is processed.
In[]:=
selfAttentionBlock=NetGraph[{"key"->NetMapOperator[LinearLayer[{headsDim,hiddenDim}]],"value"->NetMapOperator[LinearLayer[{headsDim,hiddenDim}]],"query"->NetMapOperator[LinearLayer[{headsDim,hiddenDim}]],"attention"->AttentionLayer["Dot","MultiHead"->True],"merge"->NetMapOperator[hiddenDim]},{"key"->NetPort["attention","Key"],"value"->NetPort["attention","Value"],"query"->NetPort["attention","Query"],"attention"->"merge"},"Input"->{"Varying",hiddenDim},"Output"->{"Varying",hiddenDim}]
Out[]=
NetGraph
In[]:=
selfAttentionCausalBlock=NetGraph[{"key"->NetMapOperator[LinearLayer[{headsDim,hiddenDim}]],"value"->NetMapOperator[LinearLayer[{headsDim,hiddenDim}]],"query"->NetMapOperator[LinearLayer[{headsDim,hiddenDim}]],"attention"->AttentionLayer["Dot","MultiHead"->True,"Mask"->"Causal"],"merge"->NetMapOperator[hiddenDim]},{"key"->NetPort["attention","Key"],"value"->NetPort["attention","Value"],"query"->NetPort["attention","Query"],"attention"->"merge"},"Input"->{"Varying",hiddenDim},"Output"->{"Varying",hiddenDim}]
Out[]=
NetGraph
In the decoderBlock, we combine all our parts (except the embedding) to form a unit, which we then stack to form the main part of our model.
In[]:=
decoderBlock=NetGraph[{"selfAttention"->selfAttentionBlock,"plus1"->ThreadingLayer[Plus],"normalize1"->BatchNormalizationLayer["Interleaving"->True],"pointwise"->NetMapOperator[pointwiseBlock],"plus2"->ThreadingLayer[Plus],"normalize2"->BatchNormalizationLayer[Interleaving->True]},{NetPort["Input"]->"selfAttention",{NetPort["Input"],"selfAttention"}->"plus1"->"normalize1"->"pointwise",{"normalize1","pointwise"}->"plus2"->"normalize2"},"Input"->{"Varying",hiddenDim},"Output"->{"Varying",hiddenDim}]
Out[]=
NetGraph
In[]:=
decoderCausalBlock=NetGraph[{"selfAttention"->selfAttentionCausalBlock,"plus1"->ThreadingLayer[Plus],"normalize1"->BatchNormalizationLayer["Interleaving"->True],"pointwise"->NetMapOperator[pointwiseBlock],"plus2"->ThreadingLayer[Plus],"normalize2"->BatchNormalizationLayer[Interleaving->True]},{NetPort["Input"]->"selfAttention",{NetPort["Input"],"selfAttention"}->"plus1"->"normalize1"->"pointwise",{"normalize1","pointwise"}->"plus2"->"normalize2"},"Input"->{"Varying",hiddenDim},"Output"->{"Varying",hiddenDim}]
Out[]=
NetGraph
Thanks to the classificationHead, we can turn any latent representation of a token into a prediction and thus complete our model.
In[]:=
classificationHeadLM=NetChain[{LinearLayer[hiddenDim],Ramp,DropoutLayer[predictorDropoutRate],LinearLayer[hiddenDim],Ramp,DropoutLayer[predictorDropoutRate],BatchNormalizationLayer[],LinearLayer[tokensDimLM],SoftmaxLayer[]}]
Out[]=
NetChain
Finally, we can stack the Transformer into a single piece.
In[]:=
transformerLM=NetChain[{embeddingPositionalEncoding[tokensDimLM],decoderCausalBlock,decoderCausalBlock,NetMapOperator[classificationHeadLM]},"Input"->"Varying"]
Here we create an auxiliary model to help train the Transformer on an autoregressive task.
The next step is to generate the dataset, we can generate a repeated pattern of matches parentheses.
We can now actually train the model to predict the sequences provided in the dataset.
Warning: “Mixed” is only supported on GPU.
We can define a function to call our trained model.
Building on our previous function, we can define a new one to autoregressively complete our prompt.
Case-study: Sentiment analysis
Case-study: Sentiment analysis
We can now explore how to build a bi-directional Transformer model, we’ll do this by tackling the task of sentiment analysis.
As in our previous case-study, we have to start by defining the net encoder and decoder.
Now we have to define a new classification head.
We can now put all of the pieces together.
Our dataset will map a sentence to a negative/positive sentiment.
Then, we can train our model.
Warning: “Mixed” is only supported on GPU.
Case-study: Vision Transformer
Case-study: Vision Transformer
We can now explore how to build a Vision Transformer to perform image classification.
We first have to load the dataset and define some utility functions.
We can now proceed with defining the net encoder.
Now we have to define a new classification head.
We can now put all of the pieces together.
Then, we can train our model.
Warning: “Mixed” is only supported on GPU.
We can now evaluate the performances of the model on the test set.