Information Gradient for Nonlinear Gaussian Channel with Applications to Task-Oriented Communication

TL;DR

The paper tackles optimizing information flow through nonlinear Gaussian channels by maximizing mutual information with respect to a parametric front-end. It derives a tractable information gradient based on the marginal score via the score-to-Fisher bridge, enabling DSM-based score learning and VJP-based gradient computation. The framework extends to task-oriented objectives and the information bottleneck, with alternating optimization to manage score updates and parameter learning. Empirical results on linear and nonlinear channels corroborate gradient accuracy and show effective MI and IB optimization, highlighting practical end-to-end optimization for semantic and goal-oriented communication systems.

Abstract

We propose a gradient-based framework for optimizing parametric nonlinear Gaussian channels via mutual information maximization. Leveraging the score-to-Fisher bridge (SFB) methodology, we derive a computationally tractable formula for the information gradient that is the gradient of mutual information with respect to the parameters of the nonlinear front-end. Our formula expresses this gradient in terms of two key components: the score function of the marginal output distribution, which can be learned via denoising score matching (DSM), and the Jacobian of the front-end function, which is handled efficiently using the vector-Jacobian product (VJP) within automatic differentiation frameworks. This enables practical parameter optimization through gradient ascent. Furthermore, we extend this framework to task-oriented scenarios, deriving gradients for both task-specific mutual information, where a task variable depends on the channel input, and the information bottleneck (IB) objective. A key advantage of our approach is that it facilitates end-to-end optimization of the nonlinear front-end without requiring explicit computation on the output distribution. Extensive experimental validation confirms the correctness of our information gradient formula against analytical solutions and demonstrates its effectiveness in optimizing both linear and nonlinear channels toward their objectives.

Figures (6)

• Figure 1: Gradient of mutual information $\partial I / \partial \alpha$ versus $\alpha$. The proposed VJP/score estimator (orange dashed) and the analytical gradient (blue solid) are presented in the figure.
• Figure 2: Mutual information $I(X; Y_t)$ versus channel gain $\alpha$ for the one-dimensional linear-Gaussian channel. The analytical solution (blue solid) and the path-integral estimate from gradient integration (orange dashed).
• Figure 3: Comparisons of the information gradient $\partial I/\partial\alpha$ versus $\alpha$ for the linear vector channel with Gaussian input. The analytic gradient (blue solid), the true-score VJP (orange dashed), and the learned-score VJP (green dash-dotted) are compared.
• Figure 4: Projected gradient ascent with unconditional DSM (fixed step). Iteration vs. mutual information $I(\bm A^{(k)})$ for $n{=}m{=}8$, $\sigma_x^2{=}1$, $t{=}0.5$, $\|\bm A\|_F\le 5.0$. The dashed line is the theoretical optimum $I^\star$ under the Frobenius constraint.
• Figure 5: MI estimate versus iteration for $Y=\tanh(AX)+Z$ ($n{=}12$).

Theorems & Definitions (14)

• Proposition 1: Information gradient
• proof
• Remark 1
• Proposition 2: VJP identity
• proof
• Remark 2
• Remark 3: Reduction to the standard SFB when $T=X$
• Remark 4: Stochastic tasks $T\!\sim\! p(T|X)$
• Proposition 3: Task-oriented information gradient
• proof