Landscape Learning for Neural Network Inversion

TL;DR

The paper addresses the bottleneck of slow, non-convex optimization in optimization-based inference by learning a latent-space remapping that yields a smoother loss landscape. A mapping network, trained via trajectory rollouts and an experience replay buffer, enables fast, stable gradient-based inversion for generative and discriminative tasks without retraining forward models. Empirical results across GAN inversion, 3D human pose reconstruction, and adversarial defense show substantial speedups and performance gains, including orders-of-magnitude faster convergence and improved robustness, especially on out-of-distribution data. The approach offers a general framework to accelerate OBI with broad applicability, while acknowledging training overhead and potential biases inherent to learned forward models.

Abstract

Many machine learning methods operate by inverting a neural network at inference time, which has become a popular technique for solving inverse problems in computer vision, robotics, and graphics. However, these methods often involve gradient descent through a highly non-convex loss landscape, causing the optimization process to be unstable and slow. We introduce a method that learns a loss landscape where gradient descent is efficient, bringing massive improvement and acceleration to the inversion process. We demonstrate this advantage on a number of methods for both generative and discriminative tasks, including GAN inversion, adversarial defense, and 3D human pose reconstruction.

Figures (9)

• Figure 1: Loss Landscapes Comparison. The loss landscape of optimization-based inference (OBI) is often highly non-convex. We propose to learn a smoother loss landscape through a mapping network to accelerate the optimization procedure. Plotted from real data.
• Figure 2: Method. The left and middle figure show the loss landscape for our latent space and the original latent space, respectively. While walking to the optimal solution in a few steps is hard in $X$ space, it can be done in our learned loss landscapes.
• Figure 3: Landscape Learning. An optimization trajectory $\{{\bf z}_t\}_{t=0}^{5}$ collected is used to train $\theta$. ${\bf z}_i$ that corresponds to a higher $L_i$ will yield a higher gradient when training $\theta$. Optimization over multiple steps along the trajectory causes $\theta$ to learn patterns of trajectories and create a smoother loss landscape.
• Figure 4: Optimization Performance. We visualize the trends of optimization performance compared with the baseline. In GAN Inversion (Left), we evaluate all models on test splits of CelebA-HQ karras2017progressive and LSUN-cat yu2015lsun (OOD) with loss defined in Eq. \ref{['eq:loss_gan']}. Since encoder-based inference doesn't involve optimization, we use a flat line to represent it. We perform 2000 steps of gradient descent for all models except encoder-based models. In 3D Human Pose Reconstruction (Right), we evaluate all models on test splits of GRAB taheri2020grab and PROX hassan2019resolving (OOD) with loss defined in Eq. \ref{['eq:loss_human']}. We perform 200 steps of gradient descent for all models. For each step, we plot the average loss value of test splits.
• Figure 5: Optimization Process for GAN Inversion. Comparing optimization process of our method and the baseline in order to reconstruct the ground truth image. Left shows the results from the baseline where optimization is done in the original input space $X$. Middle shows the results from our method where optimization is done in our space $Z$. Right column contains the ground truth image to each example. Each row corresponds to the same example.