Mutual Information Optimization via K-Recursion and Automatic Differentiation for Linear Gaussian Wireless Networks

Abstract

We present a differentiable framework for end-to-end mutual information (MI) optimization over linear Gaussian directed acyclic graphs (DAGs). The framework targets network-wide design under global constraints, such as a total transmit power budget, and covers MIMO precoding, amplify-and-forward relays, RIS-aided channels, and branching/merging topologies within a common linear Gaussian model. Its core ingredient is a \emph{K-recursion} that analytically propagates all node-pair covariances along the DAG in topological order, including non-adjacent cross-covariances that are necessary for correctly handling branching and merging paths. The resulting covariances yield a closed-form log-determinant expression for the end-to-end MI as a smooth function of the controllable factors. Complex-valued reverse-mode automatic differentiation on this K-recursion then returns the exact Wirtinger gradient at every controllable factor in a single backward sweep, and projected gradient ascent (PGA) is used to maximize the MI under the global constraints. Because no closed-form gradient expression per topology is required, the same topology-agnostic implementation applies to any linear Gaussian DAG. A single topology-agnostic implementation is applied to four representative DAG classes: single-link MIMO, a diamond DAG, a two-hop AF relay, and input-covariance shaping. The same implementation reaches the classical water-filling optimum in the settings where it is available and yields MI improvements in non-single-link topologies without using topology-specific gradient formulas. A further experiment on a multi-layer Gaussian network (11 nodes, 5 layers) illustrates applicability to nontrivial multi-layer topologies for which no closed-form gradient is available.