Amazon’s Build on Trainium Initiative Brings AI to Academia

Generative AI is now increasingly used to synthesize images, video, text and code. Now Amazon says it will invest US $110 million in university-led research into generative AI to help drive breakthroughs in the field, the company announced.

Generative AI systems such as DALL-E, Midjourney, and Stable Diffusion now regularly conjure photorealistic images. ChatGPT, perhaps the most well-known generative AI chatbot, has passed law school and business school exams, successfully answered interview questions for software-coding jobs, written real estate listings, and developed ad content.

However, developing novel generative AI applications increasingly requires a lot of computing power. Such resources are often well beyond academic researchers, as two researchers noted in a tongue-in-cheek paper that ran in the Proceedings of the IEEE in January.

“AI academic research today is severely bottlenecked by a lack of resources and as such, the academic sector is falling behind quickly,” says Gadi Hutt, senior director of business development at Amazon Web Services’ Annapurna Labs.

Amazon’s Build on Tranium Initiative Brings AI to Academia

As part of the new initiative from Amazon, called Build on Trainium, Amazon Web Services (AWS) has created a computer cluster where researchers can make reservations to access up to 40,000 Trainium chips, Hutt says. AWS developed these processors for high-performance, low-cost deep learning.

“With Build on Trainium, AWS is investing in a new wave of AI research and classes guided by leading AI research in universities that will advance the state of generative AI applications, libraries, and optimizations,” Hutt says.

A new programming interface for Trainium called the Neuron Kernal Interface gives researchers “bare-metal” chip programming capabilities, allowing direct access to the chip’s instruction set and enabling users to build compute kernels for new model operations and performance optimizations, Amazon says.

“That’s great,” says Julian Togelius, an associate professor of computer science and engineering at New York University. “A lot of companies like to share resources, but only if you use a specific proprietary tool chain. Giving researchers low-level access to tune features of the hardware itself sounds awesome.” Togelius was one of the two researchers who authored the Proceedings of the IEEE piece.

The AWS initiative will develop strategic partnerships with universities, including Carnegie Mellon University and the University of California, Berkeley, with more to be announced in coming weeks, Hutt says. In addition, the company will provide grant-level allocations of Trainium credits to the broader research community through multiple rounds of Amazon research awards calling for papers over the next three to five years. Researchers can start applying for these grants immediately, Hutt says. All in all, researchers will get the opportunity to build new AI architectures, machine learning libraries, and performance optimizations for large-scale distributed Trainium clusters.

Next Steps for Build on Trainium

All code developed through the new initiative will be available via open-source machine learning software libraries. “Amazon has been very aggressively reaching out to startups, so now targeting academics makes perfect sense,” Togelius says. “They’re getting more people to do research with their cloud system, getting researchers to found their own companies who may work with AWS.”

Amazon notes the initiative will also help train future AI experts. For instance, those participating in Build on Trainium will receive access to AWS’s technical education and enablement programs for Trainium, Hutt says, in partnership with the Neuron Data Science community led by Amazon’s chip developer Annapurna.

A key factor to the new initiative’s success will be how it chooses who it supports, Togelius says.

“The question is how will Amazon distribute these resources so people do things that big tech could or would not have done?” Togelius says. “If Amazon distributes these resources like it would distribute compute within Amazon, then it hasn’t gained much. You want to focus on what the large tech companies are not investing in for one reason or another—it’s too narrow, it’s experimental, it’s got weird optics—the kind of things that you can’t sell to a venture capitalist or to shareholders. Academics need to double down on the weirdness.”