![](\"https://github.com/Lightning-AI/lightning-tutorials/raw/main/course_UvA-DL/05-transformers-and-MH-attention/attention_example.svg\")
![](\"https://github.com/Lightning-AI/lightning-tutorials/raw/main/course_UvA-DL/05-transformers-and-MH-attention/scaled_dot_product_attn.svg\")
![](\"https://github.com/Lightning-AI/lightning-tutorials/raw/main/course_UvA-DL/05-transformers-and-MH-attention/multihead_attention.svg\")
![](\"https://github.com/Lightning-AI/lightning-tutorials/raw/main/course_UvA-DL/05-transformers-and-MH-attention/comparison_conv_rnn.svg\")
\n", "\n", "Next, we will look at how to apply the multi-head attention blog inside the Transformer architecture.\n", "Originally, the Transformer model was designed for machine translation.\n", "Hence, it got an encoder-decoder structure where the encoder takes as input the sentence in the original language\n", "and generates an attention-based representation.\n", "On the other hand, the decoder attends over the encoded information and generates the translated sentence\n", "in an autoregressive manner, as in a standard RNN.\n", "While this structure is extremely useful for Sequence-to-Sequence tasks with the necessity of autoregressive decoding,\n", "we will focus here on the encoder part.\n", "Many advances in NLP have been made using pure encoder-based Transformer models (if interested, models include the\n", "[BERT](https://arxiv.org/abs/1810.04805)-family,\n", "the [Vision Transformer](https://arxiv.org/abs/2010.11929), and more),\n", "and in our tutorial, we will also mainly focus on the encoder part.\n", "If you have understood the encoder architecture, the decoder is a very small step to implement as well.\n", "The full Transformer architecture looks as follows\n", "(figure credit - [Vaswani et al., 2017](https://arxiv.org/abs/1706.03762)).\n", ":\n", "\n", "
![](\"https://github.com/Lightning-AI/lightning-tutorials/raw/main/course_UvA-DL/05-transformers-and-MH-attention/transformer_architecture.svg\"){kind=spacer}
![](\"https://github.com/Lightning-AI/lightning-tutorials/raw/main/course_UvA-DL/05-transformers-and-MH-attention/warmup_loss_plot.svg\")
\n", "\n", "After having finished the implementation of the Transformer architecture, we can start experimenting\n", "and apply it to various tasks.\n", "In this notebook, we will focus on two tasks: parallel Sequence-to-Sequence, and set anomaly detection.\n", "The two tasks focus on different properties of the Transformer architecture, and we go through them below.\n", "\n", "### Sequence to Sequence\n", "\n", "A Sequence-to-Sequence task represents a task where the input _and_ the output is a sequence,\n", "not necessarily of the same length.\n", "Popular tasks in this domain include machine translation and summarization.\n", "For this, we usually have a Transformer encoder for interpreting the input sequence,\n", "and a decoder for generating the output in an autoregressive manner.\n", "Here, however, we will go back to a much simpler example task and use only the encoder.\n", "Given a sequence of $N$ numbers between $0$ and $M$, the task is to reverse the input sequence.\n", "In Numpy notation, if our input is $x$, the output should be $x$[::-1].\n", "Although this task sounds very simple, RNNs can have issues with such because the task requires long-term dependencies.\n", "Transformers are built to support such, and hence, we expect it to perform very well.\n", "\n", "First, let's create a dataset class below."]}, {"cell_type": "code", "execution_count": 15, "id": "b4fa2b28", "metadata": {"execution": {"iopub.execute_input": "2025-04-03T19:15:35.326194Z", "iopub.status.busy": "2025-04-03T19:15:35.325389Z", "iopub.status.idle": "2025-04-03T19:15:35.332545Z", "shell.execute_reply": "2025-04-03T19:15:35.331610Z"}, "papermill": {"duration": 0.024055, "end_time": "2025-04-03T19:15:35.333844", "exception": false, "start_time": "2025-04-03T19:15:35.309789", "status": "completed"}, "tags": []}, "outputs": [], "source": ["class ReverseDataset(data.Dataset):\n", " def __init__(self, num_categories, seq_len, size):\n", " super().__init__()\n", " self.num_categories = num_categories\n", " self.seq_len = seq_len\n", " self.size = size\n", "\n", " self.data = torch.randint(self.num_categories, size=(self.size, self.seq_len))\n", "\n", " def __len__(self):\n", " return self.size\n", "\n", " def __getitem__(self, idx):\n", " inp_data = self.data[idx]\n", " labels = torch.flip(inp_data, dims=(0,))\n", " return inp_data, labels"]}, {"cell_type": "markdown", "id": "685b4287", "metadata": {"papermill": {"duration": 0.014768, "end_time": "2025-04-03T19:15:35.363575", "exception": false, "start_time": "2025-04-03T19:15:35.348807", "status": "completed"}, "tags": []}, "source": ["We create an arbitrary number of random sequences of numbers between 0 and `num_categories-1`.\n", "The label is simply the tensor flipped over the sequence dimension.\n", "We can create the corresponding data loaders below."]}, {"cell_type": "code", "execution_count": 16, "id": "93c8c5d1", "metadata": {"execution": {"iopub.execute_input": "2025-04-03T19:15:35.394416Z", "iopub.status.busy": "2025-04-03T19:15:35.394078Z", "iopub.status.idle": "2025-04-03T19:15:35.427865Z", "shell.execute_reply": "2025-04-03T19:15:35.426849Z"}, "papermill": {"duration": 0.050881, "end_time": "2025-04-03T19:15:35.429192", "exception": false, "start_time": "2025-04-03T19:15:35.378311", "status": "completed"}, "tags": []}, "outputs": [], "source": ["dataset = partial(ReverseDataset, 10, 16)\n", "train_loader = data.DataLoader(dataset(50000), batch_size=128, shuffle=True, drop_last=True, pin_memory=True)\n", "val_loader = data.DataLoader(dataset(1000), batch_size=128)\n", "test_loader = data.DataLoader(dataset(10000), batch_size=128)"]}, {"cell_type": "markdown", "id": "d0905e8b", "metadata": {"papermill": {"duration": 0.015126, "end_time": "2025-04-03T19:15:35.459276", "exception": false, "start_time": "2025-04-03T19:15:35.444150", "status": "completed"}, "tags": []}, "source": ["Let's look at an arbitrary sample of the dataset:"]}, {"cell_type": "code", "execution_count": 17, "id": "a8cc2297", "metadata": {"execution": {"iopub.execute_input": "2025-04-03T19:15:35.490598Z", "iopub.status.busy": "2025-04-03T19:15:35.490154Z", "iopub.status.idle": "2025-04-03T19:15:35.497096Z", "shell.execute_reply": "2025-04-03T19:15:35.496164Z"}, "papermill": {"duration": 0.024052, "end_time": "2025-04-03T19:15:35.498332", "exception": false, "start_time": "2025-04-03T19:15:35.474280", "status": "completed"}, "tags": []}, "outputs": [{"name": "stdout", "output_type": "stream", "text": ["Input data: tensor([9, 6, 2, 0, 6, 2, 7, 9, 7, 3, 3, 4, 3, 7, 0, 9])\n", "Labels: tensor([9, 0, 7, 3, 4, 3, 3, 7, 9, 7, 2, 6, 0, 2, 6, 9])\n"]}], "source": ["inp_data, labels = train_loader.dataset[0]\n", "print(\"Input data:\", inp_data)\n", "print(\"Labels: \", labels)"]}, {"cell_type": "markdown", "id": "c774b751", "metadata": {"lines_to_next_cell": 2, "papermill": {"duration": 0.014845, "end_time": "2025-04-03T19:15:35.528170", "exception": false, "start_time": "2025-04-03T19:15:35.513325", "status": "completed"}, "tags": []}, "source": ["During training, we pass the input sequence through the Transformer encoder and predict the output for each input token.\n", "We use the standard Cross-Entropy loss to perform this.\n", "Every number is represented as a one-hot vector.\n", "Remember that representing the categories as single scalars decreases the expressiveness of the model extremely\n", "as $0$ and $1$ are not closer related than $0$ and $9$ in our example.\n", "An alternative to a one-hot vector is using a learned embedding vector as it is provided by the PyTorch module `nn.Embedding`.\n", "However, using a one-hot vector with an additional linear layer as in our case has the same effect\n", "as an embedding layer (`self.input_net` maps one-hot vector to a dense vector,\n", "where each row of the weight matrix represents the embedding for a specific category).\n", "\n", "To implement the training dynamic, we create a new class inheriting from `TransformerPredictor`\n", "and overwriting the training, validation and test step functions."]}, {"cell_type": "code", "execution_count": 18, "id": "160d5145", "metadata": {"execution": {"iopub.execute_input": "2025-04-03T19:15:35.559353Z", "iopub.status.busy": "2025-04-03T19:15:35.559015Z", "iopub.status.idle": "2025-04-03T19:15:35.568518Z", "shell.execute_reply": "2025-04-03T19:15:35.567670Z"}, "lines_to_next_cell": 2, "papermill": {"duration": 0.026622, "end_time": "2025-04-03T19:15:35.569722", "exception": false, "start_time": "2025-04-03T19:15:35.543100", "status": "completed"}, "tags": []}, "outputs": [], "source": ["class ReversePredictor(TransformerPredictor):\n", " def _calculate_loss(self, batch, mode=\"train\"):\n", " # Fetch data and transform categories to one-hot vectors\n", " inp_data, labels = batch\n", " inp_data = F.one_hot(inp_data, num_classes=self.hparams.num_classes).float()\n", "\n", " # Perform prediction and calculate loss and accuracy\n", " preds = self.forward(inp_data, add_positional_encoding=True)\n", " loss = F.cross_entropy(preds.view(-1, preds.size(-1)), labels.view(-1))\n", " acc = (preds.argmax(dim=-1) == labels).float().mean()\n", "\n", " # Logging\n", " self.log(f\"{mode}_loss\", loss)\n", " self.log(f\"{mode}_acc\", acc)\n", " return loss, acc\n", "\n", " def training_step(self, batch, batch_idx):\n", " loss, _ = self._calculate_loss(batch, mode=\"train\")\n", " return loss\n", "\n", " def validation_step(self, batch, batch_idx):\n", " _ = self._calculate_loss(batch, mode=\"val\")\n", "\n", " def test_step(self, batch, batch_idx):\n", " _ = self._calculate_loss(batch, mode=\"test\")"]}, {"cell_type": "markdown", "id": "7cb031f0", "metadata": {"lines_to_next_cell": 2, "papermill": {"duration": 0.014941, "end_time": "2025-04-03T19:15:35.599798", "exception": false, "start_time": "2025-04-03T19:15:35.584857", "status": "completed"}, "tags": []}, "source": ["Finally, we can create a training function similar to the one we have seen in Tutorial 5 for PyTorch Lightning.\n", "We create a `pl.Trainer` object, running for $N$ epochs, logging in TensorBoard, and saving our best model based on the validation.\n", "Afterward, we test our models on the test set.\n", "An additional parameter we pass to the trainer here is `gradient_clip_val`.\n", "This clips the norm of the gradients for all parameters before taking an optimizer step and prevents the model\n", "from diverging if we obtain very high gradients at, for instance, sharp loss surfaces (see many good blog posts\n", "on gradient clipping, like [DeepAI glossary](https://deepai.org/machine-learning-glossary-and-terms/gradient-clipping)).\n", "For Transformers, gradient clipping can help to further stabilize the training during the first few iterations, and also afterward.\n", "In plain PyTorch, you can apply gradient clipping via `torch.nn.utils.clip_grad_norm_(...)`\n", "(see [documentation](https://pytorch.org/docs/stable/generated/torch.nn.utils.clip_grad_norm_.html#torch.nn.utils.clip_grad_norm_)).\n", "The clip value is usually between 0.5 and 10, depending on how harsh you want to clip large gradients.\n", "After having explained this, let's implement the training function:"]}, {"cell_type": "code", "execution_count": 19, "id": "d65a25d3", "metadata": {"execution": {"iopub.execute_input": "2025-04-03T19:15:35.631137Z", "iopub.status.busy": "2025-04-03T19:15:35.630553Z", "iopub.status.idle": "2025-04-03T19:15:35.640173Z", "shell.execute_reply": "2025-04-03T19:15:35.639234Z"}, "papermill": {"duration": 0.026761, "end_time": "2025-04-03T19:15:35.641430", "exception": false, "start_time": "2025-04-03T19:15:35.614669", "status": "completed"}, "tags": []}, "outputs": [], "source": ["def train_reverse(**kwargs):\n", " # Create a PyTorch Lightning trainer with the generation callback\n", " root_dir = os.path.join(CHECKPOINT_PATH, \"ReverseTask\")\n", " os.makedirs(root_dir, exist_ok=True)\n", " trainer = pl.Trainer(\n", " default_root_dir=root_dir,\n", " callbacks=[ModelCheckpoint(save_weights_only=True, mode=\"max\", monitor=\"val_acc\")],\n", " accelerator=\"auto\",\n", " devices=1,\n", " max_epochs=10,\n", " gradient_clip_val=5,\n", " )\n", " trainer.logger._default_hp_metric = None # Optional logging argument that we don't need\n", "\n", " # Check whether pretrained model exists. If yes, load it and skip training\n", " pretrained_filename = os.path.join(CHECKPOINT_PATH, \"ReverseTask.ckpt\")\n", " if os.path.isfile(pretrained_filename):\n", " print(\"Found pretrained model, loading...\")\n", " model = ReversePredictor.load_from_checkpoint(pretrained_filename)\n", " else:\n", " model = ReversePredictor(max_iters=trainer.max_epochs * len(train_loader), **kwargs)\n", " trainer.fit(model, train_loader, val_loader)\n", "\n", " # Test best model on validation and test set\n", " val_result = trainer.test(model, dataloaders=val_loader, verbose=False)\n", " test_result = trainer.test(model, dataloaders=test_loader, verbose=False)\n", " result = {\"test_acc\": test_result[0][\"test_acc\"], \"val_acc\": val_result[0][\"test_acc\"]}\n", "\n", " model = model.to(device)\n", " return model, result"]}, {"cell_type": "markdown", "id": "fed2c30b", "metadata": {"papermill": {"duration": 0.014945, "end_time": "2025-04-03T19:15:35.671426", "exception": false, "start_time": "2025-04-03T19:15:35.656481", "status": "completed"}, "tags": []}, "source": ["Finally, we can train the model.\n", "In this setup, we will use a single encoder block and a single head in the Multi-Head Attention.\n", "This is chosen because of the simplicity of the task, and in this case, the attention can actually be interpreted\n", "as an \"explanation\" of the predictions (compared to the other papers above dealing with deep Transformers)."]}, {"cell_type": "code", "execution_count": 20, "id": "b56ec207", "metadata": {"execution": {"iopub.execute_input": "2025-04-03T19:15:35.703259Z", "iopub.status.busy": "2025-04-03T19:15:35.702460Z", "iopub.status.idle": "2025-04-03T19:15:36.945433Z", "shell.execute_reply": "2025-04-03T19:15:36.944718Z"}, "papermill": {"duration": 1.260723, "end_time": "2025-04-03T19:15:36.947230", "exception": false, "start_time": "2025-04-03T19:15:35.686507", "status": "completed"}, "tags": []}, "outputs": [{"name": "stderr", "output_type": "stream", "text": ["GPU available: True (cuda), used: True\n"]}, {"name": "stderr", "output_type": "stream", "text": ["TPU available: False, using: 0 TPU cores\n"]}, {"name": "stderr", "output_type": "stream", "text": ["HPU available: False, using: 0 HPUs\n"]}, {"name": "stderr", "output_type": "stream", "text": ["/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/connectors/logger_connector/logger_connector.py:75: Starting from v1.9.0, `tensorboardX` has been removed as a dependency of the `pytorch_lightning` package, due to potential conflicts with other packages in the ML ecosystem. For this reason, `logger=True` will use `CSVLogger` as the default logger, unless the `tensorboard` or `tensorboardX` packages are found. Please `pip install lightning[extra]` or one of them to enable TensorBoard support by default\n"]}, {"name": "stdout", "output_type": "stream", "text": ["Found pretrained model, loading...\n"]}, {"name": "stderr", "output_type": "stream", "text": ["Lightning automatically upgraded your loaded checkpoint from v1.0.2 to v2.4.0. To apply the upgrade to your files permanently, run `python -m pytorch_lightning.utilities.upgrade_checkpoint saved_models/Transformers/ReverseTask.ckpt`\n"]}, {"name": "stderr", "output_type": "stream", "text": ["You are using a CUDA device ('NVIDIA GeForce RTX 3090') that has Tensor Cores. To properly utilize them, you should set `torch.set_float32_matmul_precision('medium' | 'high')` which will trade-off precision for performance. For more details, read https://pytorch.org/docs/stable/generated/torch.set_float32_matmul_precision.html#torch.set_float32_matmul_precision\n"]}, {"name": "stderr", "output_type": "stream", "text": ["LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [2,3]\n"]}, {"name": "stderr", "output_type": "stream", "text": ["/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/connectors/data_connector.py:424: The 'test_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=63` in the `DataLoader` to improve performance.\n"]}, {"data": {"application/vnd.jupyter.widget-view+json": {"model_id": "3716048b615a4b309aa31f7db611a25f", "version_major": 2, "version_minor": 0}, "text/plain": ["Testing: | | 0/? [00:00\n", "\n", "\n"], "text/plain": ["
![](\"https://github.com/Lightning-AI/lightning-tutorials/raw/main/course_UvA-DL/05-transformers-and-MH-attention/cifar100_example_anomaly.png\")