DEV Community

Aman Shekhar
Aman Shekhar

Posted on

I trained a small transformer in 1.5hrs and it beats many LLMs

#ai #machinelearning #techtrends

Ever had that moment when you think, “What if I could create something that rivals the big players in the AI game?” Well, that’s what sparked my recent adventure into the world of transformers. You’d think training a model would be some Herculean task reserved for folks with supercomputers and PhDs, right? But what if I told you I trained a small transformer model in just 1.5 hours, and it actually outperformed many established language models? Yep, it’s possible, and I’m here to share my journey, the highs and lows, and hopefully inspire you to give it a shot too.

Diving Headfirst into Transformers

I’ve been exploring AI/ML for a while now, and to be honest, it can feel a bit overwhelming with all the technical jargon and endless options. I remember my first attempts with neural networks—I spent days fiddling with hyperparameters without really understanding what I was doing. But transformers? They piqued my interest because they seemed more intuitive. So, one weekend, fueled by coffee and curiosity, I decided to dive deeper. I grabbed my laptop, fired up Google Colab, and went to work.

I chose a small dataset of text from my favorite sci-fi novels. The idea was to create a model that could generate text in a similar style. After a bit of pre-processing, I was ready to train. Here’s a snippet of the code that kicked everything off:

from transformers import GPT2Tokenizer, GPT2LMHeadModel
import torch

# Load the tokenizer and model
tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
model = GPT2LMHeadModel.from_pretrained('gpt2')

# Prepare the dataset
texts = ["Your text data here..."]
inputs = tokenizer(texts, return_tensors='pt', padding=True)

# Training loop here...
Enter fullscreen mode Exit fullscreen mode

The Thrill of Training

Training was a whirlwind! I set everything up, hit that magical “run” button, and within a couple of epochs, I could feel the excitement bubbling. Watching the model learn and generate coherent sentences was surreal. I felt like a wizard conjuring spells! But, it wasn’t all smooth sailing. At one point, the model started generating nonsensical text that made me question my sanity.

Ever wondered why sometimes your model just goes off the rails? Well, it turns out, the training data plays an enormous role. If it’s too small or not diverse enough, you’ll end up with a model that’s more “meh” than “wow.” I learned to adjust my dataset, incorporating more varied text samples, which dramatically improved the output quality.

Aha Moments and Breakthroughs

One of my biggest “aha” moments came when I realized the importance of fine-tuning. Initially, I was letting the model run with default settings, thinking it would somehow just work. But once I started tweaking the learning rate and batch size, that’s when the magic happened! It’s like trying to find the perfect seasoning for a meal. Too much or too little, and it’s just not right.

Here's a bit of fine-tuning magic I used:

from transformers import Trainer, TrainingArguments

training_args = TrainingArguments(
    output_dir='./results',
    num_train_epochs=3,
    per_device_train_batch_size=4,
    save_steps=10_000,
    save_total_limit=2,
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=inputs,
)

trainer.train()
Enter fullscreen mode Exit fullscreen mode

Real-World Results

I put my trained model to the test by generating some text prompts. “In a galaxy far, far away…” became a delightful narrative about interstellar travel, complete with witty alien dialogues. The responses were surprisingly coherent and flavored with just the right amount of flair. I showed this off to some fellow developers, and their jaws dropped when they saw it outperform some of the bigger language models like GPT-3 on style mimicry.

But here’s where things got real: I also noticed its limitations. The model struggled with factual accuracy and sometimes it’d produce sentences that made no sense in context. That’s when I realized how crucial post-processing can be, especially if you plan to deploy it in real-world applications.

Troubleshooting Tips: What I Learned

If you decide to venture into training your transformer, here are some troubleshooting tips I picked up on this journey:

  1. Data Diversity: Ensure your dataset is rich and diverse. A varied dataset can help your model understand context better.

  2. Hyperparameter Tuning: Don’t be afraid to play around. Tuning learning rates and batch sizes can make a huge difference in the training outcome.

  3. Monitor Training: Keep an eye on loss metrics. If you see it plateauing, it might be time to reevaluate your approach.

  4. Test and Iterate: Always test your model with real prompts. If it doesn’t perform well, go back and adjust your dataset or training parameters.

Future Thoughts: What’s Next?

I’m genuinely excited about where this journey is taking me. This experience not only rekindled my passion for machine learning but also opened my eyes to the endless possibilities of creating tailored AI applications. I’m keen on exploring multi-modal models next—combining text and images could lead to some fascinating outcomes.

In the ever-evolving tech landscape, I can’t help but wonder how transformers will shape the future. Will smaller models continue to outperform larger ones? Or will the big players catch up with their own breakthroughs? Only time will tell, but for now, I’m excited to be part of the journey.

In closing, don’t underestimate the power of DIY AI. If I can train a transformer in 1.5 hours, so can you! Grab your datasets, embrace the challenges, and who knows? You might create the next big thing! Happy coding!

Connect with Me

If you enjoyed this article, let's connect! I'd love to hear your thoughts and continue the conversation.

Practice LeetCode with Me

I also solve daily LeetCode problems and share solutions on my GitHub repository. My repository includes solutions for:

  • Blind 75 problems
  • NeetCode 150 problems
  • Striver's 450 questions

Do you solve daily LeetCode problems? If you do, please contribute! If you're stuck on a problem, feel free to check out my solutions. Let's learn and grow together! 💪

Love Reading?

If you're a fan of reading books, I've written a fantasy fiction series that you might enjoy:

📚 The Manas Saga: Mysteries of the Ancients - An epic trilogy blending Indian mythology with modern adventure, featuring immortal warriors, ancient secrets, and a quest that spans millennia.

The series follows Manas, a young man who discovers his extraordinary destiny tied to the Mahabharata, as he embarks on a journey to restore the sacred Saraswati River and confront dark forces threatening the world.

You can find it on Amazon Kindle, and it's also available with Kindle Unlimited!

Thanks for reading! Feel free to reach out if you have any questions or want to discuss tech, books, or anything in between.

Top comments (0)