Geeky Gadgets

The Latest Technology News



30M Parameter AI Running on an ESP32 Microcontroller

By

Text output from a custom AI chatbot running at 1.4 tokens per second.

Deploying a 30-million-parameter AI model on an ESP32 microcontroller is no small feat, but Marina in the video below demonstrates how it’s possible by combining careful design and optimization techniques. The project involved building a decoder-only transformer model with 12 layers and a 6,000-token vocabulary, tailored to fit the ESP32’s limited 32 MB of flash memory and 16 MB of PS RAM. By employing methods like quantization to reduce model weights and using the SentencePiece tokenizer for efficient text processing, the model was adapted to handle conversational tasks while staying within the hardware’s constraints.

In this deep dive, you’ll gain insight into the key steps involved, from training the model on 300 million tokens to deploying it on resource-constrained hardware. Discover how techniques like gradient clipping and dual-core processing were used to optimize performance and explore alternative approaches such as using external storage or cloud-based solutions. Whether you’re interested in building AI for IoT devices or learning about the challenges of running AI on microcontrollers, this guide provides a comprehensive look at the process and its practical applications.

AI Model Design and Architecture

TL;DR Key Takeaways :

  • Developing AI on resource-constrained devices like the ESP32 requires designing and optimizing a 30-million-parameter model to fit within 32 MB of flash memory and 16 MB of PS RAM.
  • A decoder-only transformer model with 12 layers and a 6,000-token vocabulary was used, optimized for conversational tasks through techniques like quantization and streamlined architecture.
  • The model was trained on 300 million tokens using gradient-based optimization methods, achieving efficient processing and generalization despite hardware limitations.
  • Deployment challenges, such as limited processing speed (1.4 tokens per second), were addressed through dual-core processing and external storage solutions for smaller devices.
  • Local AI deployment offers privacy and independence, with applications in IoT, offline agents and monitoring systems, while alternative approaches like cloud-based services provide scalability and flexibility.

Creating a chatbot capable of understanding and generating coherent text requires a well-thought-out model architecture. For this project, a decoder-only transformer model was selected due to its balance between computational efficiency and performance. The architecture included:

  • 12 layers to ensure sufficient processing depth
  • A vocabulary of 6,000 tokens for language comprehension

This design was tailored to meet the constraints of the ESP32 while maintaining reasonable performance. Tokenization, a critical preprocessing step, was performed using the SentencePiece tool, which ensured compatibility with the model’s vocabulary size and structure. This allowed the model to process text effectively while adhering to hardware limitations.

The model was optimized for conversational understanding, focusing on sentence structure and contextual relevance. By carefully limiting the vocabulary size and adjusting the number of layers, computational demands were reduced without compromising the model’s ability to handle conversational tasks.

Training Process

Training a transformer model requires meticulous planning and execution. The model was trained on publicly available datasets, processing approximately 300 million tokens. Tokenization converted raw text into a machine-readable format, allowing efficient processing. Batch training was employed to handle the large dataset, while cross-entropy loss was used to evaluate the model’s performance. To stabilize training and improve convergence, gradient-based optimization techniques were applied, including:

  • Backpropagation to update model weights effectively
  • Gradient clipping to prevent exploding gradients during training

As training progressed, validation loss decreased significantly, indicating improved accuracy and generalization. These results demonstrated the model’s ability to handle conversational tasks effectively, even when operating within the constraints of a resource-limited environment.

Hardware Constraints and Optimization

The ESP32 microcontroller’s limited resources posed significant challenges. With only 32 MB of flash memory and 16 MB of PS RAM, fitting a 30-million-parameter model required aggressive optimization techniques. Key strategies included:

  • Quantization: Reducing model weights to 4-bit and 16-bit representations to minimize memory usage
  • Streamlining architecture: Adjusting the model’s structure to fit within the hardware’s constraints

Memory allocation and weight storage were carefully managed to ensure the model operated efficiently within the ESP32’s limitations. These optimizations allowed the AI to function effectively despite the device’s constrained resources, demonstrating the feasibility of deploying sophisticated models on microcontrollers.

Deep dive into the latest in ESP32 by exploring our other resources and articles.

Deployment Challenges and Solutions

Running the trained model on the ESP32 required innovative solutions to address performance bottlenecks. The microcontroller achieved a processing speed of 1.4 tokens per second. While modest, this performance demonstrated the feasibility of deploying AI on constrained hardware. To enhance efficiency, dual-core processing was utilized, allowing parallel operations and improving overall performance.

For even smaller devices, such as the Xiao ESP32S3, external storage via a microSD card was employed to load larger weights. Although this approach slowed processing speeds, it allowed the model to function on devices with even tighter constraints. By combining dual-core processing and external storage, you can overcome the limitations of smaller boards while maintaining functionality.

Alternative Approaches

While deploying AI models locally on constrained devices is an impressive achievement, alternative approaches may offer greater flexibility and scalability. These include:

  • Cloud-based AI services such as Google Gemini for natural language processing tasks
  • Local deployment of larger models like LLaMA on more capable hardware

Each approach has its trade-offs. Local deployment ensures data privacy and independence from external servers but requires significant optimization to fit within hardware constraints. Conversely, cloud-based solutions offer scalability and ease of use but may involve recurring costs and reliance on internet connectivity. Choosing the right approach depends on your specific use case, priorities and available resources.

Practical Applications and Future Directions

This project highlights the potential of deploying AI on resource-constrained devices like the ESP32. By using techniques such as quantization, gradient-based optimization and dual-core processing, you can overcome hardware limitations and achieve functional AI on microcontrollers. This opens up possibilities for a wide range of applications, including:

  • Edge computing for IoT devices
  • Offline conversational agents
  • AI-powered monitoring systems

While alternative approaches like cloud-based services or more capable hardware exist, running models locally provides unique advantages in terms of privacy, control and independence. Plans are underway to share detailed instructions and resources on GitHub, allowing others to replicate or adapt this process for their own projects. Whether you’re exploring AI on microcontrollers or considering alternative deployment methods, this project underscores the versatility and potential of AI in resource-constrained environments.

Media Credit: Marina

Filed Under: DIY Projects, Guides



Disclosure: Some of our articles include affiliate links. If you buy something through one of these links, Geeky Gadgets may earn an affiliate commission. Learn about our Disclosure Policy.

Geeky Gadgets