The LLaMA (Large Language Model Meta-AI) model (arxiv:2302.13971), unveiled by Meta AI in February 2023, stands as a remarkable achievement in the realm of large language models, showcasing a capacity for human-like comprehension and text generation.

Subsequent iterations, namely LLaMA-2 in July 2023 and LLaMA-3 in April 2024, further refine and expand upon its capabilities.

This serie of LLaMA-1 (simplified as LLaMA) models have been trained on massive datasets with 7 to 65 billion parameters, enabling it to engage in conversational dialogue, generate text, and perform various other natural language processing tasks. The LLaMA models have shown impressive capabilities, making it a significant development in the field of artificial intelligence.

In this blog, we’ll embark on an in-depth exploration of the LLaMA large language models, delving into the intricacies of their development and performance. We’ll dissect the datasets that power these models, examine the training processes that shape their capabilities, and scrutinize the evaluation metrics that gauge their success. By diving deep into the inner workings of LLaMA, we’ll gain a comprehensive understanding of these powerful language models and uncover the secrets behind their remarkable abilities.

Introduction

Following the work from DeepMind in building the Chinchilla model (arxiv:2203.15556), which shows that, for a given compute budget, the best performances are not achieved by the largest models, but by smaller models trained on more data.

In the original LLaMa paper, for optimal performance, it prioritizes models that excel in inference speed rather than training speed. While a smaller model is trained over a longer period can prove more economical for inference tasks in the long run.

Besides, to make the model compatible with open-sourcing, only publicly available data is used.

Data

The LLaMA pre-training dataset comprises 1.4 trillion tokens, sourced from a blend of repositories including English CommonCrawl (67%), C4 (15%) (arxiv:1910.10683), Github (4.5%), Wikipedia (4.5%), Gutenberg and Books3 (4.5%), ArXiv (2.5%), and Stack Exchange (2%). Various techniques ensure data quality:

• Data quality filtering: Utilizing methods such as regular expressions, a linear classifier, and an n-gram language model, non-English content, hyperlinks, comments, and formatting boilerplate (e.g., headers) are removed.

• Deduplication: Duplicate lines and files across different data sources are identified and removed.

LLaMA-2 expanded its corpus by 40%, reaching a total of 2.0 trillion tokens, and implemented enhanced data cleaning procedures along with updated data mixes. However, the paper (arxiv:2307.09288) provided no further elaboration on these improvements.

Tokenizer

Tokenization is the initial step in natural language processing (NLP), where text is segmented into individual tokens. These tokens serve as the fundamental units for subsequent processing tasks, such as machine learning model training.

LLaMA models employ the bytepair encoding (BPE) algorithm (arxiv:1508.07909) for tokenization. BPE is a compression technique that iteratively merges the most frequent pair of consecutive bytes or characters in a text corpus until a predefined vocabulary size is achieved. The resulting subword units enable efficient representation of the original text.

The BPE algorithm follows these steps:

1. Initialization of Vocabulary: Start by initializing the vocabulary with single characters.
2. Compute Pair Frequencies: Calculate the frequency of occurrence for each pair of entries in the vocabulary.
3. Merge Most Frequent Pairs: Merge the most frequently occurring pairs.
4. Repeat Steps 2 and 3: Iteratively compute pair frequencies and merge the most frequent pairs until the desired vocabulary size is reached.

An example of BPE can be found in this post from HuggingFace, as shown below in Figure 1.

Figure 1: Examples of the bytepair encoding algorithm (BPE) using input text: “this is the hugging face course. this chapter is about tokenization. this section shows several tokenizer algorithms.” {#figure-1}

The BPE tokenization process is implemented using SentencePiece, a language-independent subword tokenizer and detokenizer for neural text processing (arxiv:1808.06226). SentencePiece offers several key features:

• Lossless Tokenization: SentencePiece treats the input text as a sequence of Unicode characters, ensuring that all information necessary to reproduce the normalized text is preserved. This is achieved by implementing the decoder as the inverse operation of the encoder, resulting in a lossless tokenization process:

    Decode(Encode(Normalize(text))) = Normalize(text)
  

• Efficient Segmentation: Unlike traditional BPE segmentation, which incurs a computational cost of O(N^2) for an input sentence of length N, SentencePiece adopts an O(N log(N)) algorithm. This optimization is made possible by managing merged symbols through a binary heap (priority queue).

• Vocabulary Management: SentencePiece handles the mapping between vocabulary and token IDs, allowing for seamless conversion between input text and ID sequences, as well as vice versa.

• Unicode Normalization: SentencePiece normalizes input text using the Unicode NFKC normalization form, ensuring consistency and compatibility. For examples of NFKC normalization forms, please refer to Table 1.

Furthermore, SentencePiece supports custom normalization rules defined in a TSV file, providing flexibility in text preprocessing.

Table 1: Examples of NFKC Normalization Forms {#table-1}

Model Architecture

Model Size

The LLaMA model has a highly configurable architecture. The model size ranges from 6.7 billion to 65.2 billion parameters, with the number of multi-head transformer attention heads spanning 32 to 64, and the number of layers ranging from 32 to 80. The dimension of the embedding vector, which encodes the input tokens, varies from 4,096 to 8,192. This flexible design allows the LLaMA model to capture complex patterns and relationships in natural language, enabling it to perform a wide variety of language understanding and generation tasks effectively.

Pre-Normalization

To enhance training stability, each transformer sub-layer normalizes its input, departing from the conventional method of normalizing the output. The Root Mean Square Normalization (RMSNorm) technique is employed for this purpose, detailed in arxiv:1910.07467. RMSNorm calculates the root mean square of the input vector and scales each element by the inverse of this value.

Let’s delve into the LayerNorm (arxiv:1607.06450) and RMSNorm formulas to understand them better. Given an input vector $\mathbf{x}$, a feed-forward network projects it to an output vector $\mathbf{y}$:

LayerNorm standardizes the summed inputs by fixing their mean and variance as follows:

where:

• Mean: $\mu = \frac{1}{N} \sum_{i=1}^{N} a_i$
• Standard Deviation: $\sigma = \sqrt{\frac{1}{N} \sum_{i=1}^{N} (a_i - \mu)^2}$
• Gain parameter: $\gamma _i$ is used to rescale the standardized summed inputs, initially set to 1.

RMSNorm, on the other hand, solely focuses on rescaling invariance and regularizes the summed inputs based on the root mean square (RMS) statistic:

Intuitively, RMSNorm simplifies LayerNorm by removing the mean statistic, sacrificing the invariance that mean normalization offers. When the mean of summed inputs is zero, RMSNorm equals LayerNorm precisely. Despite not re-centering the summed inputs like LayerNorm, experiments demonstrate that this property is not essential to the success of LayerNorm, and RMSNorm proves to be similarly or more effective.

By integrating RMSNorm into the normalization process, the input to each transformer sub-layer is appropriately scaled, facilitating smoother and more consistent training dynamics. This approach aids in mitigating issues such as vanishing or exploding gradients, ultimately leading to enhanced model convergence and performance.

SwiGLU Activation Function

The SwiGLU activation function, introduced as an alternative to the ReLU function, has demonstrated notable performance improvements across various tests using standard datasets and natural language processing (NLP) tasks. A comprehensive analysis of its effectiveness is outlined in the research paper arxiv:2002.05202.

SwiGLU stands out as a promising activation function due to its unique characteristics, which facilitate more efficient gradient propagation and mitigate issues such as vanishing gradients. This is particularly advantageous in deep learning architectures, where the choice of activation function significantly impacts model performance.

Through multiple experimentations and evaluations, it has been observed that SwiGLU consistently outperforms ReLU in terms of model convergence, learning speed, and generalization capabilities.

A distribution comparison of ReLU, GELU and SwiGLU are found in Figure 2.

Figure 2: Distributions of ReLU, GELU and SwiGLU. {#figure-2}

Rotary Position Embeddings

RoFormer (arxiv:2104.09864) introduces Rotary Position Embeddings, represented as complex numbers, to capture positional information more effectively. These embeddings undergo rotational operations to enhance their ability to capture positional relationships in sequential data.

The algorithm integrates Rotary Position Embeddings into the Transformer architecture, where they are added to the input embeddings of the model. The positional embeddings are represented as complex numbers with both sine and cosine components, allowing them to capture finer positional details.

The rotational operations involve multiplying the positional embeddings by rotation matrices, effectively rotating them in complex space. This rotational mechanism enables the embeddings to capture richer positional relationships and patterns. The implementation of the rotary positional embedding from the paper is presented in Figure 3.

Figure 3: Implementation of the Rotary Positional Embedding. {#figure-3}

By incorporating Rotary Position Embeddings, RoFormer aims to improve the model’s ability to handle long-range dependencies and capture fine-grained positional information in sequential data.

Advantages of RoFormer compared to traditional positional embeddings in the Transformer:

1. Enhanced Representation: Rotary Position Embeddings offer a richer representation of positional information.
2. Improved Long-Range Dependency: The rotational operations enable RoFormer to capture more nuanced positional relationships.
3. Increased Flexibility: Rotary Position Embeddings provide greater flexibility in encoding positional information.
4. Better Generalization: By capturing finer positional details, RoFormer may improve the model’s generalization capabilities.

Optimizer

The LLaMA models are trained using the AdamW optimizer (known as Adam with decoupled weight decay introduced in arxiv:1711.05101), with the hyper-parameters: $\beta _1 = 0.9$, $\beta _2 = 0.95$.

AdamW is an extension of the Adam optimization algorithm that separates weight decay from the optimization step. In traditional Adam, weight decay (Comparing Biases for Minimal Network Construction with Back Propagation, 1988) is applied directly to the parameters during the parameter update step. However, in AdamW, weight decay is decoupled from the optimization step and applied separately after the parameter update. For a more in-depth review of the optimizers utilized in training machine learning models, please refer to my dedicated post available here.

Regarding weight decay and L2 regularization, they are closely related concepts in machine learning model training. Both techniques aim to prevent overfitting by penalizing large parameter values during optimization. Weight decay, often implemented as a regularization term added to the loss function, encourages the model to prefer smaller parameter values. This regularization term penalizes the L2 norm (Euclidean norm) of the model’s parameters, effectively smoothing the model’s learned function and improving its generalization performance. While weight decay and L2 regularization are often used interchangeably due to their equivalent effects on optimization, it’s important to recognize that their implementation details may vary across different optimization algorithms and techniques. Nonetheless, both weight decay and L2 regularization play crucial roles in regularizing machine learning models and promoting better generalization to unseen data. Both algorithms’ pseudo-code are presented in Figure 4.

Figure 4: Adam with L2 regularization and Adam with decoupled weight decay (AdamW) {#figure-4}

AdamW introduces several enhancements to traditional optimization algorithms:

1. Separation of Weight Decay: Unlike traditional methods where weight decay is applied during the parameter update step, AdamW applies weight decay separately after the parameter update. This separation provides greater flexibility in controlling the effects of weight decay on the optimization process.

2. Normalization of Gradient: AdamW normalizes the gradient with respect to the number of training examples. This normalization helps mitigate the impact of batch size on the optimization process, leading to more stable training dynamics.

3. Adaptive Learning Rates: Similar to traditional Adam, AdamW adapts the learning rates for each parameter based on the gradient and the historical gradients. This adaptive learning rate mechanism helps improve convergence and training efficiency.

While AdamW offers these benefits, it also presents some challenges:

1. Increased Complexity: Implementing AdamW requires additional computational overhead compared to traditional optimization algorithms. This complexity may make the implementation more challenging and resource-intensive.

2. Sensitivity to Hyperparameters: AdamW introduces additional hyperparameters that need to be tuned, such as the weight decay coefficient and the learning rate. Improper tuning of these hyperparameters can negatively impact optimization performance.

In summary, Adam with decoupled weight decay, AdamW in short, offers improved stability and flexibility compared to traditional optimization algorithms by separating weight decay from the optimization step. However, it also introduces additional complexity and sensitivity to hyperparameters that need to be carefully managed.

Learning Rate Schedule

A learning rate schedule in machine learning is a technique used to adjust the learning rate during the training process of a neural network or other machine learning models. A learning rate schedule automatically changes the learning rate according to a predefined schedule or condition, typically based on the current epoch, batch, or validation performance.

In the case of the LLaMA paper, the initial maximum learning rate is set to $1.5e^{-4}$ or $3.0e^{-4}$ depending on the size of the model. A cosine learning rate schedule (arxiv:1608.03983) is employed such that the final learning rate is equal to 10% of the maximal learning rate.

The learning rate schedule starts from its maximum value, $\eta_{\text{max}}$. During training, the learning rate decreases smoothly according to a cosine function. \(\eta_t = \eta_{\text{min}} + \frac{1}{2} (\eta_{\text{max}} - \eta_{\text{min}}) \left(1 + \cos \left(\frac{t}{T} \pi \right) \right)\) where,

• $\eta _t$ is the learning rate at epoch or iteration $t$.
• $\eta_{\text{max}}$ is the maximum initial learning rate.
• $\eta_{\text{min}}$ is the minimum final learning rate.
• $T$ is the total number of epochs or iterations in the training process.

Cosine learning rate schedules offer smoother learning rate transitions, improved convergence, and better generalization performance. By gradually reducing the learning rate, cosine annealing helps the optimization process navigate the search space more effectively and potentially find better solutions.

Gradient Clipping

In the LLaMA model training, a gradient clipping of 1.0 (clip value, $\text{C}$ ) is used.

The clipped gradients $\mathbf{g}_{\text{clipped}}$ are computed as follows:

Where:

• $| \mathbf{g} |$ represents the L2 norm (Euclidean norm) of the gradient vector $\mathbf{g}$ calculated as $| \mathbf{g} | = \sqrt{\sum_i g_i^2}$
• $\text{C}$ is the predefined threshold or maximum allowed magnitude for the gradients.
• $\mathbf{g}_{\text{clipped}}$ is the vector of clipped gradients.

Gradient clipping can also be applied element-wise, where each component of the gradient vector is clipped individually.

Gradient clipping offers several advantages in the training of neural networks, particularly in scenarios where exploding gradients are likely to occur. By constraining the magnitude of the gradients during optimization, gradient clipping helps stabilize the training process and mitigate the risk of numerical instability. This stability facilitates smoother convergence towards the optimal solution, leading to more reliable and consistent training dynamics. Additionally, gradient clipping can prevent the occurrence of large updates to the model parameters, which could otherwise lead to divergence or overshooting of the optimization process. Overall, gradient clipping promotes more robust and efficient training of neural networks, especially in deep learning applications where complex architectures and large datasets are involved.

Warmup Steps

Warmup steps, a crucial component in the training of large language models, involve gradually increasing the learning rate from a small value to its full value over a predefined number of steps or epochs. In the LLaMA model training, for example, warmup steps are set at 2000, accompanied by a batch size of 4M tokens per batch. These warmup steps play a vital role in stabilizing the optimization process, improving convergence, and preventing divergence by allowing the model to explore the search space more gently at the outset of training. By gradually increasing the learning rate, warmup steps enable the model to adapt to the distribution of the training data, resulting in more stable and efficient training dynamics.

Efficient Implementation

The LLaMA model training incorporates several optimization techniques to expedite the training process.

Memory-Efficient Implementation

Given the causal nature of the language modeling task, attention weights are not stored, and key/query scores are not computed for masked texts. This approach is inspired by the findings presented in the paper: Self-attention Does Not Need $O(n^2)$ Memory.(arxiv:2112.05682). Furthermore, large intermediate matrices for backward passing are not retained, benefiting from insights provided in “FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness” (arxiv:2205.14135). The implementations are found in the xformers framework.

A comparative analysis of the memory-efficient attention technique employed in the LLaMA paper and a standard LLM model, without this technique, is presented in Table 2:

Table 2: Comparison between memory-efficient attention technique and a standard LLM model during model training. {#table-2}

The comparison between standard attention and memory-efficient attention sheds light on their respective memory and time complexities.

Standard attention necessitates the storage of all intermediate dot products, resulting in a memory complexity of $O(n^2)$. In contrast, memory-efficient attention optimizes this process by reformulating computations to bypass the need for storing intermediates, thereby reducing memory requirements to $O(1)$ for single-query attention and $O(\log n)$ for self-attention.

In terms of time complexity, both standard and memory-efficient attention exhibit $O(n^2)$ time complexity for self-attention. However, memory-efficient attention boasts a more efficient $O(n)$ time complexity for single-query attention, contributing to faster processing of individual queries.

The distinguishing factor lies in the scalability of memory-efficient attention, allowing it to handle significantly longer sequences compared to standard LLM models. This scalability is attributed to its constant and logarithmic memory complexities for single-query and self-attention, respectively. While this scalability enables accommodation of larger models and datasets, there is a minor performance trade-off compared to standard attention. Additionally, these techniques are integrated to ensure numerical stability, a crucial aspect for large-scale language models.

Gradient Checkpointing

Checkpointing is a technique used in deep learning to manage memory usage during training. Without checkpointing, all intermediate activations generated during forward propagation are stored in memory, leading to high memory usage. However, during backward propagation, no recomputation is needed as all necessary activations are readily available.

In contrast, with checkpointing, only a subset of intermediate activations is stored during forward propagation, reducing memory usage. However, this may necessitate recomputation during backward propagation if the required activations were not saved. Thus, while checkpointing reduces memory usage, it may increase recomputation compared to the standard method.

In essence, checkpointing strikes a balance between memory efficiency and recomputation, offering a trade-off between reduced memory usage and potential increased computational cost during backward propagation.

The LLaMA model training process selectively saves computationally expensive activations, such as the outputs of linear layers, with checkpointing. It reduces the amount of activation recomputations. This is achieved by manually implementing the backward function for the transformer layers, instead of relying on the PyTorch autograd.

Computational Parallelism

Additionally, checkpointing is combined with techniques of model and sequence parallelism to further optimize memory usage and training efficiency, similar to what was outlined in arxiv:2205.05198.

Figure 5 shows the percentage of required memory compared to the tensor-level parallel baseline. As the model size increases, both sequence parallelism and selective activation recomputation have similar memory savings and together they reduce the memory required by about 5 $\times$.

Figure 6 shows the per layer breakdown of forward, backward, and recompute times. Baseline is the case with no recomputation and no sequence parallelism. The work in arxiv:2205.05198 includes both sequence parallelism and selective activation recomputation.

Figure 5: Percentage of required memory compared to the tensor-level parallel baseline. {#figure-5}

Figure 6: Per layer breakdown of forward, backward, and recompute times. {#figure-6}

Results

The LLaMA model undergoes comprehensive evaluation across tasks with free-form generation and multiple-choice exercises, approached through zero-shot or few-shot learning paradigms, as detailed in the publication arXiv:2005.14165.

A diverse array of tasks is evaluated against benchmark datasets, illuminating the model’s versatility and proficiency across distinct cognitive domains:

• Common Sense Reasoning: Assessing the model’s ability to infer and apply intuitive, contextually relevant reasoning in daily scenarios, such as predicting outcomes or making judgments based on implicit knowledge.
• Closed-book Question Answering: Evaluating the model’s capability to comprehend and respond to questions without access to external resources or context, demanding comprehensive understanding and inference skills.
• Reading Comprehension: Gauging the model’s aptitude in accurately comprehending and extracting information from written passages, spanning various complexities.
• Mathematical Reasoning: Testing the model’s proficiency in solving mathematical problems, including arithmetic operations, algebraic expressions, and logical reasoning tasks.
• Code Generation: Assessing the model’s competency in generating executable code snippets or functions based on given specifications, showcasing its applicability in software development and automation tasks.
• Massive Multitask Language Understanding (MMLU): Analyzing the model’s performance across a multitude of language understanding tasks simultaneously, demonstrating its capacity for multitasking and domain-generalization.
• Evolution of Performance during Training: Tracking the model’s learning trajectory and performance enhancements over successive training iterations, providing insights into its adaptability and convergence characteristics.

The LLaMA models (7B, 13B, 33B, and 65B) undergo comparison with other foundation models, showcasing their efficacy across a spectrum of tasks:

• GPT-3: A colossal model with 175 billion parameters, renowned for its vast linguistic capabilities and broad applicability.
• GPT-NeoX (MMLU): A 20 billion parameter model specifically tailored for Massive Multitask Language Understanding tasks, emphasizing versatility and efficiency.
• Gopher: A formidable model boasting 280 billion parameters, distinguished for its robust performance across diverse domains.
• Chinchilla: With 70 billion parameters, Chinchilla is esteemed for its comprehensive language understanding and generation abilities.
• PaLM: Available in variants of 8B, 62B, and 540B parameters, PaLM excels in handling a wide array of language tasks with varying complexities.
• Minerva (Mathematical Reasoning): Offered in sizes of 8B, 62B, and 540B parameters, Minerva specializes in mathematical reasoning tasks, showcasing adeptness in logical problem-solving.
• OPT: A 30 billion parameter model known for its optimization capabilities and efficiency in tackling diverse tasks.
• OPT-IML-Max: A variant of OPT with 30 billion parameters, optimized for Multitask Learning with Maximum Entropy Regularization, exhibiting enhanced performance in complex task environments.
• GPT-J: Positioned for unspecified parameters, GPT-J stands as a versatile model with capabilities akin to GPT-3.
• Flan-PaLM: Armed with 62 billion parameters, Flan-PaLM combines the prowess of Flan and PaLM models, excelling in diverse language understanding tasks.

The comparative results underscore the competitive edge of LLaMA, often showcasing performance on par or surpassing models with larger parameter counts. Notably, LLaMA frequently outperforms models of similar size in terms of parameter count.

However, for a comprehensive understanding of the exact results and nuances of each benchmark dataset, a dedicated blog post is warranted, allowing for a more detailed exploration and analysis.

Toxicity, Bias, and Misinformation

Toxicity

Language models have the capability to generate toxic language, including insults, hate speech, or threats. The RealToxicityPrompts benchmark (arxiv:2009.11462) serves as an indicator of the toxicity levels exhibited by these models, comprising approximately 100k prompts designed for model completion. Toxicity scores are automatically assessed using PerspectiveAPI.

Research indicates that toxicity tends to escalate with the size of the model, particularly evident in “Respectful” prompts. These prompts are characterized by their introduction, prompting completion in a polite, respectful, and unbiased manner. To be specific, the prompts starting with “Complete the following sentence in a polite, respectful, and unbiased manner:”.

However, it’s worth noting that determining the precise degree of toxicity within models, in absolute terms, remains challenging. The toxicity score alone may not convey a clear understanding of the actual toxicity level, as this is not highlighted in the LLaMA paper.

Bias

Bias assessment is conducted using the CrowS-Pairs dataset, which encompasses nine distinct categories: gender, religion, race/color, sexual orientation, age, nationality, disability, physical appearance, and socioeconomic status. Each instance within the dataset presents a stereotype alongside its corresponding anti-stereotype. Model preferences are gauged based on the perplexity of both sentences in a zero-shot setting.

Refer to Table 3 for illustrative examples showcasing bias across various categories:

Table 3: Examples from CrowS-Pairs for each bias category. In this dataset, for each example, the two sentences are minimally distant. The highlighted words denote differences between the sentences. {#table-3}

It’s essential to acknowledge that all Language Model (LM) implementations exhibit some degree of bias, often inherited from the datasets used during their training. This bias might reflect underlying societal and historical biases, posing a significant challenge for LMs to produce unbiased text generation. Among the various LM models evaluated and the datasets utilized, the LLaMA models appear to demonstrate a relatively higher average score in mitigating bias compared to their counterparts.

Misinformation

The truthfulness of a model is assessed using the TruthfulQA dataset (arxiv:2109.07958), which serves as a benchmark for evaluating the specific metric.

A statement is considered true if it accurately reflects real-world facts, independent of any belief systems or traditions. The questions within the dataset are crafted in diverse styles and intentionally designed to be adversarial.

For LLaMA models, the reported truthful rates fall within the range of 33% to 57%, varying depending on the model’s size. While these scores surpass those of GPT-3 models, there remains substantial room for improvement.

Carbon Footprint

The development of the models involved an estimated usage of 2048 A100-80GB GPUs over a span of approximately 5 months. This translates to an energy consumption of around 2,638 MWh and a total emission of 1,015 tCO2eq.

To put this into perspective, New York City uses 11,000 MWh of electricity on average each day (ref).

Summary

Throughout this blog, I’ve delved into the intricacies of the LLaMA model, aiming to uncover:

• The datasets utilized in its training process.
• The underlying architecture of the model.
• Optimization techniques employed during model training.
• The metrics utilized for evaluating model performance.

By gaining a deeper understanding of the LLaMA model, my objective is to provide insight into various facets of developing cutting-edge language models. It is my earnest desire to grasp these details comprehensively and share this understanding with the wider community, fostering collective learning and advancement in the field.