Optimisation Strategies to Speed up Transformers

Transformers have revolutionised the field of natural language processing (NLP) and have found applications in various domains such as machine translation, text generation, and more. The transformer architecture employs self-attention mechanisms to process input data in parallel, addressing the limitations of recurrent neural networks (RNNs) and achieving state-of-the-art performance on many tasks. Despite their success, transformers are computationally intensive and often require significant resources in terms of memory and processing power. This makes optimising transformers critical for practical applications, especially in environments with limited computational resources or where low-latency responses are crucial.

Types of Optimisation strategies

Optimisation strategies for transformers can be broadly categorised into:

• Algorithmic Improvements - involve modifying the transformer architecture or training procedures to enhance efficiency. Techniques such as reducing the complexity of self-attention mechanisms, employing sparse attention, and utilising knowledge distillation fall under this category.
• Hardware Accelerations - leverage specialised hardware like GPUs, TPUs, and FPGAs to speed up computations. These accelerations are often coupled with low-level optimisations and parallel processing capabilities to maximize performance.
• Software Engineering practices - focus on optimising the implementation of transformer models. This includes leveraging optimised libraries, utilising mixed-precision training, and implementing efficient data handling and batching processes.

In this exploration, we delve into Software Engineering practices and how they can be used to efficiently speed up Transformers. By examining these strategies, we aim to provide a comprehensive understanding of the current state-of-the-art methods for enhancing transformer efficiency, thus enabling their wider and more effective deployment across different applications and industries.

Optimisation Strategies

Fixed-Length Padding, Dynamic Padding and Uniform Length Batching

Fixed-length padding is the process of ensuring that all input sequences in a batch fed into the model have the same length by adding padding tokens to shorter sequences.

Let's illustrate fixed-length padding with example. Consider the following set of sentences:

1. "She enjoys reading books."
2. "He loves coding in Python."

If the maximum sequence length allowed by the model is 8 tokens, we need to pad each sentence to ensure they all have the same length. Here is how it would look:

Original Sentences:

1. "She enjoys reading books."
2. "He loves coding in Python."

Tokenised Sentences:

1. "She enjoys reading books [PAD] [PAD] [PAD] [PAD]"
2. "He loves coding in Python [PAD] [PAD] [PAD]"

Dynamic padding involves adjusting the length of input sequences by padding them to match the maximum length within each specific batch, rather than padding all sequences to the length of the longest sequence in the entire dataset. Unlike fixed-length padding, which standardizes sequence lengths across the entire dataset, dynamic padding ensures that sequences are only padded to the length of the longest sequence within their respective batch.

Uniform Length Batching involves organising sequences of similar lengths into batches for training or inference. Unlike dynamic padding, which pads sequences to the length of the longest sequence within each batch, uniform-length batching creates batches where all sequences are inherently the same length. This is done by dividing the dataset into buckets based on sequence lengths and then forming batches by sampling sequences from these buckets.

As seen from above images, the total tokens using Fixed-Length padding is 168 tokens which is then reduced to 160 tokens using Dynamic Padding and can be further dropped down to 124 tokens using Uniform Length Batching.

Numeric Precision Reduction

Numeric precision reduction speeds up computations by using lower precision floating-point numbers, a common method to enhance prediction speed.

Automatic Mixed Precision (AMP) is a technique that accelerates deep learning model training by combining single-precision (float32) and half-precision (float16) arithmetic. Modern GPUs handle float16 computations faster and with less memory than float32.

AMP dynamically changes the precision during training, using float16 for most calculations but switching to float32 when necessary to avoid numerical issues like underflow or overflow.

Parameter Sharing

Parameter sharing in transformers is a technique used to reduce the number of parameters in a model by reusing the same parameters across multiple layers. This approach can help mitigate the memory and computational overhead of deep transformer models, making them more efficient while maintaining their performance.

Freezing Embeddings

Freezing the embedding layer during the fine-tuning of a pre-trained language model is a straightforward yet effective technique to enhance training efficiency. The embedding layer, which converts input tokens into dense vectors, is computationally intensive and memory-hungry. By freezing this layer, you prevent its weights from being updated during back propagation. This approach offers several benefits:

1. Memory Efficiency: By not updating the embedding layer, you save significant GPU memory. This freed-up memory can be utilized to increase the batch size, which can lead to more stable and faster convergence during training.
2. Training Speed: Since the embedding layer's weights are not adjusted, fewer computations are required during each training step. This reduction in computational load speeds up the overall training process.
3. Stability: Pre-trained embeddings already capture a lot of useful information about the language. Freezing them ensures that this valuable information is preserved, which can lead to better performance on downstream tasks, especially when the amount of fine-tuning data is limited.
4. Simplified Optimisation: With fewer parameters to update, the optimization process becomes simpler and potentially more effective, as the model can focus on adjusting the more task-specific layers.

By freezing the embedding layer, you leverage the pre-trained knowledge effectively while optimising resource usage, making it a practical strategy for fine-tuning large language models.

Gradient Accumulation

Gradient accumulation is a strategy to optimize computational resources, especially when GPU memory is limited. It works by accumulating gradients over several smaller batches of data rather than updating model parameters after each small batch.

• Mini-batch: It's a subset of the training data, used in batches during model training to compute gradients and update parameters. Smaller mini-batches consume less memory but may lead to noisy gradient estimates.
• Gradients: These are derivatives of the loss function with respect to the model parameters. They guide the parameter updates during training, aiming to minimize the loss.

By accumulating gradients over multiple mini-batches, gradient accumulation allows for a larger effective batch size without requiring all data to be stored in memory simultaneously. This method balances memory usage with computational efficiency in training deep learning models.

Gradient Checkpointing

Gradient checkpointing is a technique that cuts down on memory usage during deep neural network training, at the cost of slightly longer computation times. It trades computational resources to save memory, making it possible to train bigger or wider models and use larger mini-batch sizes.

Here's how it works:

• Back propagation in chunks: Instead of storing the entire backpropagation graph in memory, gradient checkpointing breaks down the process into smaller parts.
• Recalculation of activations: It saves memory by recalculating intermediate activations during the backward pass, rather than storing them.

Although this approach increases computation due to extra forward passes for each chunk, it significantly reduces memory demands. Despite slowing down training, it enables handling larger batches effectively, which can be advantageous in certain training scenarios.

Conclusion

Optimisation strategies to speed up transformers are essential for harnessing their full potential in practical applications. By implementing techniques such as fixed-length and dynamic padding, uniform length batching, and numeric precision reduction, we can significantly improve the efficiency and performance of transformer models. Methods like freezing embeddings and gradient accumulation further enhance training speed and resource utilisation.

As transformer models continue to grow in complexity and application scope, these optimisation strategies will play a critical role in making them accessible and usable in real-world scenarios. By reducing computational demands and accelerating training times, these methods enable more efficient use of resources, allowing for the deployment of powerful transformer models in environments with limited computational capabilities. Ultimately, these optimizations contribute to the broader adoption and impact of transformer-based technologies across various domains.