MIT CSAIL Uses Deep Generative Model StyleGAN2 to Deliver SOTA Image Reconstruction Results

A group of researchers from MIT Computer Science & Artificial Intelligence Laboratory (CSAIL) have proposed a simple framework for performing different image reconstruction tasks using the state-of-the-art generative model StyleGAN2.

It’s common for machine learning researchers to train models in a supervised setting for solving downstream prediction and image reconstruction tasks. For example, in the task of super-resolution, which aims to obtain high-resolution output images from low-resolution versions, classical methods train models on pairs of low-resolution and high-resolution images.

However, such end-to-end methods can also require re-training whenever there is a distribution shift in the inputs or relevant latent variables. Distribution shifts can easily occur for example in the input x-ray images collected from a hospital if the hospital’s medical scanners are upgraded, or as the patients contributing the images age due to improved healthcare.

Given the prohibitively high computation resources required to re-train end-to-end approaches when distribution shifts occur, how else might researchers build ML models that are both easy to train and robust to distribution shifts?

“To this end, it is of crucial importance to introduce causal inductive biases in the model,” the CSAIL researchers explain, pointing to causal modelling as a solution. The team says adding causal inductive biases can ensure the independence of mechanisms within the framework, so that upstream changes will not necessitate retraining the downstream models.

The team says previous supervised learning image reconstruction approaches that starting with the corrupted image and generated a restored image worked in an “anti-causal direction,” regularizing the inversion process using smoothness priors or sparsity, and this tended to result in blurry images. Their proposed novel Bayesian Reconstruction through Generative Models (BRGM) approach is causal, and closely follows the data generating process.

The researchers leveraged StyleGAN2 for building robust image priors. To restore a noisy image, the pretrained SOTA StyleGAN2 generator model generates a potential clean reconstruction. A corrupted image is then generated via a corruption model processing the potential clean reconstruction. The simulated corrupted image is compared to the noisy image with a distance metric so the loss function can be minimized to avoid blurring in the final output image. In this way, the approach accounts for distribution shifts in either the image dataset or the corruption processing without requiring any model retraining.

The team examined their BRGM approach on three datasets, including challenging medical datasets, comprising:

• 60,000 images of human faces from the Flickr Faces High Quality (FFHQ) dataset
• ~240,000 chest X-ray images from MIMIC III
• 7,329 brain MRI slices from a collection of 5 neuroimaging datasets

In both qualitative and quantitative evaluations, BRGM outperformed SOTA methods PULSE, ESRGAN, SRFBN on super-resolution and inpainting reconstruction tasks.

The paper Bayesian Image Reconstruction using Deep Generative Models is on arXiv.

---

Reporter: Fangyu Cai | Editor: Michael Sarazen

---

Synced Report | A Survey of China’s Artificial Intelligence Solutions in Response to the COVID-19 Pandemic — 87 Case Studies from 700+ AI Vendors

This report offers a look at how China has leveraged artificial intelligence technologies in the battle against COVID-19. It is also available on Amazon Kindle. Along with this report, we also introduced a database covering additional 1428 artificial intelligence solutions from 12 pandemic scenarios.

Click here to find more reports from us.

---

We know you don’t want to miss any news or research breakthroughs. Subscribe to our popular newsletter Synced Global AI Weekly to get weekly AI updates.