Turbocharging AI: Unleashing Transformer Power with Mixed Precision in PyTorch

Section 1: Introduction

In our previous exploration, we delved into the theoretical foundations of mixed precision training and its potential to enhance the performance of deep learning models. Now, it’s time to put theory into practice. In this article, we will take a hands-on approach by applying mixed precision training to a real-world task: sentiment analysis using the BERT uncased model.

BERT (Bidirectional Encoder Representations from Transformers) is a state-of-the-art model widely used for Natural Language Processing (NLP) tasks, including sentiment analysis. By leveraging mixed precision training, we aim to explore how different training modes—manual mixed precision, automatic mixed precision, and no mixed precision—impact the performance of BERT in terms of accuracy, precision, recall, F1-score, and speed.

Here is an explanation of BERT model : https://www.geeksforgeeks.org/explanation-of-bert-model-nlp/

Section 2: Mixed Precision Training in Action

Now that we’ve set the stage, it’s time to put mixed precision training to the test. We’ll run three different training modes—no mixed precision, automatic mixed precision (AMP), and manual mixed precision—on the BERT uncased model for sentiment analysis using the IMDb dataset.

The goal is to compare the following metrics for each mode:

• Accuracy: How often the model gets the right prediction.
• Precision: Out of the predicted positive cases, how many were truly positive.
• Recall: Out of all actual positive cases, how many were captured by the model.
• F1-Score: The balance between precision and recall.
• Training and Evaluation Time: How long it takes to train and evaluate the model in each mode.

Also here is the GitHub link to the code with an explanation:

Link to GitHub for Mixed Precision Training

Dataset Details for Our Case:

For our experiment, we worked with a subset of the IMDb dataset consisting of 10,000 samples. This subset was balanced to ensure an equal distribution of positive (label = 1) and negative (label = 0) reviews. Here's a breakdown:

• Total Data Points: 10,000 samples
• Training Set Size: 8,000 samples (4,000 positive and 4,000 negative)
• Test Set Size: 2,000 samples (1,000 positive and 1,000 negative)
• Positive/Negative Distribution: 50% positive, 50% negative in both the training and test sets.

The balanced nature of the dataset ensures that the model does not favor one class over the other, allowing us to better assess the performance of mixed precision techniques without being biased by imbalanced data.

Note: All experiments were conducted using a T4 GPU on Google Colab

2.1 No Mixed Precision (Baseline)

Our first training method doesn't use mixed precision. In this case, the model operates entirely in FP32, offering maximum precision but potentially at the cost of speed.

Here’s how the training and evaluation went:

• Training Time: ~640 seconds
• Evaluation Time: ~22.79 seconds
• Accuracy: 86.80%
• Precision: 90.12%
• Recall: 82.27%
• F1-Score: 86.02%

We can see that the model performed quite well on the IMDb sentiment analysis task. However, the training time was substantial, indicating that FP32, while precise, can be slow.

2.2 Automatic Mixed Precision (AMP)

Next, we leveraged NVIDIA’s automatic mixed precision (AMP) functionality. This allows the model to dynamically switch between FP16 and FP32 based on the task at hand, potentially speeding up the training process without sacrificing too much precision.

• Training Time: ~297.29 seconds
• Evaluation Time: ~15.38 seconds
• Accuracy: 86.10%
• Precision: 89.26%
• Recall: 81.66%
• F1-Score: 85.29%

By incorporating AMP, we cut down the training time by more than half. While the accuracy metrics didn’t change dramatically, the time savings here are significant. This approach is particularly beneficial in real-world scenarios where time and resource efficiency are paramount.

2.3 Manual Mixed Precision

For our final run, we manually applied mixed precision by selectively using FP16 for faster operations and FP32 for stability where needed. While the goal was to balance speed and precision, the results show that this approach did not reduce the overall training time compared to the no mixed precision mode. This unexpected outcome indicates that, in our case, the overhead of managing precision manually may have outweighed the intended speed benefits.

- FP16 (Half Precision) was used during the forward pass, speeding up matrix computations during training and evaluation while reducing memory usage.

- FP32 (Full Precision) was used during backpropagation to ensure accuracy in gradient calculations and weight updates, minimizing the risk of instability.

Updated Results:

- Training Time: ~643.49 seconds

- Evaluation Time: ~22.52 seconds

- Accuracy: 86.55%

- Precision: 85.17%

- Recall: 88.02%

- F1-Score: 86.57%

Takeaway: Despite the manual control over precision, this approach did not outperform the no mixed precision mode in terms of speed. While it allowed for more precise control over FP16 and FP32 usage, the increased complexity and management likely contributed to the longer training times. In this case, manual mixed precision may not provide enough benefit to justify the extra effort, especially when automatic mixed precision can achieve similar accuracy with better efficiency.

Future Work: I plan to dive deeper into this topic, researching the potential scenarios and models where manual mixed precision might offer more tangible benefits. Through this work, I aim to publish a research paper exploring the broader applications, trade-offs, and optimizations possible with manual mixed precision in AI model training.

Section 3: Visualizing the Results

Below are the confusion matrices for each of the training modes. These visualizations give us a deeper understanding of where the model made the most correct and incorrect predictions.

1. No Mixed Precision Confusion Matrix

2. Automatic Mixed Precision Confusion Matrix

3. Manual Mixed Precision Confusion Matrix

As seen in these matrices, the manual and automatic mixed precision modes performed very similarly in terms of misclassifications, while the no mixed precision mode showed a slight improvement in capturing true positives (as seen by fewer false negatives).

• In the No Mixed Precision mode, the model captured more true negatives (924) but had slightly higher false negatives (175). However, it performed better in identifying true positives (812) compared to the other modes.
• Automatic Mixed Precision (AMP) mode had a similar performance but slightly reduced true negatives (916) and slightly more false negatives (181). It showed a marginal decrease in true positive identification (806) compared to the No Mixed Precision mode.
• In the Manual Mixed Precision mode, the model had the most false negatives (151), capturing fewer true negatives (864) compared to both No Mixed Precision and AMP modes. However, it excelled in identifying true positives (867), outperforming the other modes in this metric.

Section 4: Insights and Key Takeaways

Based on the results, we can infer the following:

• Speed Gains with Mixed Precision: Automatic mixed precision (AMP) clearly stands out in terms of reducing training time. It cut down the training time by more than half compared to the no mixed precision mode, with only a slight decrease in performance metrics. Manual mixed precision, on the other hand, took nearly as long as no mixed precision, suggesting that the complexity of manually managing precision did not deliver the intended speed benefits.
• Accuracy vs. Speed Trade-off: The no mixed precision mode showed a slight advantage in precision (90.12%) but at the cost of significantly longer training times (~640 seconds). While AMP provided a much faster training time (~297.29 seconds), its accuracy and precision were only slightly lower. Manual mixed precision, despite taking more time (~643.49 seconds), did not show significant performance gains, indicating that it may not offer enough benefit in this specific case.
• Manual vs. Automatic Mixed Precision: While manual mixed precision provides more granular control, allowing you to fine-tune which parts of the process use FP16 or FP32, the results show that this additional control didn’t translate into performance gains or speed improvements. AMP remains the more effective option for real-world tasks, providing similar results with far less effort and complexity.

Conclusion: The Pit Stop Analysis

As promised in our previous article, we took the theory to the track and tested our models in real-world conditions, much like a racer testing out different strategies to navigate the circuit.

Just as Han says in Tokyo Drift, "The first drifters invented drifting out here in the mountains by feel. They didn’t have anyone to teach them. They were slipping, falling, guessing, risking, and re-doing—until they figured it out." In the world of AI, it’s no different. You have to feel your way through the process, experimenting with different methods, making adjustments, and learning from the results.

Each method—no mixed precision, automatic, and manual—had its own strengths, much like navigating different sections of a race track. FP32 provided high precision but at the expense of training time. Automatic mixed precision significantly sped up training with only a slight dip in accuracy, while manual mixed precision aimed to fine-tune the balance between speed and control but didn't offer the expected time savings.

The results were clear: automatic mixed precision is the standout, offering a substantial performance boost with faster training times and minimal accuracy loss. Manual mixed precision, while giving more control, didn’t outperform in terms of time efficiency. Like the first drifters, we learned by doing—testing, adjusting, and refining our approach along the way.

As we cross the finish line of this experiment, the lesson is clear: hands-on experimentation is essential to mastering AI. Whether you're racing to deploy a model or fine-tuning for maximum precision, mixed precision training is your best tool to strike that perfect balance.