Skip to content

Contact sales

By filling out this form and clicking submit, you acknowledge our privacy policy.

Creating a Large Language Model from scratch: A beginner's guide

A step-by-step guide on how to create your first Large Language Model (LLM), even if you're new to natural language processing.

Feb 15, 2024 • 10 Minute Read

Please set an alt value for this image...

Imagine stepping into the world of language models as a painter stepping in front of a blank canvas. The canvas here is the vast potential of Natural Language Processing (NLP), and your paintbrush is the understanding of Large Language Models (LLMs). This article aims to guide you, a data practitioner new to NLP, in creating your first Large Language Model from scratch, focusing on the Transformer architecture and utilizing TensorFlow and Keras.

Table of contents

Understanding the basics

What is a Large Language Model?

A Large Language Model (LLM) is akin to a highly skilled linguist, capable of understanding, interpreting, and generating human language. In the world of artificial intelligence, it's a complex model trained on vast amounts of text data.

It is a type of artificial intelligence model specifically designed to understand, interpret, generate, and sometimes translate human language. These models are a subset of machine learning models and are part of the broader field of natural language processing (NLP). Let's break down the concept to understand it better:

Key Characteristics of Large Language Models:

  1. Large Scale: As the name suggests, these models are 'large' not just in their physical size in terms of the number of parameters they contain, but also in the vast amount of data they are trained on. Models like GPT-3, BERT, and T5 consist of billions of parameters and are trained on diverse datasets comprising texts from books, websites, and other sources.

  2. Understanding Context: One of the primary strengths of LLMs is their ability to understand the context. Unlike earlier models that focused on individual words or phrases in isolation, LLMs consider the entire sentence or paragraph, allowing them to comprehend nuances, ambiguities, and the flow of language.

  3. Generating Human-Like Text: LLMs are known for their ability to generate text that closely resembles human writing. This includes completing sentences, writing essays, creating poetry, or even generating code. The advanced models can maintain a theme or style over long passages.

  4. Adaptability: These models can be fine-tuned or adapted for specific tasks, like answering questions, translating languages, summarizing texts, or even creating content for specific domains like legal, medical, or technical fields.

The Transformer: The Engine Behind LLMs

At the heart of most LLMs is the Transformer architecture, introduced in the paper "Attention Is All You Need" by Vaswani et al. (2017). Imagine the Transformer as an advanced orchestra, where different instruments (layers and attention mechanisms) work in harmony to understand and generate language.

TensorFlow and Keras: Your Building Blocks

TensorFlow, with its high-level API Keras, is like the set of high-quality tools and materials you need to start painting. It simplifies building and training complex models.

Building the Transformer with TensorFlow and Keras

Step 1: Setting Up Your Environment

Before diving into code, ensure you have TensorFlow installed in your Python environment:

      pip install tensorflow
    

Step 2: The Encoder and Decoder Layers

The Transformer model consists of encoders and decoders. Think of encoders as scribes, absorbing information, and decoders as orators, producing meaningful language.

Encoder Layer:

      import tensorflow as tf
from tensorflow.keras.layers import MultiHeadAttention, LayerNormalization, Dense

class TransformerEncoderLayer(tf.keras.layers.Layer):
    def __init__(self, d_model, num_heads, dff, rate=0.1):
        super(TransformerEncoderLayer, self).__init__()
        self.mha = MultiHeadAttention(num_heads, d_model)
        self.ffn = tf.keras.Sequential([
            Dense(dff, activation='relu'), 
            Dense(d_model)
        ])

        self.layernorm1 = LayerNormalization(epsilon=1e-6)
        self.layernorm2 = LayerNormalization(epsilon=1e-6)
        self.dropout1 = tf.keras.layers.Dropout(rate)
        self.dropout2 = tf.keras.layers.Dropout(rate)
    
    def call(self, x, training):
        attn_output = self.mha(x, x, x)
        attn_output = self.dropout1(attn_output, training=training)
        out1 = self.layernorm1(x + attn_output)

        ffn_output = self.ffn(out1)
        ffn_output = self.dropout2(ffn_output, training=training)
        out2 = self.layernorm2(out1 + ffn_output)

        return out2
    

This piece of code defines a Transformer Encoder Layer using TensorFlow and Keras, which are powerful tools for building neural networks. Let’s break the code down:

Import Statements:
      import tensorflow as tf
from tensorflow.keras.layers import MultiHeadAttention, LayerNormalization, Dense
    

Here, we import TensorFlow and specific layers from Keras needed for building the encoder layer. These layers include MultiHeadAttention for handling the attention mechanism, LayerNormalization for stabilizing the neural network, and Dense for fully connected layers.

Defining the TransformerEncoderLayer Class:
      class TransformerEncoderLayer(tf.keras.layers.Layer):
    

This line begins the definition of the TransformerEncoderLayer class, which inherits from TensorFlow's Layer class. This custom layer will form one part of the Transformer model.

Initialization Method (__init__):
      def __init__(self, d_model, num_heads, dff, rate=0.1):
    super(TransformerEncoderLayer, self).__init__()
    

The __init__ method initializes the encoder layer. It takes several parameters:

  • d_model: The dimensionality of the input (and output) of the layer.
  • num_heads: The number of heads in the multi-head attention mechanism.
  • dff: The dimensionality of the inner layer in the feed-forward network.
  • rate: The dropout rate used for regularization.
Multi-Head Attention and Feed-Forward Network:
      self.mha = MultiHeadAttention(num_heads, d_model)
self.ffn = tf.keras.Sequential([Dense(dff, activation='relu'), Dense(d_model)])
    

The encoder layer consists of a multi-head attention mechanism and a feed-forward neural network. self.mha is an instance of MultiHeadAttention, and self.ffn is a simple two-layer feed-forward network with a ReLU activation in between.

Layer Normalization and Dropout:
      self.layernorm1 = LayerNormalization(epsilon=1e-6)
self.layernorm2 = LayerNormalization(epsilon=1e-6)
self.dropout1 = tf.keras.layers.Dropout(rate)
self.dropout2 = tf.keras.layers.Dropout(rate)
    

These lines create instances of layer normalization and dropout layers. Layer normalization helps in stabilizing the output of each layer, and dropout prevents overfitting.

Attention and Feed-Forward Operations:
      attn_output = self.mha(x, x, x)
attn_output = self.dropout1(attn_output, training=training)
out1 = self.layernorm1(x + attn_output)
    

Here, the layer processes its input x through the multi-head attention mechanism, applies dropout, and then layer normalization. It's followed by the feed-forward network operation and another round of dropout and normalization.

Decoder Layer:

      class TransformerDecoderLayer(tf.keras.layers.Layer):
    def __init__(self, d_model, num_heads, dff, rate=0.1):
        super(TransformerDecoderLayer, self).__init__()
        self.mha1 = MultiHeadAttention(num_heads, d_model)
        self.mha2 = MultiHeadAttention(num_heads, d_model)

        self.ffn = tf.keras.Sequential([
            Dense(dff, activation='relu'), 
            Dense(d_model)
        ])

        self.layernorm1 = LayerNormalization(epsilon=1e-6)
        self.layernorm2 = LayerNormalization(epsilon=1e-6)
        self.layernorm3 = LayerNormalization(epsilon=1e-6)
        
        self.dropout1 = tf.keras.layers.Dropout(rate)
        self.dropout2 = tf.keras.layers.Dropout(rate)
        self.dropout3 = tf.keras.layers.Dropout(rate)

    def call(self, x, enc_output, training, look_ahead_mask, padding_mask):
        attn1, attn_weights_block1 = self.mha1(x, x, x, look_ahead_mask)
        attn1 = self.dropout1(attn1, training=training)
        out1 = self.layernorm1(attn1 + x)

        attn2, attn_weights_block2 = self.mha2(enc_output, enc_output, out1, padding_mask)
        attn2 = self.dropout2(attn2, training=training)
        out2 = self.layernorm2(attn2 + out1)

        ffn_output = self.ffn(out2)
        ffn_output = self.dropout3(ffn_output, training=training)
        out3 = self.layernorm3(ffn_output + out2)

        return out3, attn_weights_block1, attn_weights_block2
    

The Transformer Decoder is an essential part of the Transformer model, often used in tasks like machine translation, text generation, and more. Let’s break down the parts of the code that are new:

Attention layers

Two multi-head attention layers (mha1 and mha2) are defined. mha1 is used for self-attention within the decoder, and mha2 is used for attention over the encoder's output. The feed-forward network (ffn) follows a similar structure to the encoder.

The call Method:
      def call(self, x, enc_output, training, look_ahead_mask, padding_mask):
    

This method is where the layer's operations are defined. It takes additional parameters compared to the encoder:

  • enc_output: Output from the encoder.

  • look_ahead_mask: To mask future tokens in a sequence (for self-attention).

  • padding_mask: To mask padded positions (for encoder-decoder attention).

Attention and Feed-Forward Operations:
      attn1, attn_weights_block1 = self.mha1(x, x, x, look_ahead_mask)
attn2, attn_weights_block2 = self.mha2(enc_output, enc_output, out1, padding_mask)
    

The decoder processes its input through two multi-head attention layers. The first one (attn1) is self-attention with a look-ahead mask, and the second one (attn2) focuses on the encoder's output. This is followed by the feed-forward network. Each step involves dropout and normalization.

Step 3: Assembling the Transformer

Think of this step as assembling your orchestra. Each encoder and decoder layer is an instrument, and you're arranging them to create harmony.

Full Transformer Model:

      class Transformer(tf.keras.Model):
    def __init__(self, num_layers, d_model, num_heads, dff, input_vocab_size, 
                 target_vocab_size, pe_input, pe_target, rate=0.1):
        super(Transformer, self).__init__()
        self.encoder = Encoder(num_layers, d_model, num_heads, dff, 
                               input_vocab_size, pe_input, rate)
        self.decoder = Decoder(num_layers, d_model, num_heads, dff, 
                               target_vocab_size, pe_target, rate)

        self.final_layer = tf.keras.layers.Dense(target_vocab_size)

    def call(self, inp, tar, training, enc_padding_mask, 
             look_ahead_mask, dec_padding_mask):
        enc_output = self.encoder(inp, training, enc_padding_mask)
        dec_output, attention_weights = self.decoder(
            tar, enc_output, training, look_ahead_mask, dec_padding_mask)

        final_output = self.final_layer(dec_output)

    

Training the model

With the Transformer model assembled, it's time to train it. This process is like teaching the orchestra to play a symphony, where the symphony is the task you want your model to perform (e.g., language translation, text generation).

Preparing the Data

Data preparation involves collecting a large dataset of text and processing it into a format suitable for training. TensorFlow's data API can be used for this purpose.

Training Loop

      for epoch in range(epochs):
    # Initialize the training step
    for (batch, (inp, tar)) in enumerate(dataset):
        # Training code here
    

However in this following section we will explore how to leverage existent LLMs by using Transfer Learning.

Implementing transfer learning with Hugging Face

Transfer learning in the context of LLMs is akin to an apprentice learning from a master craftsman. Instead of starting from scratch, you leverage a pre-trained model and fine-tune it for your specific task. Hugging Face provides an extensive library of pre-trained models which can be fine-tuned for various NLP tasks.

Setting Up Hugging Face Transformers

First, you need to install the Hugging Face transformers library:

      pip install transformers
    

Loading a Pre-Trained Model

Choose a pre-trained model from Hugging Face's model hub. For this example, let's use bert-base-uncased, a popular BERT model:

      from transformers import BertTokenizer, TFBertModel

tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = TFBertModel.from_pretrained('bert-base-uncased')
    

Preparing Data for Fine-Tuning

Suppose you're fine-tuning the model for a sentiment analysis task. First, preprocess your data:

      # Example sentences
sentences = ["I love this product!", "This is a bad product."]

# Tokenize sentences
inputs = tokenizer(sentences, padding=True, truncation=True, return_tensors="tf")
    

Notize we have to use BERT tokenizer to ensure everything is tokenized and padded exactly as BERT likes.

Fine-Tuning the Model

Now, you can add a classification layer on top of the pre-trained model and fine-tune it:

      from tensorflow.keras.layers import Input, Dense
from tensorflow.keras.models import Model

# Define input layers
input_ids = Input(shape=(None,), dtype='int32', name="input_ids")
attention_mask = Input(shape=(None,), dtype='int32', name="attention_mask")

# Load the pre-trained BERT model
bert = model(input_ids, attention_mask=attention_mask)

# Add a classification layer on top
x = bert.last_hidden_state[:, 0, :]
x = Dense(128, activation='relu')(x)
output = Dense(1, activation='sigmoid')(x)

# Construct the final model
fine_tuned_model = Model(inputs=[input_ids, attention_mask], outputs=[output])

# Compile the model
fine_tuned_model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])

# Example labels for the sentences
labels = [1, 0]  # 1 for positive, 0 for negative sentiment

# Train the model
fine_tuned_model.fit(inputs, labels, epochs=3, batch_size=32)

    

Testing the Fine-Tuned Model

Finally, test the fine-tuned model on new sentences:

      test_sentences = ["I am not sure about this product.", "Absolutely fantastic!"]
test_inputs = tokenizer(test_sentences, padding=True, truncation=True, return_tensors="tf")

predictions = fine_tuned_model.predict(test_inputs)

# Interpret the predictions
for sentence, prediction in zip(test_sentences, predictions):
    sentiment = "Positive" if prediction > 0.5 else "Negative"
    print(f"Sentence: '{sentence}' - Sentiment: {sentiment}")
    

Conclusion

Creating an LLM from scratch is an intricate yet immensely rewarding process. By understanding and building upon the Transformer architecture with TensorFlow and Keras, and leveraging transfer learning through Hugging Face, you can create a model that's not just a powerful NLP tool but a reflection of your unique approach to understanding language.

As you continue on this journey, remember that the field of NLP is ever-evolving, and there's always more to learn and explore. Happy modeling!

Further learning resources

Axel Sirota

Axel S.

Axel Sirota is a Microsoft Certified Trainer with a deep interest in Deep Learning and Machine Learning Operations. He has a Masters degree in Mathematics and after researching in Probability, Statistics and Machine Learning optimization, he works as an AI and Cloud Consultant as well as being an Author and Instructor at Pluralsight, Develop Intelligence, and O'Reilly Media.

More about this author