Parameter Estimation of Mutual Information Maximized Channels

Technical Summary: Parameter Estimation of Mutual Information Maximized Channels

Problem Statement
This paper addresses the challenge of estimating the parameters of a discrete memoryless channel, denoted as $p(y | x; \theta)$, when the transmitter operates under a specific constraint: it selects an input distribution $\pi$ to maximize the mutual information $I(X; Y)$ given the true channel parameter $\theta^*$. In this setting, the receiver observes only an independent and identically distributed (i.i.d.) sequence of channel outputs $y_{1:T}$ and possesses no knowledge of the channel parameters $\theta$ or the input distribution $\pi$.

The core difficulty arises because the input distribution $\pi$ is not an independent variable but is functionally dependent on $\theta$ via the capacity-achieving condition (i.e., $\pi^*(\theta) = \arg\max_{\pi} I(X; Y | \theta)$). Standard Maximum Likelihood Estimation (MLE) approaches that treat $\pi$ and $\theta$ as free, unconstrained parameters often fail to recover the true $\theta^*$ and $\pi^*$. This failure occurs because different pairs of $(\theta, \pi)$ can induce identical marginal output distributions, rendering the problem non-identifiable without enforcing the mutual information maximization constraint.

Methodology
The authors propose two efficient algorithms to compute the constrained maximum likelihood estimator $\hat{\theta}_{MLE}$, both leveraging the Blahut–Arimoto (BA) optimality conditions. The problem is formulated as finding $\theta$ that maximizes the marginal log-likelihood $L_T(\theta, \pi(\theta))$, where $\pi(\theta)$ satisfies the BA fixed-point condition $b(\pi, \theta) = \pi$.

1. Bilevel Fixed-Point Method:
   This approach treats the estimation as a bilevel optimization problem. The upper-level objective is the log-likelihood, and the lower-level problem is the maximization of mutual information (solved via BA). The algorithm iteratively:

   • Solves the lower-level problem to find $\pi$ for a fixed $\theta$.
   • Computes the gradient of the upper-level objective with respect to $\theta$ using implicit differentiation. This involves calculating the Jacobian of the BA fixed-point mapping to determine how $\pi$ changes with $\theta$ ($\nabla_\theta \pi(\theta)$).
   • Updates $\theta$ using a gradient step.

2. Augmented Lagrangian (AL) Method:
   This method reformulates the problem by relaxing the strict BA fixed-point constraint. Instead of solving the inner BA problem to convergence at every step, the AL method incorporates the BA residual $R(\theta, \pi) = b(\pi, \theta) - \pi$ as a penalty term in an augmented Lagrangian function. The algorithm alternates between:

   • Performing a few inner iterations of the BA update to refine $\pi$.
   • Updating the dual variables (multipliers) and the penalty parameter.
   • Updating the primal variable $\theta$ using the gradient of the augmented Lagrangian.
     This approach avoids the computationally expensive implicit differentiation step required by the bilevel method.

Key Contributions

• Problem Formulation: The paper formally defines the unsupervised learning setting where channel parameters and input distributions must be jointly estimated under the constraint that the input distribution maximizes mutual information.
• Identifiability Analysis: The authors provide a theoretical analysis of identifiability and sample complexity. They demonstrate that without the mutual information constraint, the problem is often non-identifiable (e.g., in Binary Symmetric Channels with capacity-achieving inputs). Conversely, they show that for certain channels like the Binary Erasure Channel (BEC), the parameter is identifiable even under the capacity-achieving constraint.
• Algorithmic Development: The proposal of two distinct algorithms (Bilevel and AL) that explicitly enforce the mutual information maximization constraint, contrasting them with naive Joint-ML approaches that ignore this constraint.
• Empirical Validation: Numerical experiments on a discrete memoryless Gaussian-like channel demonstrate that both proposed algorithms successfully recover the true parameter $\theta^*$ and input distribution $\pi^*$, whereas the naive Joint-ML approach converges to incorrect parameters that yield high likelihood but do not satisfy the system constraints.

Results
In simulations with $N=10$ inputs, $M=50$ outputs, and $T=200,000$ observations:

• Accuracy: Both the Bilevel and AL methods converged to the true parameter $\theta^* = 0.7$ with high precision (median absolute errors on the order of $10^{-4}$). The estimated input distributions $\pi$ were recovered with high fidelity (KL divergences $\approx 10^{-7}$).
• Failure of Naive Approach: The Joint-ML method, which does not enforce the mutual information constraint, converged to a different parameter set that maximized the likelihood but was far from the true $\theta^*$, confirming the non-identifiability issue in the unconstrained setting.
• Efficiency: While both proposed methods achieved similar accuracy, the Augmented Lagrangian method demonstrated superior computational efficiency. It required approximately 33% fewer full Blahut–Arimoto solves and resulted in a 7.7% reduction in wall-clock time compared to the Bilevel method, attributed to avoiding the implicit differentiation step and requiring fewer inner BA iterations.

Significance and Claims
The paper claims that enforcing the mutual information maximization constraint is critical for the identifiability of channel parameters in systems where the transmitter optimizes its input distribution. The proposed Augmented Lagrangian method is highlighted as a practical solution that balances theoretical rigor with computational efficiency, offering a viable alternative to more complex bilevel optimization techniques. The authors note that the method is simple to implement and compatible with existing optimization frameworks, though it relies on the uniqueness of the local BA fixed point and requires tuning of penalty updates. The work is supported by the National Science Foundation (Grant No. CCF:SHF:2417898).