Stochastic Vector Approximate Message Passing with applications to phase retrieval

Phase retrieval refers to the problem of recovering a high-dimensional vector $\boldsymbol{x} \in \mathbb{C}^N$ from the magnitude of its linear transform $\boldsymbol{z} = A \boldsymbol{x}$, observed through a noisy channel. To improve the ill-posed nature of the inverse problem, it is a common practice to observe the magnitude of linear measurements $\boldsymbol{z}^{(1)} = A^{(1)} \boldsymbol{x},..., \boldsymbol{z}^{(L)} = A^{(L)}\boldsymbol{x}$ using multiple sensing matrices $A^{(1)},..., A^{(L)}$, with ptychographic imaging being a remarkable example of such strategies. Inspired by existing algorithms for ptychographic reconstruction, we introduce stochasticity to Vector Approximate Message Passing (VAMP), a computationally efficient algorithm applicable to a wide range of Bayesian inverse problems. By testing our approach in the setup of phase retrieval, we show the superior convergence speed of the proposed algorithm.
Comment: We found that damping scheme proposed in S. Sarker et al. (2021) substantially enhances the VAMP algorithm, which changed the result shown in Fig. 4. We also added the link to our code for numerical experiment. Some typos are also corrected