Book Review: Transformers In Action

The Attention Is All You Need paper proposed the Transformer Architecrture as an improvement to the dominant encoder-decoder models of the time (both recurrent and convolutional). These models used an attention mechanism to connect the encoder and decoder parts, but the Transformer Architecture flipped the script, putting the Attention Mechanism at the center. An early implementation of the Transformer Architecture was BERT, which used the Transformer as an encoder. Later models such as BART used encoder and decoder Transformer components in a sequence-to-sequence setup. Since then, there has been an explosion of variants around this basic model, accompanied by a steady breaking of benchmarks at tasks where older recurrent and convolution sequence-to-sequence models reigned supreme.

A second major breakthrough was the emergence of decoder-only Transformer models for text generation. Early models were less than encouraging, but as researchers increased the number of parameters to train ever larger and larger models on larger and larger datasets, their text generation capabilities improved to the point where they become viable candidates for using as pre-trained general purpose inference models. These models are also based on Transformers, but are generally differentiated by calling them Large Language Models (LLM) or Foundation Models (FM).

From a user's point of view, once you get past the slightly larger computing requirements, the first category (BERT like Transformer models) is actually easier to fine tune for custom tasks than its predecessors, thanks to tooling available from libraries such as HuggingFace Transformers and SentenceTransformers. The second category (LLMs), at least initially, were the domain of compute and data rich organizations, who would create these models and make them available to others over an HTTP API as inference-only models, often for a fee. Because of the massive number of parameters and volume of training data, these models were generalized enough to do inference on diverse tasks in diverse domains without additional fine-tuning. Of course, because they were generative models, their outputs were not deterministic, prompting cautions such as the On the Dangers of Stochastic Parrots paper, and patterns to alleviate it like Retrieval Augmented Generation (RAG) and Chain of Thought (CoT) prompting. More recently, fine tuning has become practical for this class of models with the advent of Parameter Efficient Fine Tuning (PEFT) techniques. Also, with the advent of multimodal LLMs and reasoning capabilities, they are more than just Large Language Models.

Anyway, the point of this (probably incomplete) history lesson is that the Transformers in Action book by Nicole Koenigstein, that I am reviewing, primarily covers Transformers in the second category, except for the first two chapters where it covers basics of the Transformer architecture. If you were more interested in the first category, I would recommend Transformers for Natural Language Processing by Denis Rothman, which I have reviewed on Amazon previously.

Back to the review. The book is organized in three parts, with the first part consisting of Chapters 1 and 2, the second part consisting of Chapters 3-5 and the third part consisting of Chapters 6-10.

In Part 1, Chapter 1 describes the Transformers architecture at a high level, how it incorporates ideas from earlier neural models and how it is different from them. It covers the idea of in-context learning (zero-shot and few shot), the distinguishing feature of Transformer based LLMs. Chapter 2 does a deep dive into the Transformer Architecture and its components, covering ideas such as Stacked Encoder-Decoder, Add and Norm (LayerNorm) layers, the Query-Key-Value Attention Mechanism, and position wise Feed Forward Network (FFN).

In Part 2, Chapter 3 moves the discussion into decoder-only Transformers, i.e. Large Language Models, and the central theme of this book. It describes variants of the Transformer Architecture, i.e. encoder only models such as BERT versus decoder only Autoregressive models that predict the next token. It touches on Causal Attention and KV Cache as necessary ingredients for this type of model. It also touches on the use of encoder only models as Embedding models and how it relates to RAG. It also mentions Mixture of Experts (MoE) as a promising architectural variant of Decoder-only models.

Chapter 4 covers some basics about parameters that control the behavior of LLMs, such as top-k and top-p sampling, temperature, prompting styles such as Zero-shot, Few shot, CoT (Chain of Thought), Contrastive CoT where both right and wrong reasoning traces are provided, Chain of Verification (CoVe) where the model reflects on and verifies its output, Tree of Thought (ToT) which introduces intermediate steps in problem solving traces, and Thread of Thought (ThoT) that partitions the problem into sub-problems and combines the threads from the sub-solutions into the final generation.

Chapter 5 covers Preference Alignment and RAG. The first part, Preference Alignment, is aimed at fine-tuning the behavior of an LLM to a particular domain or behavior. It covers Reinforcement Learning from Human Feedback (RLHF) as a Markov Decision Process (MDP) and the use of Proximal Policy Optimization (PPO). It describes specializations of PPO such as DPO (Direct Preference Optimization) that does not need an explicit reward model, and GRPO ((Group Relative Policy Optimization) tht removes the need for an explicit value function. Both DPO and GRPO are preceded by SFT (Supervised Fine Tuning) to align a transformer to a domain. The second part covers RAG, which is more familiar to most people using LLMs -- the discussion formalizes the structure of a RAG pipeline (retriever, generator and refinement layer) and describes some popular RAG variants, i.e. Agentic RAG, Corrective RAG, Self RAG and Fusion RAG.

In Part 3, Chapter 6 discusses Multimodal models, how they are different from text-only LLMs, and how they work by projecting text and non-text input into a shared embedding space. It differentiates between Converter based alignment where all modalities are projected onto the same space and Perception based alignment where modality specific encoders produce each embedding and the LLM uses an Attention mechanism to combine them.

Chapter 7 discusses Small Language Models (SLMs) which are decoder-only transformer models with 8-13 billion parameters. These are larger than encoder-only or encoder-decoder style models, but smaller than other decoder-only models. Such models are usually better at general purpose inference than (smaller) encoder-only or encoder-decoder models, but not as good as their full size counterparts. SLMs focus on specialization and efficiency, and can be deployed as edge devices or specialized components co-existing with LLMs in RAG pipelines. They can be used to generate data to train their larger counterparts using Weak to Strong Learning, Approximate Gradient Proxies, and function as Auxilliary Reward Models for RLHF. They can be deployed as specialized tools or agents in Agentic Workflows, e.g. classifiers for sentiment analysis, compliance checks, intent detection, guard models and coding models. They also work well in privacy concious domains, where you don't want your requests going out to a third party model provider. Finally, they are more practical to fine-tune for your specific use case than the full sized LLM.

Chapter 8 discusses training and evaluating Large Language Models, and suggests the use of Ray Tune for Hyperparameter Tuning and the Weights and Biases (W&B) platform for logging and determining GPU utilization. It details various PEFT techniques such as LoRA (Low Rank Adaptation), DoRA (Weight Decomposed LoRA), Quantization, QLoRA (Quantized LoRA), QA-LoRA (Quantized Aware LoRA), and LQ-LoRA (Low Rank plus Quantized Matrix Decomposition LoRA). Unfortunately, the author has not included too many examples of this in the book, possibly based on it being perceived as out of scope for this book's average reader. However, I had expected some coverage of evaluation techniques which I did not find -- evaluation is a real problem for teams building RAG or other inference-only pipelines, and it is complex because outputs are non-deterministic. Perhaps this is an oversight that can be addressed in a future edition of the book.

Chapter 9 covers deployment issues associated with LLMs, namely around optimization and scaling. Model Optimization techniques such as pruning (removing neurons or edges) and distillation (from larger to more efficient smaller models for specific tasks) and Memory Optimization techniques such as various types of sharding (tensor, pipeline, optimizer and hybrid) are described. The chapter also describes Inference Optimization techniques such as KV Caching, Paged Attention, vLLM and Operator Fusion, GPU Optimizations such as Tiling and Flash Attention, and extensions to support Long Context such as Rotary Embeddings (RoPE), iRoPE which alternates between RoPE and NoPE (No Positional Embeddings), block sparse and linear attention, and sliding window and chunked attention.

Finally, Chapter 10 covers techniques to create Responsible and Ethical LLM based applications. It outlines some possible reasons for LLM bias based on the geographical distribution of training data, approaches to flag and filter hateful or toxic generations using pre-trained BERT class models such as RoBERTa-Toxicity and HateBERT, the use of custom logging on W&B for interpretability analysis, using perturbation models in the Captum tool to determine feature attribution, and explanations using Local Interpretable Model Agnostic Explanations (LIME). It also describes some rule-based techniques to ensure Responsible behavior of LLMs, such as adding disclaimer text, penalizing tokens if they match a blacklist, using rule and LLM driven input and output guards like the ones provided by llm-guard. It also suggests using safe/unsafe classifiers in Purple Llama to address Lifecylce vulnerabilities and prevent Jailbreaks.

As with the earlier Transformers book, what I found most useful about this book was the coverage. While I feel fortunate to have actually lived through these transformative times rather than read about them in a book, the pace of breakthroughs in the state of the art are hard to keep up with unless you are actively doing the research yourself (and maybe not even then). As a result, you end up knowing about a few things that you have used or considered using or found interesting, but are woefully ignorant about a lot of the other things in the field. Books like this not only fill out your knowledge gaps, they also give you new ideas based on things that you just learned.

In addition, this book describes many useful techniques to improve your LLM pipelines. Many of us, me included, have built traditional and neural (pre-transformer and transformer based) ML pipelines, and have been building RAG pipelines over the past couple of years. But we may not be familiar with all the latest prompting techniques, or we may not have fine-tuned a SLM because of the compute requirements. Books like these show us how to do it, and thereby make us more productive and more effective users of LLMs.

Knowledge Graph Aligned Entity Linker using SentenceTransformers

Most of us are familiar with Named Entity Recognizers (NERs) that can recognize spans in text as belonging to a small number of classes, such as Person (PER), Organization (ORG), Location (LOC), etc. These are usually multi-class classifier models, trained on input sequences to return BIO (Begin-Input-Output) tags for each token. However, recognizing entities in a Knowledge Graph (KG) using this approach is usually a much harder proposition, since a KG can contain thousands, even millions, of distinct entities, and it is just not practical to create a multi-class classifier for so many target classes. A common approach to building a NER for such a large number of entities is to use dictionary based matching. However, the approach suffers from the inability to do "fuzzy" or inexact matching, beyond standard normalization streategies such as lowercasing and stemming / lemmatizing, and requires you to specify up-front all possible synonyms that may be used to refer to a given entity.

An alternative approach may be to train another model, called a Named Entity Linker (NEL) that would take the spans recognized as candidate entities or phrases by the NER model, and then attempt to link the phrase to an entity in the KG. In this situation, the NER just learns to predict candidate phrases that may be entities of interest, which puts it on par with simpler PER/ORG/LOC style NERs in terms of complexity. The NER and NEL are pipelined together in a setup that is usually known as Named Entity Recognition and Linking (NERL).

In this post, I will describe a NEL model that I built for my 2023 Dev10 project. Our Dev10 program allows employees to use up to 10 working days per year to pursue a side-project, similar to Google's 20% program. The objective is to learn an embedding model where encodings of synonyms of a given entity are close together, and where encodings of synonyms of different entities are pushed far apart. We can then encode each entity in this space as the encoding of the centroid of the encodings of its individual synonyms. Each candidate phrase output from the NER model can then be encoded using this embedding model, and its nearest neighbors in the embedding space would correspond to the most likely entities to link to.

The idea is inspired by Self-Alignment Pretraining for Biomedical Entity Representations (Liu et al, 2021) which produced the SapBERT model (SAP == Self Aligned Pretraining). It uses Contrastive Learning to fine-tune the BiomedBERT model. In this scenario, positive pairs are pairs of synonyms for the same entity in the KG and negative pairs are synonyms from different entities. It uses the Unified Medical Language System (UMLS) as its KG, to source synonym pairs.

I follow a largely similar approach in my project, except that I use the SentenceTransformers library to fine tune the BiomedBERT model. For my initial experiments, I also used the UMLS as my source of synonym pairs, mainly for reproducibility purposes since it is a free resource available for download to anyone. I tried fine-tuning a bert-base-uncased model and the BiomedBERT models, with MultipleNegativesRanking (MNR) as well as Triplet loss, the latter with Hard Negative Mining. My findings are in line with the SapBERT paper, i.e. that BiomedBERT performs better than BERT base, and that MNR performs better than Triplet loss. The last bit was something of a dissapointment, since I had expected Triplet loss to perform better. It is possible that the Hard Negative Mining was not hard enough, or maybe I needed a higher number than 5 negatives for each positive.

You can learn more about the project in my GitHub repository sujitpal/kg-aligned-entity-linker, as well as find the code in there, in case you want to replicate it.

Here are some visualizations from my best model. The chart on the left shows the distribution of cosine similarities between known negative synonym pairs (orange curve) and known positive synonym pairs (blue curve). As you can see, there is almost no overlap. The heatmap on the right shows the cosine similarity of a set of 10 synonym pairs, where the diagonal corresponds to positive pairs and the non-diagonal elements correspond to negative pairs. As you can see, the distribution seems quite good.

I also built a small demo that shows what in my opinion is the main use case for this model. It is a NERL pipeline, where the NER component is the UMLS entity finder (en_core_sci_sm) from the SciSpacy project, and the NEL component is my best performing model (kgnel-bmbert-mnr). In order to look up nearest neighbors for a given phrase encoding, the NEL component also needs a vector store to store the centroids of the encodings of entity synonyms, I used QDrant for this purpose. The QDrant vector store needs to be populated with the centroid embeddings in advance, and in order to cut down on the index and vectorization time, I only computed embeddings for centroids for entities of type "Disease or Syndrome" and "Clinical Drug". The visualizations below show the outputs (from displacy) of the outputs of the NER component:

and that of the NEL component in my demo NERL pipeline. Note that only spans that were identified as a Disease or Drug with a confidence above a threshold were selected in this phase.

Such a NERL pipeline could be used to mine new literature for new synonyms of existing entities. Once discovered, they could be added to the synonym list for the dictionary based NER to increase its recall.

Anyway, that was all I had for this post. Today is also January 1 2024, so I wanted to wish you all a very Happy New Year and a productive 2024 filled with many Machine Learning adventures!

Fine-tuning OpenAI CLIP for different domains

In July this year, a group of us on the TWIML Slack Channel came together and participated in the Flax/JAX Community Week organized by Hugging Face and Google Cloud. Our project was about fine-tuning the CLIP Model from OpenAI with the RSICD (Remote Sensing Image Captioning Dataset), and ended up placing third.

The code for the project is available on github at arampacha/CLIP-rsicd if you are curious about how we went about doing this, or if you want to replicate our efforts. Our fine-tuned model is available on the Hugging Face model repository at flax-community/clip-rsicd-v2, you can find instructions on how to use it for inference on your own remote-sensing / satellite data. We also have a Streamlit based demo that shows its application in image search and finding features in images using text descriptions. Finally, we also have a blog post on the Hugging Face blog titled Fine tuning CLIP with Remote Sensing (Satellite) images and captions. Hope you fine these useful, do check them out.

Even before this project, I had been considering learning a joint embedding for medical images and their captions as described in the Contrastive Learning of Medical Visual Representations from Paired Images and Text (CONVIRT) paper by Zhang et al (2010), and using it to power a text-to-image image search application. Based on the RSICD project, however, CLIP looked like a better and more modern alternative.

Elsevier has a Dev-10 program for their engineers, by which they are given 10 working days (2 weeks) to build something that does not necessarily have to align with company objectives, but which is somewhat work-related. When my Dev-10 days came up in early September, I used it to fine-tune the same OpenAI CLIP baseline as we did for the Flax/JAX community week, but with the ImageCLEF 2017 Image Captioning dataset. Happily, the results were just as encouraging as fine-tuning it with RSICD, if anything, the improvement was even more dtamatic.

During the RSICD fine-tuning exercise, the fine-tuning work was done by other members of the team. My contribution to that project was the evaluation framework, the image augmentation piece, the demo, and later the blog post. On the ImageCLEF exercise, I was the only developer, so while a lot of the code in the second case was borrowed or adapted from the first, there were some important differences as well, apart from the dataset.

First, in the RSICD fine-tuning case we used JAX/Flax with a TPU enabled instance on Google Cloud, and in the second I used Pytorch on a single-GPU EC2 instance on AWS (with the Deep Learning AMI). I found that the Hugging Face wrapper for CLIP provides a lot of the support that was being done explicitly, so I tried to leverage the provided functionality as much as possible, resulting in slightly cleaner and more readable code (even if I say so myself :-)).

Second, I didn't do any image or text augmentation like we did with the RSICD fine-tuning effort. RSICD had a total of 10k images with approximately 5 captions per image, of which we were using about 7k for training. On the other hand, ImageCLEF was about 160k images and captions, of which we were using 140k for training. In addition, RSICD was training on a TPU with 4 parallel devices, and ImageCLEF was training an a single GPU. Because of this, I ended up using subsampling from the training set as a form of regularization instead, and using early stopping to terminate the training process once no improvements in validation accuracy were detected.

Third, with the benefit of hindsight, I settled on a more industry-standard metric for evaluation, the Mean Reciprocal Rank (MRR@k) compared to the less strict and somewhat ad-hoc Hits@k metric I had used for the first exercise.

And fourth, because the data volume for my second Image Search demo was much larger (200k images instad of 10k), I switched from using NMSLib to using Vespa, the open source hybrid vector + text search engine from Yahoo!. Using it, I was able to provide image search results based on lexical matches between query and caption text, vector space matches between CLIP query vector and CLIP image vectors, and hybrid search results ranked by combining the relevance of the two approaches.

Unfortunately I am not able to share the code. Since the work was done on company time with company resources, the code rightfully belongs to the company. I am also hopeful that the work could be used to power image search (or related) functionlity in some production application. For these reasons I am unable to share the code, but in general, it is similar (with the differences enumerated above) to the RSICD version.

However, just to give some idea of the kind of results you can expect from a fine-tuned CLIP model, here are couple of screenshots. The results are for the queries "computed tomography" and "computed tomography deep vein thrombosis". Both results are from doing vector matching, i.e. ranked by cosine similarity between the CLIP encoding of the query text and the CLIP encoding of each image.

As you can see, CLIP returns relevant images for both high level and detailed queries, indicating how rich the embedding is. My main takeaway from this series of exercises are twofold -- first, CLIP's joint image-text encoding is a seriously powerful idea and is super-effective, and second, transformer models trained on general data (natural images and text in this case) can be fine-tuned effectively for specialized domains using relatively small amounts of data.

Distributed Training of a Bengali ALBERT model

Even though I am from India and my mother tongue is Bengali, and I speak, read, and write both Hindi and Bengali almost as well as English, in my career with Natural Language Processing (NLP) I have worked exclusively with English. This is probably not that uncommon, because until recently, English was the language where most NLP work happened, and to a lesser extent some of the major European languages (Spanish, French, German, Russian, etc.). Fortunately or unfortunately, among these languages, English was the only one I knew well enough to work with.

As NLP work with European languages became more widespread, I secretly envied my European colleagues for being multilingual in the "right" languages. The rise of CJK (Chinese, Japanese, Korean) that followed (and its impact on NLP in CJK languages) largely passed me by as well, since I did not know any of these languages either. Lately, however, I have been encouraged by the rise of NLP with Indic languages (languages spoken in India), not the least because it has given me hope that I will finally be able to put my multilingual skills to some use after all :-).

Indic languages have largely been considered low-resource languages, because there was not enough material in electronic format to train NLP models, in spite of most of them individually having a fairly rich and evolved literature. This has changed (or least been alleviated to a large extent) with the rise of the Internet and social media, and Indian people rediscovering their roots and beginning to communicate in their native languages. Software infrastructure to support this, such as Avro keyboard has also helped, making it easier to start communicating electronically using non-English languages.

In any case, I saw this tweet inviting people that spoke Bengali to a decentralized training experiment organized by Neuropark, Hugging Face, and Yandex Research to train an ALBERT model for Bengali. Participants needed access to Colab and an Internet connection. I was curious about the distributed training part, and since I satisfied the prerequisites, I decided to join in the experiment. That was a week and a half ago, training finished today (Friday). In this post, I will describe what I learned from the experience.

The objective was to train an ALBERT-large model from scratch on the Bengali language. The ALBERT transformer model was proposed in the paper ALBERT: A lite BERT for Self-Supervised Learning of Language Representations in 2019 by Lan et al. It is based on the BERT transformer model, but has fewer parameters and better performance on many benchmark tasks. The steps involved in the training are as follows.

1. Bengali tokenizer training.
2. ALBERT Bengali Language Model (LM) training.
3. Model evaluation, both subjective and using downstream task

Tokenizer Training

The tokenizer was trained on the the Bengali subset of the multilingual OSCAR dataset. Text was normalized using the following normalizer pipeline: NMT, which converts various whitespace breaks between words to a simple space; NFKC, which does some unicode magic (see below) that unifies the way characters are encoded; lowercase, which doesn't affect Bengali as much because it doesn't have case, but does help with embedded English text, and various regexes, including one to transform a sequence of spaces to a single space. The Unigram Language Model algorithm (see Subword Regularization: Improving Neural Network Translation Models with Multiple Subword Candidates (Kudo, 2018)) wqs used for tokenization.

The open source Bengali NLP library BNLP was used for sentence segmentation in the model training step (see below). The team also tried out BLTK, another Bengali NLP library, but finally went with BNLP after testing results from both.

A previous version of the tokenizer was trained using data scraped from various Bengali language websites via the Bakya project and used Byte Pair Encoding (BPE), but this was not used in the final training. In my original post, I had mistakenly assumed that this was the tokenizer that was being used for the training.

The work around normalization happened before I joined the project, but I was around when there was a request to check the quality of sentences tokenized using BNLP versus BLTK. It was then that I realized that the team actually needed Bengali readers rather than speakers, and (mistakenly at least in my case) assumed that the latter automatically implies the former. Having grown up outside Bengal, I learned Hindi at school as a second language, so while I can read Bengali (having learnt it at home), I am not that fluent in it as I am at Hindi.

I also learned another interesting thing about Unicode character representation for Bengali (and probably other Indic languages), which is probably related to the Unicode magic around NFKC, that I want to share here. In English, the 26 letters of the alphabet are combined in different ways to form words. In the Bengali alphabet (as in Hindi and possibly other Indic languages derived from Sanskrit), there are 7 consonant groups of 5 characters each. Each group emits a sound that uses a particular section of your vocal apparatus (lips, tongue and roof of palate, throat, etc), and the sound gets softer as you step across the group. There are also 14 vowel characters that are used to modify the consonant sounds to form words. Unlike English, the vowels are overlaid on the consonants at the same character position. In addition, pairs of consonants can be conjoined to form new characters representing a transitional sound -- this is called যুক্তাক্ষর (pronounced juktakkhor) or conjoined word.

Anyway, it turns out that Unicode elegantly handles both the overlaying of vowels on to consonants as well as combining two consonants to form a third, as the following code snippet illustrates (probably more readily apparent to Bengali readers, others will need to squint a bit at the output to get it).

Model Training

The model was trained on text from Bengali Wikipedia and the Bengali portion of the OSACAR dataset combined. The model being trained was the AlbertForPreTraining model from Hugging Face. ALBERT uses two pre-training objectives. The first is Masked Language Modeling (MLM) similar to BERT, where we mask out 15% of the tokens and have the model learn to predict them. The second is Sentence Order Prediction (SOP) which in case of BERT tries to predict if one sentence follows another, but in case of ALBERT uses text segments instead of sentences, and is regarded as more efficient compared to BERT SOP.

Training was done in a distributed manner using the Hivemind project from Yandex Research. This project allows a central team to build the training script and have volunteer members on the Internet (such as myself) run it on a subset of the data, using free GPU-enabled Colab and Kaggle notebooks. I believe Hivemind can also distribute the training across hybrid non-cloud GPU instances and non-free cloud instances as well, but these were not used here. Once started, the training script on a particular Colab or Kaggle notebook will continue until the user stops it or the platform decides to time them out, either via policy (Kaggle allows maximum 9 hours continuous GPU use) or due to inactivity. The training scripts can be found in the github repository mryab/collaborative-training.

Volunteers need to opt-in to the training by adding themselves to an allow-list (requesting via the Discord channel) and signing up for a Hugging Face account. When starting up their instance, they authenticate themselves via their Hugging Face username and password. Each notebook functions as a peer in the decentralized training setup, training the model locally and creating local updates against the model, and logging its progress using the Weights and Biases (wandb) API. At the end of each training step, notebooks within the peer group share model parameters (model averaging) with each other using a process called butterfly all-reduce. After each successful training round, the peers shuffle around and find new groups to join. This ensures that the local updates are propagated to all the peers over time. If a peer leaves the group, this affects only the immediate peer group, the remaining members of which will be re-assembled into other running peer groups.

For a more technical coverage of the distributed training algorithm, please refer to Moshpit SGD: Communication-Efficient Decentralized Training on Heterogeneous Unreliable Devices (Ryabinin et al, 2021) and its predecessor Towards Crowdsourced Training of Large Neural Networks using decentralized Mixture-of-Experts (Ryabinin and Gusev, 2020).

At the point when training started, the model was reporting a loss of around 11, which came down to below 2 after one week and over 20,000 training steps, as shown in the loss curve on the left below. The alive peers on the right shows the number of simultaneous training instances over the week. At its peak there were around 50, which oscillated between 20 and 40 over the course of the training. The gradual decline towards the end of the training could be at least partially attributed to volunteers running out of Kaggle quotas (30 GPU hours per week) and being punished by Colab for hogging CPU resources.

Model Evaluation

Of course, for a language model such as Bengali ALBERT, a better metric than the loss decreasing from 11 to 1.97, is how well it does on some downstream task. As the model trained, its checkpoints were subjected to two forms of evaluation.

First, the model was fine-tuned for an NER task (WikiNER) using the Bengali subset of the multi-lingual Wiki-ANN dataset, a dataset annotated with LOC (location), PER (person), and ORG (organization) tags in IOB format. The charts below the Precision, Recall, and F1 values by model checkpoints over the course of the training. The final scores were 97.5% accuracy, 95.6% F1, 95.4% Precision, and 95.8% Recall.

In addition, model checkpoints were used to test the model's capability to predict masked words in provided sentences. This evaluation was more subjective in nature, manually looking at the top 5 masked word predictions for given sentences and checking out their relevance, but it was observed that the final model made almost perfect masked word predictions, compared to previous checkpoints with more variable behavior.

Conclusion

This experience has been of immense educational value for me. I got to use and see a distributed training environment close up, and got to interact with a lot of very smart and committed developers and researchers and fellow volunteers who I will not list by name, because I am sure I will forget someone. I also got to see a lot of code that I am sure I will use for inspiration later. For example, I am also a bit embarrassed to say that this was my first experience using the Weights and Biases (wandb) API, but I liked what I saw, so I plan to use it in the future.

In addition, the progress that has been made in Bengali NLP (and other Indic languages) was a real eye opener for me. In fact, the current model is not even the first transformer based model for Bengali, there is already a multi-language IndicBERT which has shown promising results on some tasks. However, this is the first transformer based model for Bengali that was trained in a distributed manner.

The model (tentatively called SahajBERT) and tokenizer will shortly be available for download on Hugging Face. I will provide the links to them as they become available.

Finally, many thanks to Nilavya Das, Max Ryabinin, Tanmoy Sarkar, and Lucile Saulnier for their valuable comments and for fact-checking the draft version of this post.

Updates (2021-05-24)

1. Updated description of tokenizer training process.
2. Added links to papers that provide more information about the distributed training approach.

Update (2021-06-01) -- The trained tokenizer and model described above has been published and is now available for download at neuropark/sahajBERT on the Huggingface models site.

Word Sense Disambiguation using BERT as a Language Model

The BERT (Bidirectional Encoder Representation from Transformers) model was proposed in the paper BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding (Devlin, et al, 2019). BERT is the encoder part of an encoder-decoder architecture called Transformers, that was proposed in Attention is all you need (Vaswani, et al., 2017). The BERT model is pre-trained on two tasks against a large corpus of text in a self-supervised manner -- first, to predict masked words in a sentence, and second, to predict a sentence given the previous one, and are called Masked Language Modeling and Next Sentence Prediction tasks respectively. These pre-trained models can be further fine-tuned for tasks as diverse as classification, sequence prediction, and question answering.

Since the release of BERT, the research community has done a lot of work around Transformers and BERT-like architectures, so much so, that HuggingFace has its enormously popular transformers library dedicated to helping people work efficiently and easily with popular Transformer architectures. Among other things, the HuggingFace transformers library provides a unified interface to working with different kinds of Transformer architectures (with slightly different details), as well as provide weights for many pre-trained Transformer architectures.

Most of my work with transformers so far has been around fine-tuning them for Question Answering and Sequence Prediction. I recently came across a blog post Examining BERT's raw embeddings by Ajit Rajasekharan, where he describes how one can use a plain BERT model (pre-trained only, no fine-tuning required) and later a BERT Masked Language Model (MLM), as a Language Model, to help with Word Sense Disambiguation (WSD).

The idea is rooted in the model's ability to produce contextual embeddings for a words in a sentence. A pre-trained model has learned enough about the language it is trained on, to produce different embeddings for a homonym based on different sentence contexts it appears in. For example, a pre-trained model would produce a different vector representation for the word "bank" if it is used in the context of a bank robbery versus a river bank. This is different from how the older word embeddings such as word2vec work, in that case a word has a single embedding, regardless of the sentence context in which it appears.

An important point here is that there is no fine-tuning, we will leverage the knowledge inherent in the pre-trained models for our WSD experiments, and use these models in inference mode.

In this post, I will summarize these ideas from Ajit Rajasekharan's blog post, and provide Jupyter notebooks with implementations of these ideas using the HuggingFace transformers library.

WSD using raw BERT embeddings

Our first experiment uses a pre-trained BERT model initialized with the weights of a bert-base-cased model. We extract a matrix of "base" embeddings for each word in the model's vocabulary. We then pass in sentences containing our ambiguous word into the pre-trained BERT model, and capture the input embedding and output embedding for our ambiguous word. Our first sentence uses the word "bank" in the context of banking, and our second sentence uses it in the context of a river bank.

We then compute the cosine similarity between the embedding (input and output) for our ambiguous word against all the words in the vocabulary, and plot the histogram of cosine similarities. We notice that in both cases, the histogram shows a long tail, but the histogram for the output embedding seems to have a shorter tail, perhaps because there is less uncertainty once the context is known.

We then identify the words in the vocabulary whose embeddings are most similar (cosine similarity) to the embedding for our ambiguous word. As expected, the similar words for both input embeddings relate to banking (presumably because this may be the dominant usage of the word in the language). For the output embeddings, also as expected, similar words for our ambiguous word relate to banking in the first sentence, and rivers in the second.

The notebook WSD Using BERT Raw Embeddings contains the implementation described above.

WSD using BERT MLM

In our second experiment, we mask out the word "bank" in our two sentences and replace it with the [MASK] token. We then pass these sentences through a BERT Masked Language Model (MLM) initialized with weights from a bert-base-cased model. The output of the MLM is a 3-dimensional tensor of logits, where the first dimension is the number of sentences in the batch (1), the second dimension is the number of tokens in the input sentence, and the third domension is the number of words in the vocabulary. Effectively, the output provides log probabilities for predictions across the entire vocabulary for each token position in the input.

As before, we identify the logits corresponding to our masked position in the input (and output) sequence, then compute the softmax of the logits to convert them to probabilities. We then extract the top k (k=20) terms with the highest probabilities.

Again, as expected, predictions for the masked word are predominantly around banking for the first sentence, and predominantly around rivers for the second sentence.

The notebook WSD Using BERT Masked Language Model contains the implementation described above.

So thats all I had for today. Even though I understood the idea in Ajit Rajasekharan's blog post at a high level, and had even attempted something similar for WSD using non-contextual word embeddings (using the average of word embeddings across a span of text around the ambiguous word), it was interesting to actually go into the transformer model and figure out how to make things work. I hope you found it interesting as well.