Researchers at MIT, the MIT-IBM Watson AI Lab, and other institutions have developed a more efficient computer vision model that dramatically reduces the computational complexity of this task. Their model can perform semantic segmentation accurately in real time on a device with limited hardware resources, such as the onboard computers that enable autonomous vehicles to make split-second decisions.

The latest semantic segmentation models learn the interaction between each pair of pixels in an image directly, so their calculations grow quadratically as the image resolution increases. Therefore, although these models are accurate, they are too slow to process high-resolution images in real time on a peripheral device such as a sensor or a mobile phone.

MIT researchers designed a new building block for semantic segmentation models that achieves the same capabilities as these state-of-the-art models, but with only linear computational complexity and hardware-efficient operations.

The result is a new series of high-resolution computer vision models that operate up to nine times faster than previous models when deployed on a mobile device. More importantly, this new series of models offers accuracy equal to or better than alternative models.

This technique could not only be used to help autonomous vehicles make real-time decisions, but it could also improve the efficiency of other high-resolution computer vision tasks, such as medical image segmentation.

"Although researchers have long used traditional vision transformers, and they produce amazing results, we want people to also pay attention to the efficiency aspect of these models. Our work demonstrates that it is possible to drastically reduce the computation so that real-time image segmentation can be performed locally on a device," says Song Han, associate professor in the Department of Electrical and Computer Engineering (EECS), a member of the MIT-IBM Watson Artificial Intelligence Lab, and lead author of the paper describing the new model.

He is joined on the article by Han Cai, a graduate student at EECS; Junyan Li, a student at Zhejiang University; Muyan Hu, a student at Tsinghua University; and Chuang Gan, a senior research fellow at the MIT-IBM Watson Artificial Intelligence Lab. The research will be presented at the International Conference on Computer Vision.

A simplified solution

Categorizing each pixel in a high-resolution image that can have millions of pixels is a difficult task for a machine learning model. Recently, a powerful new type of model, known as a vision transformer, has been used effectively.

Transformers were originally developed for natural language processing. In that context, they encode each word in a sentence as a token and then generate an attention map, which captures the relationships of each token to all other tokens. This attention map helps the model understand the context when making predictions.

Using the same concept, a vision transformer slices an image into pixel patches and encodes each small patch into a token before generating an attention map. When generating this attention map, the model uses a similarity function that directly learns the interaction between each pair of pixels. In this way, the model develops what is known as a global receptive field, meaning it can access all relevant parts of the image.

Since a high-resolution image can contain millions of pixels, divided into thousands of patches, the attention map quickly becomes enormous. Therefore, the amount of computation increases quadratically as the image resolution increases.

In their new series of models, called EfficientViT, the MIT researchers use a simpler mechanism to build the attention map: they replace the nonlinear similarity function with a linear one. This allows them to rearrange the order of operations to reduce the total number of calculations without changing the functionality or losing the overall receptive field. With their model, the number of calculations required for a prediction grows linearly as image resolution increases.

"But there's no such thing as a free lunch. Linear attention only captures the overall context of the image and misses local information, which worsens accuracy," Han explains.

To compensate for that loss of precision, the researchers included two additional components in their model, each of which adds only a small amount of computation.

One of these elements helps the model capture the interactions of local features, mitigating the weakness of the linear function in extracting local information. The second, a module that enables multiscale learning, helps the model recognize large and small objects.

"The most important thing is to carefully balance performance and efficiency," Cai says.

They designed EfficientViT with a hardware-adapted architecture, making it easier to run on various devices, such as virtual reality headsets or edge computers in autonomous vehicles. Their model could also be applied to other computer vision tasks, such as image classification.

Rationalization of semantic segmentation

When they tested their model on datasets used for semantic segmentation, they found that it performed up to nine times faster on an Nvidia graphics processing unit (GPU) than other popular vision transformer models, with the same or better accuracy.

"Now we can get the best of both worlds and reduce the computing power to make it fast enough to run on mobile devices and in the cloud," Han says.

Based on these results, the researchers want to apply this technique to accelerate generative machine learning models, such as those used to generate new images. They also want to further expand EfficientViT to other vision tasks.

“Model compression and lightweight model design are crucial research topics for achieving efficient AI computing, especially in the context of large-database models. Professor Song Han’s group has shown remarkable progress in compressing and accelerating modern deep learning models, particularly vision transformers,” adds Jay Jackson, global vice president of artificial intelligence and machine learning at Oracle, who was not involved in this research. “Oracle Cloud Infrastructure has been supporting his team in advancing this impactful line of research toward efficient and green AI.”.

###

Written by Adam Zewe, MIT News