\subsection{Lemmas and Proofs}
\label{sec:omitted_proofs}

 \setcounter{lemma}{0}
 \setcounter{theorem}{0}

\begin{lemma}
Let $X\sim p_X(x)$ and $Y\sim p_{Y}(y|x)$ be the channel input and output, respectively. Let $D(\cdot)$ be a positive function. If $\mathcal{J}_{\gamma}(D)$, $\gamma>0$, is a value function defined as 
\begin{align}
\mathcal{J}_{\gamma}(D) = \; & \gamma \cdot \mathbb{E}_{(x,y) \sim p_{XY}(x,y)}\biggl[\log \biggl(D\bigl(x,y\bigr)\biggr)\biggr] \nonumber \\ 
& + \mathbb{E}_{(x,y) \sim p_{X}(x)p_{Y}(y)}\biggl[- D^{\gamma}\bigl(x,y\bigr)\biggr],
\end{align}
then
\begin{equation}
I(X;Y) \geq \tilde{I}_{\gamma DIME}(X;Y) = \mathcal{J}_{\gamma}(D^*)+1,
\end{equation}
where
\begin{equation}
D^*(x,y) = \biggl(\frac{p_{XY}(x,y)}{p_{X}(x)\cdot p_Y(y)}\biggr)^{1/\gamma} = \arg \max_D \mathcal{J}_{\gamma}(D).
\end{equation}
\end{lemma}

\begin{proof}
Consider a scaled generator $f(u) = \frac{u}{\gamma}\log u$ and for simplicity of notation, denote $p_{XY}$ and $p_Xp_Y$ with $p$ and $q$, respectively. Then
\begin{equation}
D_{KL}(p||q) = \gamma \int_{x}{q(x) \frac{p(x)}{\gamma q(x)}\log\biggl(\frac{p(x)}{q(x)}\biggr) \diff x},
\end{equation}
with conjugate $f^*(t)$, with $t \in \mathbb{R}$ given by
\begin{equation}
f^*(t) = \frac{e^{\gamma t -1}}{\gamma}.
\end{equation}
Substituting in \eqref{eq:f_bound} yields to
\begin{equation}
D_{KL}(p||q) \geq \sup_{T\in \mathbb{R}} \biggl\{ \gamma \mathbb{E}_{x \sim p(x)} \bigl[T(x)\bigr]-\mathbb{E}_{x\sim q(x)}\bigl[e^{\gamma T(x)-1} \bigr]\biggr\}.
\end{equation}
Using \eqref{eq:optimal_ratio} it is easy to verify that the optimal value of $T$ is the log-likelihood ratio rather than the density ratio. Indeed,
\begin{equation}
T^*(x) = \frac{1}{\gamma}\log\biggl(\frac{p(x)}{q(x)}\biggr) + \frac{1}{\gamma}.
\end{equation}
Finally, with the change of variable $T(x)=\log(D(x))+1/{\gamma}$, the optimal discriminator has form
\begin{equation}
D^*(x) = \biggl(\frac{p(x)}{q(x)}\biggr)^{1/\gamma}
\end{equation}
and
\small
\begin{equation}
D_{KL}(p||q) \geq 1 + \sup_{D\in \mathbb{R}_+} \gamma \mathbb{E}_{x \sim p(x)} \bigl[\log \bigl( D(x) \bigr) \bigr] \nonumber \\
 -\mathbb{E}_{x\sim q(x)}\bigl[D^{\gamma}(x) \bigr].
\end{equation}
Denoting the two terms inside the supremum with $J_{\gamma}(D)$ and replacing $p$ and $q$ with the joint and marginal distributions, concludes the proof.
\end{proof}

\section{Introduction}
\label{sec:introduction}
In digital communication systems, data detection and decoding is a fundamental task implemented at the receiver. The maximum a-posteriori (MAP) decoding approach is the optimal one to minimize the symbol or sequence error probability \cite{Proakis2001,Bahl1974}. If the channel model is known, the MAP decoder can be realized according to known algorithms. For instance, for a linear time invariant channel (LTI) with additive Gaussian noise, the decoder comprises a first stage where the channel impulse response is estimated. Then, a max-log-MAP sequence estimator algorithm is implemented. The decoding metric is essentially related to the Euclidean distance between the received and hypothesized transmitted message filtered by the channel impulse response \cite{Proakis2001}. The task becomes more intriguing for channels that cannot be easily modeled or that are unknown. In such a case, a learning strategy turns out to be an attractive solution and it can be enabled by recent advances in machine learning for communications \cite{Oshea2017, Nachmani2018, Dorner2018}.

In this paper, we derive a neural architecture for optimal decoding in an unknown channel. We propose to use the mutual information (MI) of the input-output channel pair as decoding metric. The computation of the MI is a formidable task and it is unknown in many instances. Therefore, the MI has to be learned, which can be done in principle with neural architectures \cite{Mine2018, Wunder2019, Letizia2021, Song2020, LetiziaNIPS}. 
Additionally, in the problem at hand, the decoder has a more specific task: it has to learn and minimize the a-posteriori information $-\log_2{(p_{X|Y}(\mathbf{x}|\mathbf{y}))}$ of the transmitted codeword $\mathbf{x}$ given the observed/received signal vector $\mathbf{y}$. The direct estimation of such a-posteriori information becomes then the objective. Such an estimation is possible by directly learning the conditional probability density function (pdf) $p_{X|Y}(\mathbf{x}|\mathbf{y})$ with a new estimator that exploits a discriminative model referred to as mutual information neural decoder (MIND). MIND allows not only to perform the decoding task, but also to return an estimate of the channel achievable rate and of the decoding error probability. 

This paper is organized as follows. Section \ref{sec:mid} presents the decoding approach principle. Section \ref{sec:ndec} describes the core idea about MIND. Section \ref{sec:implementation} discusses two practical neural architectures. Section \ref{sec:estimation} explains how MIND can be exploited to estimate both the achievable rate and the decoding error probability. Numerical results for a number of scenarios are reported in Section \ref{sec:results}. Finally, the conclusions are drawn.


\section{Mutual Information Decoding Principle}
\label{sec:mid}

The considered communication system comprises an encoder, a channel, and a decoder. The encoder maps $M=2^k$ messages of $k$ bits into $M$ encoded messages of $n$ (not necessarily binary) symbols. The uncoded message is denoted with the vector $\mathbf{b}$, while the coded message with the vector $\mathbf{x}$ belonging to the set $\mathcal{A}_x$. Thus, the code rate is $R=k/n$ bits per channel use. The channel conveys the coded message to the receiver. Without loss of generality, we can write that the received message is $\mathbf{y}=H(\mathbf{x},\mathbf{h},\mathbf{n})$, where the channel transfer function $H$ implicitly models the channel internal state $\mathbf{h}$ including stochastic variables, e.g. noise, $\mathbf{n}$. For instance, for an LTI channel with additive noise, we obtain $\mathbf{y}=\mathbf{h}*\mathbf{x}+\mathbf{n}$, where $*$ denotes convolution.

The aim of the decoder is to retrieve the transmitted message so that to minimize the decoding error, or equivalently to maximize the received information for each transmitted message. Indeed, if we consider the mutual information of the pair $(\mathbf{y},\mathbf{x})$, with joint pdf $p_{XY}(\mathbf{x},\mathbf{y})$, and marginals $p_X(\mathbf{x})$ and $p_Y(\mathbf{y})$, this can be written as

\begin{equation}
i(\mathbf{y};\mathbf{x})=i(\mathbf{x})-i(\mathbf{x}|\mathbf{y})
\end{equation}
where $i(\mathbf{x})$ is the information of the transmitted message, $i(\mathbf{x}|\mathbf{y})$ is the a-posteriori information (residual uncertainty observing $\mathbf{y}$) \cite{Gallager1968}. It follows that for each transmitted message, the mutual information is maximized when the a-posteriori information is minimized. The a-posteriori information is given by
\begin{equation}
\label{eq:PI}
i(\mathbf{x}|\mathbf{y})=-\log_2 p_{X|Y}(\mathbf{x}|\mathbf{y}) = -\log_2{\frac{p_{XY}(\mathbf{x},\mathbf{y})}{p_{Y}(\mathbf{y})}}
\end{equation}
so that the decoding criterion becomes
\begin{equation}
\mathbf{\hat{x}}=\argmin_{\mathbf{x}\in \mathcal{A}_x}{i(\mathbf{x}|\mathbf{y})} = \argmax_{\mathbf{x}\in \mathcal{A}_x}{\log_2 p_{X|Y}(\mathbf{x}|\mathbf{y})}
\end{equation}
which corresponds to the log-MAP decoding principle \cite{Proakis2001,Bahl1974}.

To compute the a-posteriori information, we need to know the a-posteriori probability $p_{X|Y}(\mathbf{x}|\mathbf{y})$. If the channel transfer function can be modeled, then the decoder is realized according to known approaches. In fact, the a-posteriori probability can be explicitly calculated as $p_{X|Y}(\mathbf{x}|\mathbf{y})=p_{Y|X}(\mathbf{y}|\mathbf{x})p_{X}(\mathbf{x})$, i.e., from the conditional channel pdf and the coded message a-priori probability. However, for unknown channels and sources, we have to resort to a learning strategy as explained in the next section.

\section{Discriminative Formulation of MIND}
\label{sec:ndec}

In the following, we show that it is possible to design a parametric discriminator whose output leads to the estimation of the a-posteriori information. The optimal discriminator is found by maximizing an appropriately defined loss function.
Before describing the proposed solution, it should be noted that in generative adversarial networks (GANs) \cite{Goodfellow2014} a discriminator $D$ is used to distinguish if a sample is an original one coming from real data rather than a fake one generated by the generator $G$. The adversarial training pushes the discriminator $D(\mathbf{a})$ towards the optimum value \cite{Goodfellow2014} 
\begin{equation}
D^*(\mathbf{a}) = \frac{p_{data}(\mathbf{a})}{p_{data}(\mathbf{a})+p_{gen}(\mathbf{a})}
\end{equation} 
where $p_{data}(\mathbf{a})$ is the collected data pdf and $p_{gen}(\mathbf{a})$ is the generated data pdf.
Hence, the output of the optimal discriminator can be used to estimate the probability density ratio 
\begin{equation}
\frac{p_{gen}(\mathbf{a})}{p_{data}(\mathbf{a})} = \frac{1-D^*(\mathbf{a})}{D^*(\mathbf{a})}.
\end{equation}
 
\begin{comment}
The training procedure for $G$ is to maximize the probability of $D$ making a mistake. GANs implement a min-max two-player game such that the Nash equilibrium is reached. Given an input noise vector $\mathbf{z}$ with joint pdf $p_{noise}(\mathbf{z})$, the mapping to the data space is achieved through $G(\mathbf{z};\theta_{gen})$. It has been proved that the generator implicitly learns the true distribution when the following value function $V(G,D)$ is used
\begin{align}
\label{eq:GAN}
V(G,D) &=  \mathbb{E}_{\mathbf{x} \sim p_{data}(\mathbf{x})}[\log D(\mathbf{x})] \\ \nonumber
&+ \mathbb{E}_{\mathbf{z} \sim p_{noise}(\mathbf{z})}[\log(1-D(G(\mathbf{z})))],
\end{align}
an the equilibrium is reached so that $p_{gen} = p_{data}$.

The adversarial training pushes the discriminator $D(\mathbf{x})$ towards the optimum value (see \cite{Goodfellow2014})
\begin{equation}
D^*(\mathbf{x}) = \frac{p_{data}(\mathbf{x})}{p_{data}(\mathbf{x})+p_{gen}(\mathbf{x})},
\end{equation}
thus, we notice that the output of the optimal discriminator indirectly estimates the probability density ratio $p_{gen}/p_{data}$ since we obtain
\begin{equation}
\frac{p_{gen}(\mathbf{x})}{p_{data}(\mathbf{x})} = \frac{1-D^*(\mathbf{x})}{D^*(\mathbf{x})}.
\end{equation}
\end{comment}
Now, the idea is to design a discriminator that estimates the probability density ratio $p_{XY}(\mathbf{x},\mathbf{y})/p_{Y}(\mathbf{y})$ and consequently the a-posteriori information as shown in \eqref{eq:PI}. 

The following Lemma provides an a-posteriori information estimator that exploits a discriminator trained to minimize a certain loss function. At the equilibrium, a transformation of the discriminator output yields the searched density ratio.   

\begin{comment}
take only the optimal discriminator in the formulation of \eqref{eq:GAN} and use it as indirect a-posteriori probability and mutual information estimator. More in detail, since $p_{X|Y}(\mathbf{x}|\mathbf{y})=p_{XY}(\mathbf{x},\mathbf{y})/p_{Y}(\mathbf{y})$, then $i(\mathbf{x}|\mathbf{y})$ can be indirectly estimated using the optimum discriminator $D^*$ when $p_{data}\equiv  p_Y(\mathbf{y})$ and $p_{gen} \equiv p_{XY}(\mathbf{x}, \mathbf{y})$. 
\end{comment}

\begin{lemma}
\label{lemma:Lemma1}
Let $\mathbf{x}\sim p_X(\mathbf{x})$ and $\mathbf{y}\sim p_Y(\mathbf{y})$ be the channel input and output, respectively. Let $H(\cdot)$ be the channel transfer function, in general stochastic, such that $\mathbf{y} = H(\mathbf{x}, \mathbf{h}, \mathbf{n} )$, with $\mathbf{h}$ and $\mathbf{n}$ being internal state and noise variables, respectively. If $\mathcal{J}_{MIND}(D)$ is the value function defined as
\begin{align}
\label{eq:discriminator_function}
\mathcal{J}_{MIND}(D) &= \; \mathbb{E}_{(\mathbf{x},\mathbf{y}) \sim p_{X}(\mathbf{x})p_{Y}(\mathbf{y})}\biggl[\log \biggl(D\bigl(K,\mathbf{y}\bigr)\biggr)\biggr] \nonumber \\ 
& + \mathbb{E}_{(\mathbf{x},\mathbf{y}) \sim p_{XY}(\mathbf{x},\mathbf{y})}\biggl[\log \biggl(1-D\bigl(\mathbf{x},\mathbf{y}\bigr)\biggr)\biggr],
\end{align}
where $K$ is a constant, then
\begin{equation}
\label{eq:optimal_discriminator_1}
D^*(\mathbf{x},\mathbf{y}) = \frac{p_{Y}(\mathbf{y})}{p_{Y}(\mathbf{y}) + p_{XY}(\mathbf{x},\mathbf{y})}= \arg \max_D \mathcal{J}_{MIND}(D),
\end{equation}
and
\begin{equation}
\label{eq:i-DMIE}
i(\mathbf{x}|\mathbf{y}) = -\log \biggl(\frac{1-D^*(\mathbf{x},\mathbf{y})}{D^*(\mathbf{x},\mathbf{y})}\biggr).
\end{equation}
\end{lemma}

\begin{proof}
From the hypothesis of the Lemma, the value function can be rewritten as
\begin{align}
\mathcal{J}_{MIND}(D) &= \int_{\mathcal{T}_x} \int_{\mathcal{T}_y}\biggl[p_{XY}(K(\mathbf{x}),\mathbf{y}) \log \biggl(D(\mathbf{x},\mathbf{y})\biggr) \nonumber \\ 
&+ p_{XY}(\mathbf{x},\mathbf{y}) \log \biggl(1-D(\mathbf{x},\mathbf{y})\biggr)\biggr] \diff \mathbf{x} \diff \mathbf{y}
\end{align}
where $\mathcal{T} = \mathcal{T}_x \times \mathcal{T}_y $ is the domain of definition of the joint density function and $K(\mathbf{x})$ is a function that maps any $\mathbf{x}$ into the constant value $K$ with probability one.

To maximize $\mathcal{J}_{MIND}(D)$, a necessary and sufficient condition requires to take the derivative of the integrand with respect to $D$ and to set it to $0$, yielding the following equation in $D$
\begin{equation}
\frac{p_{XY}(K(\mathbf{x}),\mathbf{y})}{D(\mathbf{x},\mathbf{y})} -\frac{p_{XY}(\mathbf{x},\mathbf{y})}{1-D(\mathbf{x},\mathbf{y})} =0.
\end{equation}
The solution of the above equation yields the optimum discriminator $D^*(\mathbf{x},\mathbf{y})$ since: $p_{XY}(K(\mathbf{x}),\mathbf{y})=p_Y(\mathbf{y})$ and $\mathcal{J}_{MIND}(D^*)$ is a maximum being the second derivative w.r.t. $D$ a non-positive function.

Finally, at the equilibrium,
\begin{equation}
p_{X|Y}(\mathbf{x}|\mathbf{y})=\frac{p_{XY}(\mathbf{x},\mathbf{y})}{p_{Y}(\mathbf{y})} = \frac{1-D^*(\mathbf{x},\mathbf{y})}{D^*(\mathbf{x},\mathbf{y})},
\end{equation}
therefore the thesis follows.
\end{proof}

\section{Parametric (Neural Network) Implementation}
\label{sec:implementation}
The practical implementation of MIND is realized through a neural network-based learning approach. That is, firstly a training phase is implemented. It consists of transmitting repeatedly all coded messages such that at equilibrium the a-posteriori information is estimated. Then, in the testing phase, decoding can take place for each received new message by minimizing the a-posteriori conditional information. However, when $D(\mathbf{x},\mathbf{y})$ is parameterized by a neural network, it may be difficult to maximize the value function in \eqref{eq:discriminator_function}. One of the reason resides in the type of learning problem formulated in \eqref{eq:discriminator_function}. In the following, we propose two different interpretations and discuss their differences.

\begin{comment}
In addition, when $D(\mathbf{x},\mathbf{y})$ is parameterized by a neural network, the activation function of the last layer of a GAN discriminator is often chosen as the sigmoid function.  Therefore,  at the equilibrium
\begin{equation}
 D^*(\mathbf{x},\mathbf{y}; \theta) = \frac{1}{1+\exp(-\mathbf{W}^*_L \cdot \mathbf{z}_{L-1}-b^*_L)},
\end{equation}
where $\theta$ represents the parameters (weights and biases) of the neural network and $ \mathbf{z}_{L-1}$ is the output of the $(L-1)$-th layer.  Substituting the last layer output into the a posterior MI estimator in \eqref{eq:i-DMIE} yields
\begin{equation}
i(\textbf{x}|\textbf{y}) =  \biggl[ \mathbf{W}^*_L \cdot \mathbf{z}_{L-1}+b^*_L  \biggr],
\end{equation}
thus, the second-to-last layer of a GAN discriminator provides the a-posteriori MI.
\end{comment}

\subsection{Unsupervised Approach}


The value function in \eqref{eq:discriminator_function}, as it stands, requires the discriminator to be a scalar parameterized function which takes as input both the transmitted and received samples $(\mathbf{x},\mathbf{y})$. The resulting architecture is presented in Fig.\ref{fig:unsupervised}. Such formulation is general and can be applied to any coding scheme. In fact, it is sufficient during training to alternate the input joint samples $(\mathbf{x},\mathbf{y})$ with marginal samples $(K(\mathbf{x}),\mathbf{y})$. Then, during the testing phase, decoding is accomplished by finding $\mathbf{x}$ that minimizes \eqref{eq:i-DMIE}, for all possible coded messages $\mathbf{x}$. This  method can be thought as an unsupervised learning approach since no known labels are exploited in the loss function. Hence, the objective is to implicitly estimate both the joint and marginal distributions, $p_{XY}(\mathbf{x},\mathbf{y})$ and $p_{Y}(\mathbf{y})$, respectively, as shown in \eqref{eq:optimal_discriminator_1}.
Nevertheless, it has been numerically verified that the unsupervised approach may lead to unstable training, especially for low signal-to-noise ratio (SNR) values. Therefore, another architecture is proposed below.

\begin{figure}[h]
	\centering
	\includegraphics[scale=0.43]{pics/ND_unsupervised_new.pdf}
	\caption{Unsupervised discriminator architecture and relative value function. Input is fed with paired samples from the joint distribution while the output consists of a single probability ratio value.}
	\label{fig:unsupervised}
\end{figure} 

\begin{comment}
the single output network, which has to represent a transformation of a probability mass function, tends to collapse the information from previous layers in a unique node and empirically can be verified that it can lead to training problems; lastly, the knowledge of the true label (the transmitted sample) does not explicitly appear in the loss function, thus, it does not guide the training towards optimality. In summary, learning the joint distribution $p_{XY}(\mathbf{x},\mathbf{y})$ may be not convenient for our purpose. All these empirical and heuristic considerations made us rethink about the architecture to use, as explained next. 
\end{comment}

\subsection{Supervised Approach}

Another idea consists of decomposing the value function in \eqref{eq:discriminator_function} as a sum of independent elements for every input code $\mathbf{x}_i$, for $i\in \{1,\dots,M\}$ as follows
\begin{equation}
\mathcal{J}_{MIND} = \sum_{\mathbf{x}_i \in \mathcal{A}_x}{\mathcal{J}_i},
\end{equation}
where 
\begin{align}
\mathcal{J}_{i}(D_i) &= \; \mathbb{E}_{\mathbf{y} \sim p_{Y}(\mathbf{y})}\biggl[\log \biggl(D_i\bigl(\mathbf{y}\bigr)\biggr)\biggr] \\ \nonumber
&+ \mathbb{E}_{(\mathbf{x}_i,\mathbf{y}) \sim p_{Y|X}(\mathbf{y}|\mathbf{x}_i)}\biggl[\log \biggl(1-D_i\bigl(\mathbf{y}\bigr)\biggr)\biggr].
\end{align}
Given the channel input $\mathbf{x}_i$, we can feed the discriminator $D_i$ with only the output channel $\mathbf{y}$. It is then simple to verify that the maximum of $\mathcal{J}_{MIND}$ is obtained when
\begin{equation}
\label{eq:metric}
\frac{1-D^*_i}{D^*_i} = \frac{p_{XY}(\mathbf{x}_i,\mathbf{y})}{p_Y(\mathbf{y})} = p_{X|Y}(\mathbf{x}_i|\mathbf{y}).
\end{equation}
From an implementation perspective, every $D_i$ corresponds to the $i$-th output of a vector discriminator (see Fig.\ref{fig:supervised}) $\mathbf{D}(\cdot)$. Therefore the number of outputs in this setup depends on $M$ and each output can be interpreted as a likelihood ratio.
Finally, the multiple output architecture value function can be rewritten in a vector form as follows
\begin{align}
\label{eq:new_discriminator_function}
\mathcal{J}_{MIND}(\mathbf{D}) &= \; \mathbb{E}_{\mathbf{y} \sim p_{Y}(\mathbf{y})}\biggl[\log \biggl(\mathbf{D}\bigl(\mathbf{y}\bigr)\biggr) \cdot \mathbf{1}_M\biggr] \\ \nonumber
&+ \mathbb{E}_{(\mathbf{x},\mathbf{y}) \sim p_{XY}(\mathbf{x},\mathbf{y})}\biggl[\log \biggl(1-\mathbf{D}\bigl(\mathbf{y}\bigr)\biggr)\cdot \mathbf{1}_{M,\mathbf{x}}\biggr],
\end{align}
where $\mathbf{1}_M = [1,1,\dots,1]$ is a vector of $M$ elements and $\mathbf{1}_{M,\mathbf{x}}$ is the one-hot code of $\mathbf{x}$.
The new value function in \eqref{eq:new_discriminator_function} possesses now the label information in the scalar product with the one-hot code and therefore enjoys the property of supervised learning. In addition, the discriminator does not need to implicitly learn the joint distribution but only the conditional $p_{X|Y}(\mathbf{x}|\mathbf{y})$, which is done by the $M$ outputs $D_i$.

\begin{figure}
	\centering
	\includegraphics[scale=0.43]{pics/ND_supervised_new.pdf}
	\caption{Supervised discriminator architecture and relative value function. Input is fed with the received channel samples while the output consists of multiple probability ratio values, one for each transmitted message.}
	\label{fig:supervised}
\end{figure} 

\begin{figure*}
\begin{minipage}[t]{0.49\textwidth}
	\centering
  \includegraphics[scale=0.2]{pics/non-uniform_source_ALL.pdf}
  \caption{a) Symbol error rate for a 4-PAM modulation with non-uniform source distribution over an AWGN channel. Comparison among the optimal MAP decoder, MaxL decoder, MIND decoder and the estimated probability of error provided by MIND. b) Estimated source and conditional entropy (top) and estimated average mutual information (bottom) using MIND.}
	\label{fig:non-uniform_source}
\end{minipage}%
\hfill
\begin{minipage}[t]{0.49\textwidth}
	\centering
   \includegraphics[scale=0.2]{pics/non-linear_channel_ALL.pdf}
  \caption{a) Symbol error rate for a 4-PAM modulation with uniform source distribution over a non-linear channel affected by additive Gaussian noise. Comparison among MaxL decoder with no CSI, MaxL decoder with perfect CSI, MIND decoder and the estimated probability of error provided by MIND. b) Estimated source and conditional entropy (top) and estimated average mutual information (bottom) using MIND.}
  \label{fig:non-linear_channel}
\end{minipage}%
\end{figure*}

Among the $M$ discriminator outputs \eqref{eq:metric}, and similarly to the unsupervised MAP scenario, the final decoding stage implements
\begin{equation}
\mathbf{\hat{x}}_i = \argmax_{i \in \{1,...M\}} \log_2 \frac{1-D^*_i}{D^*_i} = \argmin_{\mathbf{x}_i \in \mathcal{A}_x} i(\mathbf{x}_i|\mathbf{y}).
\end{equation}
From an information perspective, MIND aims at finding the codeword $\mathbf{x}_i$ which minimizes the residual uncertainty $i(\mathbf{x}_i|\mathbf{y})$ after the observation of $\mathbf{y}$.

\begin{comment}
\subsection{Remarks}
It should be noted that the neural decoder does not assume any a-priori knowledge of either the channel model or the encoded message distribution. The estimation of the a-posteriori information implicitly includes the estimation of both the channel transfer function $H$ and the encoded message distribution. The complexity grows with the message length. NON SO SE CANCELLARE O SPIEGARE MEGLIO.
\end{comment}

\section{Estimation of Achieved Information Rate and Decoding Error Probability}
\label{sec:estimation}
Dealing with an unknown channel, a relevant question is the estimation of the information rate achieved with the used coding scheme. MIND can be exploited for such a goal. In fact, the normalized average mutual information (in bits per channel use) is given by \cite{Gallager1968}
\begin{equation}
    I(X;Y) = \frac{1}{n}(H(X)-H(X|Y))
    \end{equation}
whose computation requires the entropy of the source $H(X)$ and the conditional entropy $H(X|Y)$:
\begin{equation}
\label{eq:source_entropy}
H(X) = \mathbb{E}_{\mathbf{x} \sim p_X(\mathbf{x})}[i(\mathbf{x})] = -\sum_{\mathbf{x}_i \in \mathcal{A}_x}{p_X(\mathbf{x}_i)\log_2 p_X(\mathbf{x}_i)}
\end{equation}
\begin{equation}
\label{eq:conditional_entropy}
H(X|Y) = -\mathbb{E}_{\mathbf{y} \sim p_Y(y)}\biggl[\sum_{\mathbf{x}_i \in \mathcal{A}_x}{p_{X|Y}(\mathbf{x}_i|\mathbf{y})\log_2 p_{X|Y}(\mathbf{x}_i|\mathbf{y})}\biggr].
\end{equation}
The discriminator in MIND, for a given coding scheme, returns an estimate of the a-posteriori probability $p_{X|Y}(\mathbf{x}_i|\mathbf{y})=(1-D_i)/D_i \; \forall i=\{1,\dots,M\}$, therefore \eqref{eq:conditional_entropy} can be directly computed. About \eqref{eq:source_entropy}, it is worth mentioning that the source distribution can be obtained using Monte Carlo integration from the a-posteriori probability given by MIND, so that with a number of realizations $N\to \infty$ the source entropy becomes  
\begin{equation}
\small
\label{eq:monte-carlo}
H(X) \approx -\frac{1}{N}\sum_{\mathbf{x}_i \in \mathcal{A}_x}{\sum_{j=1}^{N}{p_{X|Y}(\mathbf{x}_i|\mathbf{y}_j)}\log_2 \biggl(\frac{1}{N} \sum_{j=1}^{N}{p_{X|Y}(\mathbf{x}_i|\mathbf{y}_j)}\biggr)},
\end{equation}
and similarly
\begin{equation}
H(X|Y) \approx -\frac{1}{N} \sum_{j=1}^{N}{\sum_{\mathbf{x}_i \in \mathcal{A}_x}{p_{X|Y}(\mathbf{x}_i|\mathbf{y}_j)\log_2 p_{X|Y}(\mathbf{x}_i|\mathbf{y}_j)}}
\end{equation}
where again the "instantaneous" conditional pdf at time $j$ is given by the MIND discriminator output.


MIND can be used also to estimate the decoding error probability $P_e$. Indeed, 
\begin{equation}
    P_e = 1-P[\mathbf{x}=\hat{\mathbf{x}}] = 1-\mathbb{E}_{\mathbf{y} \sim p_Y(y)}[p_{X|Y}(\mathbf{x}=\mathbf{\hat{x}}|\mathbf{y})],
\end{equation}
where the probability $p_{X|Y}(\mathbf{x}=\mathbf{\hat{x}}|\mathbf{y})$ comes from the decision criterion $\max_{\mathbf{x}_i} p_{X|Y}(\mathbf{x}_i|\mathbf{y})$ which is the instantaneous probability of correct decoding. Hence, using again Monte Carlo integration
\begin{equation}
    P_e \stackrel{N\to \infty}{=} 1-\frac{1}{N} \sum_{j=1}^{N}{\max_{\mathbf{x}_i \in \mathcal{A}_x} p_{X|Y}(\mathbf{x}_i|\mathbf{y}_j)}.
\end{equation}

In the next section, we compare MIND to other well-known decoding criteria in scenarios for which an analytic solution is available. Furthermore, thanks to the network ability to estimate the a-posteriori probability, the average mutual information between pairs of input-output channel samples is also reported.


\section{Numerical Results}
\label{sec:results}

The assessment of the proposed decoding approach is done by considering the following three representative scenarios: a) Uncoded transmission of symbols that are produced by a non-uniform source over an additive Gaussian noise channel; b) Uncoded transmission in a non-linear channel with additive Gaussian noise; c) Short block coded transmission in an additive Middleton noise channel.

Details about the architecture and hyper parameters of the neural networks are reported in the GitHub repository \cite{MIND_github}.

\subsection{Non-uniform Source}
To show the ability of the decoder to learn and exploit data dependence at the source, we firstly consider 4-PAM transmission with symbols $\mathbf{x}_i \in \{-3,-1,1,3\}$ and mass function given by a non-uniform source $p(\mathbf{x}) = [(1-P)/2, P/2, (1-P)/2, P/2]$, with $P=0.05$. It should be noted that data dependence at the source can be either created with a coding stage from a source that emits i.i.d. symbols, or intrinsically by the source of data traffic having certain statistics, as in this case. In Fig. \ref{fig:non-uniform_source}a the probability of symbol error is shown. MIND is compared with the optimal MAP decoder that knows the a-priori distribution of the transmitted data symbols $p(\mathbf{x})$ and with the maximum-likelihood (MaxL) decoder that only knows the Gaussian nature of the channel and assumes a source with uniform distribution. Moreover, the estimated $P_e$ is also shown. In Fig. \ref{fig:non-uniform_source}b, the source entropy, the conditional entropy and the average achieved mutual information, as estimated by MIND, are reported. It is worth noticing that the average mutual information increases with the SNR, or alternatively, the residual uncertainty introduced by the channel decreases. In addition, the estimated source entropy tends to stabilize around its real value given by $P\log_2(P/2)-(1-P)\log_2((1-P)/2)$.
\begin{figure*}[h]
	\centering
	\includegraphics[scale=0.25]{pics/middleton_info.pdf}
	\caption{Bit error rate and mutual information of a short block coded transmission in an additive Middleton noise channel: a) Repetition code; b) Hamming code; c) Convolutional code.}
	\label{fig:middleton}
\end{figure*} 
\subsection{Non-linear Channel}
As a second example, we consider 4-PAM uncoded transmission with uniform source distribution over a non-linear channel with additive Gaussian noise. The objective of this experiment is to show the ability of MIND to discover such a channel non-linearity. In particular, the channel introduces a non-linearity (for instance because of the presence of non-linear amplifiers) modeled as $y_k=\text{sign}(x_k)\sqrt{|x_k|}+n_k$, where $k$ denotes the $k$-th time instant.
Fig.\ref{fig:non-linear_channel}a demonstrates how the MIND decoder has managed to implicitly learn the non-linear channel model during the training phase and effectively use such information during decoding. A comparison with the MaxL decoder with and without perfect channel state information (CSI) knowledge is also conducted. Results show that MIND exhibits performance close to optimality. Fig. \ref{fig:non-linear_channel}b illustrates both the behaviour of the entropies and the average mutual information as function of the SNR.

\subsection{Additive Non-gaussian Noise Channel}
Lastly, and perhaps as most interesting example, we consider a short block coded transmission over an additive non-gaussian channel. The aim is to assess the ability of MIND to learn and exploit the presence of non-gaussian noise. In particular, we suppose that the noise follows a truncated Middleton distribution \cite{Middleton1977}, also called Bernoulli-Gaussian noise model, so that at any given time instant $k$ it is obtained as $n_k = (1-\epsilon_k)n_{1,k}+\epsilon_k n_{2,k}$ where $n_{1,k} \sim \mathcal{N}(0,\sigma_b^2)$ is a zero-mean Gaussian random variable with variance $\sigma_b^2$ and $n_{2,k} \sim \mathcal{N}(0,B\sigma_b^2)$ is also a zero-mean Gaussian random variable but with variance $B$ times larger. Instead, $\epsilon_k$ is a Bernoulli random variable with probability of success $P$. The pdf of the noise samples is then given by
\begin{equation}
p_{N}(\mathbf{n}_k) = (1-P)\mathcal{N}(0,\sigma_b^2)+P\mathcal{N}(0,B\sigma_b^2).
\end{equation} 
Assuming that the noise model is known and that BPSK symbols are transmitted, two decoding strategies can be devised. The first, denoted with MaxL Middleton, uses maximum likelihood decoding with the known conditional pdf $p(\mathbf{y}|\mathbf{x})=p_{N}(\mathbf{y}-\mathbf{x})$. The second strategy is a genie decoder that knows the outcome of the Bernoulli event for every time instance. That is, it knows whether a received sample is hit by Gaussian noise with variance $\sigma_b^2$ or $B\sigma_b^2$. A third decoding strategy is offered by MIND that learns the channel statistics.
We distinguish among three different types of codes: a) a binary repetition code with length $5$; b) a $(7,4)$ Hamming code; c) a rate $1/2$ convolutional code with memory $2$ and block-length $18$. For each of them, the bit error rate (BER) obtained with the genie, the MaxL Middleton and the MIND decoders is reported. The parameter $B$ was set to $5$ in all the experiments involving Middleton noise, wherein we also set $P=0.05$. 
Fig.\ref{fig:middleton} shows the gain provided by the neural-based decoder scheme over the classical maximum likelihood one that exploits the Middleton distribution in (24). With MIND, the BER performance gets closer to the genie decoder. Furthermore, Fig.\ref{fig:middleton} also reports an estimate of the average mutual information provided by MIND.

\begin{comment}
\begin{enumerate}
\item Uncoded transmission in additive Gaussian noise with block transmission of binary PSK data symbols that are dependent; this is to show the ability of the decoder to learn and exploit data dependence at the source. The MI decoder is compared with the optimal MAP decoder that knows (a-priori) that joint distribution of the block of transmitted data symbols $p(\mathbf{x})$. In addition the ML decoder that assumes i.i.d. transmitted data symbols is considered. It should be noted that data dependency at the source can be either created with a coding stage from a source that emits i.i.d. symbols, or intrinsically by the source of data traffic having certain semantics.
\item Uncoded transmission in a non-linear channel with additive Gaussian noise; this is to show the ability of the MI decoder to discover such a channel non linearity. We still assume to transmit i.i.d. 4-PAM symbols. However, the channel introduces an non linearity (for instance because of the presence of non linear amplifiers). Thus, the considered channel model is $y_i=sign(\mathbf{x}_i)\sqrt{|\mathbf{x}_i|}+n_i$.
\item Uncoded transmission in an additive Middleton noise channel; this is to assess the ability of the MI decoder to learn and exploit the presence of non Gaussian noise. A comparison in terms of BER is made between the MI decoder, the genie optimal decoder that knows the noise statistics, and the one that assumes Gaussian noise. The considered modulation is binary PSK with equally probable symbols.

\item Short block coded transmission in additive white Gaussian noise. In particular, a rate $1/2$ convolutional code with memory 2 is assumed. The input block has length 20 bits. In addition, the distribution of the source of $2^{20}$ information blocks is not uniform but Poissonian. A comparison between the MI decoder and the Viterbi decoding algorithm is made.
\end{enumerate}
\end{comment}

\section{Conclusions}
\label{sec:conclusions}
In this paper, MIND, a neural decoder that uses the mutual information as decoding criterion, has been proposed. Two specific architectures have been described and they are capable of learning the conditional information $i(\mathbf{x}|\mathbf{y})$ of the codeword $\mathbf{x}$ given the channel output observation $\mathbf{y}$ in unknown channels. Additionally, MIND allows the estimation of the achieved information rate with the used coding scheme as well as the decoding error probability. Several numerical results obtained in illustrative channels show that MIND can achieve the performance of the genie MAP decoder that perfectly knows the channel model and source distribution. It outperforms the conventional max-log decoder that assumes the presence of Gaussian noise.        

\bibliographystyle{IEEEtran}

\section{Discriminative Mutual Information Estimators}
\label{sec:theory}
The mutual information $I(X;Y)$ quantifies the statistical dependence between two random variables, $X$ and $Y$, by measuring the amount of information of $X$ via observation of $Y$. Unlike the correlation, the mutual information measures linear and non-linear associations \cite{Letizia2020} since its estimation requires the knowledge of three distributions
\begin{equation}
I(X;Y) = \mathbb{E}_{(x,y)\sim p_{XY}(x,y)}\biggl[\log\frac{p_{XY}(x,y)}{p_X(x)p_Y(y)}\biggr],
\end{equation}
where $p_X(x)$ and $p_Y(y)$ are the distributions of $X$ and $Y$, respectively, and $p_{XY}(x,y)$ is the joint probability distribution. The presence of these three components renders the mutual information estimation a challenging task and several approaches have been proposed to tackle the problem. Standard and traditional approaches rely on binning, density and kernel estimation \cite{Moon1995} and $k$-nearest neighbours \cite{Kraskov2004}. However, for applications such as channel coding, high-dimensional data is often involved and classic approaches fail to scale. Moreover, they can not easily be integrated in an optimization framework that targets the channel capacity. To overcome such limitations, neural networks can be used to obtain a parametric representations of the variational lower bounds on the mutual information, where the objective of the network is to estimate the density ratio. 

One of the first examples of neural estimators is MINE \cite{Mine2018}. The idea behind MINE is to exploit the Donsker-Varadhan dual representation of the Kullback-Leibler (KL) divergence to produce the bound to optimize
\begin{align}
\label{eq:MINE}
I(X;Y) \geq & I_{MINE}(X;Y) \nonumber \\
= & \sup_{\theta \in \Theta} \mathbb{E}_{(x,y)\sim p_{XY}(x,y)}[T_{\theta}(x,y)]  \nonumber \\
& - \log(\mathbb{E}_{(x,y)\sim p_X(x) p_Y(y)}[e^{T_{\theta}(x,y)}]),
\end{align}
where $\theta \in \Theta$ parameterizes a family of functions $T_{\theta} : \mathcal{X}\times \mathcal{Y} \to \mathbb{R}$ via a deep neural network. MINE is a biased estimator, consequently, the authors suggested to replace the expectation in the denominator of the gradient with an exponential moving average. However, MINE suffers from high-variance estimations and a recent solution named SMILE \cite{Song2020} tries to alleviate such issue by clipping the density ratio between $-e^{-\tau}$ and $e^{\tau}$, with $\tau \geq 0$. SMILE is equivalent to MINE in the limit $\tau \rightarrow +\infty$.

A mutual information unbiased variational lower bound relies on the $f$-divergence representation proposed in \cite{Nguyen2010}
\begin{align}
\label{eq:NWJ}
I(X;Y) \geq & I_{NWJ}(X;Y) \nonumber \\
= & \sup_{\theta \in \Theta} \mathbb{E}_{(x,y)\sim p_{XY}(x,y)}[T_{\theta}(x,y)]  \nonumber \\
& -\mathbb{E}_{(x,y)\sim p_X(x) p_Y(y)}[e^{T_{\theta}(x,y)-1}].
\end{align}
where it is easy to verify that $I_{NWJ}\leq I_{MINE}$.

In \cite{LetiziaNIPS}, the discriminative mutual information estimator (DIME) was proposed. It is a neural estimator that mimics the task of the discriminator in GANs \cite{Goodfellow2014}. Indeed, the optimal discriminator in the GAN or $f$-GAN framework can be used either as a direct or indirect mutual information estimator when it is fed with paired and unpaired input data of distribution $p_{XY}$ and $p_X p_Y$, respectively. We refer to as direct estimator, the one where the optimal discriminator $D^*$ is proportional to the density ratio in \eqref{eq:density_ratio}. An example of this is the d-DIME estimator, a mutual information variational lower bound that reads as follows
\begin{equation}
\label{eq:d_DIME}
I(X;Y) \geq \tilde{I}_{dDIME}(X;Y) = \frac{\mathcal{J}_{\alpha}(D^*)}{\alpha}+1-\log(\alpha),
\end{equation}
where $\mathcal{J}_{\alpha}(D)$, $\alpha>0$, is a value function defined as
\begin{align}
\mathcal{J}_{\alpha}(D) = \; & \alpha \cdot \mathbb{E}_{(x,y) \sim p_{XY}(x,y)}\biggl[\log \biggl(D\bigl(x,y\bigr)\biggr)\biggr] \nonumber \\ 
& +\mathbb{E}_{(x,y) \sim p_{X}(x)p_{Y}(y)}\biggl[-D\bigl(x,y\bigr)\biggr],
\label{eq:discriminator_function}
\end{align}
and
\begin{equation}
\label{eq:optimal_discriminator_2}
D^*(x,y) = \alpha \cdot \frac{p_{XY}(x,y)}{p_{X}(x)\cdot p_Y(y)} = \arg \max_D \mathcal{J}_{\alpha}(D).
\end{equation}
When the positive function $D(\cdot)$ is parameterized by a neural network, the NWJ estimator can be thought as a special case of d-DIME with $\alpha=1$ and $T_{\theta}(x,y) = \log(D_{\theta}(x,y))+1$.

In the following section we extend DIME, in particular, we introduce $f$-DIME, a discriminative mutual information estimator based on the $f$-divergence measure.   

\section{$f$-DIME}
\label{sec:f-DIME}
In this section, we describe a general methodology to estimate the mutual information by applying the variational representation of $f$-divergence functionals $D_f(P||Q)$. In detail, if the measures $P$ and $Q$ are absolutely continuous w.r.t. $\diff x$ and possess densities $p$ and $q$, then the $f$-divergence reads as follows
\begin{equation}
D_f(P||Q) = \int_{\mathcal{X}}{q(x)f\biggl(\frac{p(x)}{q(x)}\biggr)\diff x},
\end{equation}
where $\mathcal{X}$ is a compact domain and the function $f:\mathbb{R}_+ \to \mathbb{R}$ is convex, lower semicontinuous and satisfies $f(1)=0$. Under these conditions, the authors of \cite{Nguyen2010} (and \cite{Nowozin2016} for the $f$-GAN) exploited the Fenchel convex duality to derive a lower bound on $D_f$
\begin{equation}
\label{eq:f_bound}
D_f(P||Q) \geq \sup_{T\in \mathbb{R}} \biggl\{ \mathbb{E}_{x \sim p(x)} \bigl[T(x)\bigr]-\mathbb{E}_{x\sim q(x)}\bigl[f^*\bigl(T(x)\bigr)\bigr]\biggr\},
\end{equation}
where $T: \mathcal{X} \to \mathbb{R}$ and $f^*$ is the Fenchel conjugate of $f$ defined as
\begin{equation}
f^*(t) := \sup_{u\in \mathbb{R}} \{ ut -f(u)\}.
\end{equation}
Therein, it was shown that the bound in \eqref{eq:f_bound} is tight for optimal values of $T(x)$ and it takes the following form
\begin{equation}
\label{eq:optimal_ratio}
T^*(x) = f^{\prime} \biggl(\frac{p(x)}{q(x)}\biggr),
\end{equation}
where $f^{\prime}$ is the derivative of $f$.

To create a novel neural estimator based on the representation in \eqref{eq:f_bound}, we firstly substitute the mutual information $I(X;Y)$ with its KL divergence representation 
\begin{equation}
I(X;Y) = D_{KL}(p_{XY}||p_X p_Y),
\end{equation}
and by defining 
\begin{align}
\mathcal{J}_{f}(T) = \; & \mathbb{E}_{(x,y) \sim p_{XY}(x,y)}\bigl[T\bigl(x,y\bigr)\bigr] \nonumber \\ 
& -\mathbb{E}_{(x,y) \sim p_{X}(x)p_{Y}(y)}\bigl[f^*\bigl(T\bigl(x,y\bigr)\bigr)\bigr],
\label{eq:value_function_f}
\end{align}
we propose to parametrize $T(x,y)$ with a deep neural network $T_{\theta}$ of parameters $\theta$ and solve with gradient ascent and back-propagation
\begin{equation}
\theta^* = \arg \max_{\theta} \mathcal{J}_f(T_{\theta}).
\end{equation}
The main idea is the following: \textit{it is possible to exploit the relation \eqref{eq:optimal_ratio} to estimate the mutual information}. Indeed, if the derivative of $f$ is invertible, 
\begin{equation}
\bigl(f^{\prime}\bigr)^{-1}\bigl(T_{\theta^*}(x,y)\bigr) = \frac{p_{XY}(x,y)}{p_X(x)p_Y(y)},
\end{equation}
thus, having access to $T_{\theta^*}$ and using the density ratio estimated above, we obtain $f$-DIME as
\begin{equation}
\label{eq:f_dime}
I_{fDIME}(X;Y) = \mathbb{E}_{(x,y) \sim p_{XY}(x,y)}\biggl[ \log \biggl(\bigl(f^{\prime}\bigr)^{-1}\bigl(T_{\theta^{*}}(x,y)\bigr) \biggr) \biggr].
\end{equation}

We argue that any value function $\mathcal{J}_f$ of the form in \eqref{eq:value_function_f}, seen as dual representation of a certain $f$-divergence $D_f$, can be maximized to estimate the mutual information using \eqref{eq:f_dime}. As a notable example, the GAN-based i-DIME estimator in \cite{LetiziaNIPS} can be derived using $f$-DIME with generator 
\begin{equation}
f(u) = u\log u-(u+1)\log(u+1)+\log4
\end{equation}
and $T(x)=\log(D(x))$.
However, only a subset of $\mathcal{J}_f$ can be used inside the capacity-driven autoencoder loss function since a maximization of the mutual information is also required. In particular, only the generator $f$ of the KL divergence is suitable for the mutual information maximization, since in such case
\begin{equation}
\max_{p_X(x)} D_f = \max_{p_X(x)} D_{KL} = \max_{p_X(x)} I(X;Y).
\end{equation}

The generator function of the KL divergence is $f(u) = u\log u$ with conjugate $f^*(t) = e^{t-1}$. It easy to notice that such choice of $f$ coincides with the NWJ estimator or d-DIME when $T(x) = \log(D(x))$ and $\alpha=1$ (see \eqref{eq:d_DIME}). Nevertheless, as shown in \cite{LSGAN}, the value function plays a fundamental role during the training process of a discriminator. Consequently, we combined the approach described for the derivation of $f$-DIME with d-DIME to obtain a family of lower bounds on the mutual information that may result in more robust estimators according to specific case studies.

\begin{lemma}
\label{lemma:Lemma1}
Let $X\sim p_X(x)$ and $Y\sim p_{Y}(y|x)$ be the channel input and output, respectively. Let $D(\cdot)$ be a positive function. If $\mathcal{J}_{\gamma}(D)$, $\gamma>0$, is a value function defined as 
\begin{align}
\mathcal{J}_{\gamma}(D) = \; & \gamma \cdot \mathbb{E}_{(x,y) \sim p_{XY}(x,y)}\biggl[\log \biggl(D\bigl(x,y\bigr)\biggr)\biggr] \nonumber \\ 
& + \mathbb{E}_{(x,y) \sim p_{X}(x)p_{Y}(y)}\biggl[- D^{\gamma}\bigl(x,y\bigr)\biggr],
\end{align}
then
\begin{equation}
I(X;Y) \geq \tilde{I}_{\gamma DIME}(X;Y) = \mathcal{J}_{\gamma}(D^*)+1,
\end{equation}
where
\begin{equation}
\label{eq:optimal_discriminator_gamma}
D^*(x,y) = \biggl(\frac{p_{XY}(x,y)}{p_{X}(x)\cdot p_Y(y)}\biggr)^{1/\gamma} = \arg \max_D \mathcal{J}_{\gamma}(D).
\end{equation}
\end{lemma}
\begin{comment}
\begin{proof}
Consider a scaled generator $f(u) = \frac{u}{\gamma}\log u$ and for simplicity of notation, denote $p_{XY}$ and $p_Xp_Y$ with $p$ and $q$, respectively. Then
\begin{equation}
D_{KL}(p||q) = \gamma \int_{x}{q(x) \frac{p(x)}{\gamma q(x)}\log\biggl(\frac{p(x)}{q(x)}\biggr) \diff x},
\end{equation}
with conjugate $f^*(t)$, with $t \in \mathbb{R}$ given by
\begin{equation}
f^*(t) = \frac{e^{\gamma t -1}}{\gamma}.
\end{equation}
Substituting in \eqref{eq:f_bound} yields to
\begin{equation}
D_{KL}(p||q) \geq \gamma \sup_{T\in \mathbb{R}} \biggl\{ \mathbb{E}_{x \sim p(x)} \bigl[T(x)\bigr]-\frac{1}{\gamma}\mathbb{E}_{x\sim q(x)}\bigl[e^{\gamma T(x)-1} \bigr]\biggr\}.
\end{equation}
Using \eqref{eq:optimal_ratio} it is easy to verify that the optimal value of $T$ is the log-likelihood ratio rather than the density ratio itself. Indeed,
\begin{equation}
T^*(x) = \frac{1}{\gamma}\log\biggl(\frac{p(x)}{q(x)}\biggr) + \frac{1}{\gamma}.
\end{equation}
Finally, with the change of variable $T(x)=\log(D(x))+1/{\gamma}$, the optimal discriminator has form
\begin{equation}
D^*(x) = \biggl(\frac{p(x)}{q(x)}\biggr)^{1/\gamma}
\end{equation}
and
\begin{align}
D_{KL}(p||q) \geq \; & \sup_{D\in \mathbb{R}_+} \biggl\{ \gamma \mathbb{E}_{x \sim p(x)} \bigl[\log \bigl( D(x) \bigr) \bigr] \nonumber \\
& -\mathbb{E}_{x\sim q(x)}\bigl[D^{\gamma}(x) \bigr]\biggr\} + 1.
\end{align}
Denoting the two terms inside the supremum with $J_{\gamma}(D)$ and replacing $p$ and $q$ with the joint and marginal distributions, concludes the proof.
\end{proof}
\end{comment}
\begin{figure}
\centering
  	\includegraphics[scale = 0.255]{pics/gammaDIME}
  	\caption{Behaviour of the functional $J_{\gamma}$ varying the parameter $\gamma$. Concavity of the neighbourhood of the maximal value depends on $\gamma$.}
  	\label{fig:gammaDIME}
\end{figure}

Fig.~\ref{fig:gammaDIME} illustrates how the value function in \eqref{eq:value_function_f} (divided by $\gamma$) behaves according to the choice of the positive parameter $\gamma$. The position of the maximum depends on the density ratio as described in \eqref{eq:optimal_discriminator_gamma} ($R = 1$ in the figure). The role of $\gamma$ is to modify the concavity of the value function to improve the convergence rate of the gradient ascent algorithm. Moreover, when $D$ is parameterized by a neural network, the choice of the learning rate and activation functions may be related with $\gamma$. Indeed, notice that $\gamma$-DIME can also be derived from d-DIME using the substitution $D(x,y) \rightarrow D^{\gamma}(x,y)$.
\section{Results}
\label{sec:results}
In the experimental results, we consider the transmission of $M$ messages at rate $R$ over an AWGN channel, for which we know the channel capacity under an
average power constraint. The choice of AWGN channel shall not be seen as a limitation since the formulation transcends the channel characteristics. Moreover, it provides a reference curve for the validation of the estimators and it allows comparisons across several application domains \cite{Mine2018, LetiziaNIPS}. 

We describe two scenarios: a) the transmission at a rate $R>1$ and short code-length $n$ and b) the transmission at a low rate with larger code-length. For both scenarios, we train a rate-driven autoencoder (see \cite{Letizia2021}) and compare its performance for different estimators in terms of BLER and accuracy of the mutual information estimation. 

For all the experiments and as a proof of concept, we use simple encoder, decoder and neural mutual information estimator architectures. In particular, the encoder is a shallow neural network with $M$ neurons in the input layer, $M$ in the hidden one and $2n$ neurons at the output while the decoder is its complementary.  
The discriminative estimator possesses instead a fixed architecture, a
two layer multilayer perceptron neural network with $200$
neurons in each layer with ReLU and sigmoid activation functions. The only
difference in the estimators resides in the final layer where we use a linear dense layer for MINE, while a softplus layer for the DIME based estimators, defined as
\begin{equation}
 D^*(x,y; \theta) = \log(1+\exp(W^*_L \cdot z_{L-1}+b^*_L)),
\end{equation}
where $W$ and $b$ are weights and biases, respectively, $z$ is the output of the previous layer and $L$ is the number of layers of the network.
\begin{figure}
\centering

  	\includegraphics[scale=0.255]{pics/ber_all_1}
  	\caption{Comparison of the BLER obtained by the autoencoder AE($3,18$) and AE($6,3$) with an AWGN intermediate layer, using different estimators with $\beta = 0.2$, $\epsilon = 0.2$.}
  	\label{fig:ber_all}
\end{figure}
The training of the autoencoder was conducted at a fixed $E_b/N_0$ ratio of $7$ dB. We minimized the loss function alternating encoder and decoder weights update with the discriminator weights update. The number of iterations was set to $10$k with a learning rate of $0.01$ for both the components. For more implementation details, we refer to the code publicly available. \footnote{\url{https://github.com/tonellolab/capacity-approaching-autoencoders}}

The first scenario describes a digital transmission scheme with input alphabet size $M=64$ over $3$ channel uses. We denote such system as AE($6,3$) since $k = \log_2(M)$. The communication rate is $R=2$. Fig.~\ref{fig:ber_all} shows the autoencoder performance in terms of BLER for different estimators. From the curves we can infer that the $\gamma$-DIME estimators, and in particular, higher values of $\gamma$ seem to provide improved performance in terms of BLER. The result of Fig.~\ref{fig:MI_6_6} is insightful: none of the estimators provides stable estimations, especially at high SNRs where it is expected that the mutual information saturates to the information rate $R=2$. The DIME based estimators are more accurate compared to the divergent MINE. Nevertheless, the greater is $\gamma$, the greater is the variance (see theorem 2 of \cite{Song2020}) of the estimation. It is also interesting to notice that the concavity of $J_{\gamma}$ (shown in Fig.~\ref{fig:gammaDIME}) influences the mutual information estimation as it appears in the results. Indeed, for low values of $\gamma$, fixed learning rate and training iterations, the estimators are trained by moving on the flat curves of Fig.~\ref{fig:gammaDIME}, which is equivalent as using a small learning rate to upgrade the gradients. For completeness and to show the advantage of the proposed estimators under fading, we report in Fig.~\ref{fig:MI_6_6_Ray} the achieved estimated information rate with Rayleigh fading. Also in this case, MINE diverges for high SNR values while the DIME estimators are much more accurate and tend to saturate to the rate $R=2$. 

\begin{figure}
\centering

  	\includegraphics[scale=0.255]{pics/MI_6_6_1}
  	\caption{Comparison of different mutual information estimators for the autoencoder AE($6,3$) with $\beta = 0.2$, $\epsilon = 0.2$ on the AWGN channel.}
  	\label{fig:MI_6_6}
\end{figure}

\begin{figure}
\centering

  	\includegraphics[scale=0.255]{pics/MI_6_6_Ray_1}
  	\caption{Comparison of different mutual information estimators for the autoencoder AE($6,3$) with $\beta = 0.2$, $\epsilon = 0.2$ on the Rayleigh fading channel.}
  	\label{fig:MI_6_6_Ray}
\end{figure}

The second scenario describes a digital transmission scheme with input alphabet size $M=8$ over $9$ channel uses. We denote such system as AE($3,9$). The communication rate is $R=1/3$. Fig.~\ref{fig:ber_all} illustrates the advantage of using MINE for low values of the mutual information. Under the same training iterations, MINE is a tighter bound compared to $\gamma$-DIME and therefore converges faster and provides stronger gradients for the encoder. However, Fig.~\ref{fig:MI_3_18} highlights the fact that for low SNRs and alphabet dimension and for relatively small $n$, the estimators perform similarly. We can conclude that the choice of the estimator to use varies according to input constraints such as alphabet dimension, code-length and SNR. We did not comment about the architecture to use although we remark that it is one of the main aspects in modern machine learning applications and therefore we expect significant improvements. 

\begin{figure}
\centering

  	\includegraphics[scale=0.255]{pics/MI_3_18_1}
  	\caption{Comparison of the achieved information rate using different mutual information estimators for the autoencoder AE($3,18$) with $\beta = 0.2$, $\epsilon = 0.2$ on the AWGN channel.}
  	\label{fig:MI_3_18}
\end{figure}
\section{Capacity-driven Autoencoders}
\label{sec:autoencoders}

The considered communication system comprise an encoder, a channel a decoder. The encoder maps $M=2^k$ messages of $k$ bits into $M$ encoded messages of $n$ symbols with $N-ary$ alphabet. The uncoded message is denoted with the vector $\mathbf{b}$, while the coded message with the vector $\mathbf{x}$. Thus, the code rate is $R=k/n$ bits per channel use. The channel conveys the coded message to the receiver. Without loss of generality, we can write that the received message $\mathbf{b}=H(\mathbf{x},\mathbf{h},\mathbf{n})$, where the channel transfer function $H$ implicitly takes into account the channel internal state $\mathbf{h}$ and the noise $\mathbf{n}$. 
For instance, for an LTI channel with additive noise, we obtain $\mathbf{y}=\mathbf{h}*\mathbf{x}+\mathbf{n}$. 

The aim of the decoder is to retrieve the transmitted message so that to minimize error, or equivalently to maximize the received information for each transmitted message. In fact, if we consider the mutual information of the pair $(\mathbf{y},\mathbf{x})$ this can be written as
 
\begin{equation}
I(\mathbf{y};\mathbf{x})=i(\mathbf{x})-i(\mathbf{x}|\mathbf{y})
\end{equation}
where $i(\mathbf{x})$ is the information of the transmitted message, $i(\mathbf{x}|\mathbf{y})$ is the a posterior information (residual uncertainty observing $\mathbf{y}$. It follows that for each transmitted message, the mutual information is maximized when the a posterior information is minimized. The a posteriori information is given by 

\begin{equation}
i(\mathbf{x}|\mathbf{y})=-log_2p(\mathbf{x}|\mathbf{y})
\end{equation}
so that the decoding criteria becomes

\begin{equation}
\mathbf{x}=\argmin{i(\mathbf{x}|\mathbf{y})}=\argmax{p(\mathbf{x}|\mathbf{y})
\end{equation}

which corresponds to the maximum a posteriori decoding principle. 

To compute the a posterior information we need to know the a posterior probability $p(\textbf{x}|\textbf{y}$. If the channel transfer function can be modeled, then the decoder can be realized according to known approaches. For instance, for an LTI channel with additive Gaussian noise, the decoder is realized by firstly estimating the channel impulse response and then implementing a MAP sequence estimator that uses as the main metric the Euclidean distance between the received and transmitted message. The task becomes more intriguing for channels that cannot be easily modelled. In such a case a learning strategy turn out to be an attractive solution. Herein, a supervised learning approach is considered. That is, firstly a training phase is implemented. It consists of transmitting repeatedly all coded messages. The a posterior information is estimated. 
Then, in a second phase, decoding can take place for each received message by minimizing the a posterior conditional information. 

The task is formidable because learning the posterior information is not trivial. To do so we can exploit the discriminator used in generative adverbial neural network architectures (GANs). In fact the discriminator     
  

 



\section{Introduction}
\label{sec:introduction}
In digital communication systems, data detection and decoding is a fundamental task at the receiver. The maximum a-posteriori (MAP) decoding approach is the optimal one to minimize the symbol or sequence error probability. If the channel model is known, the decoder can be realized according to known algorithms. For instance, for a linear time invariant (LTI) channel with additive Gaussian noise, the decoder comprises a first stage where the channel impulse response is estimated. Then, a max-log-MAP sequence estimator algorithm is implemented. The decoding metric is the Euclidean distance between the received and hypothesized transmitted message filtered by the channel impulse response. The task becomes more intriguing for channels that cannot be easily modelled or are unknown. In such a case, a learning strategy turns out to be an attractive solution. 
 


Recent advancements in machine learning for communications have fostered the design of systems with a data-driven paradigm. As a consequence, physical layer design has been reinterpreted using machine learning techniques to improve the coding and decoding scheme performance \cite{Oshea2017, Nachmani2018, Dorner2018, Alberge2019, Stark2019}. In particular, the work in \cite{Oshea2017} introduced the concept of autoencoder-based communication systems. In contrast to traditional bottom-up approaches, the encoder and decoder blocks can be jointly learned during the end-to-end autoencoder training phase \cite{Oshea2017}. The autoencoder is a deep neural network that maps a sequence of input bits $\mathbf{s}$ into a sequence of output bits $\hat{\mathbf{s}}$. The input sequence is generally transformed into data symbols $x$ that are then fed into an intermediate channel layer (or network), which introduces constraints, distortions and uncertainties. Given the intermediate channel model, the autoencoder is typically trained by minimizing the cross-entropy loss function, so that it essentially performs a classification task. However, it is known that training a classifier via cross-entropy often suffers from overfitting issues, especially for large networks \cite{Zhang2017}. Moreover, autoencoders for communications shall consider the channel capacity during the learning process in order to produce optimal channel input samples (latent codes). In this direction, the work in \cite{Letizia2021} proposed to use a mutual information regularizer to control and estimate the amount of information stored in the latent representation. In addition, it proposed to use label smoothing to improve the network decoding ability. In \cite{Letizia2021}, the mutual information was computed with the neural estimator MINE \cite{Mine2018} by inserting an appropriate block in the autoencoder architecture. MINE was also exploited in \cite{Wunder2019} to study optimal coding schemes when no channel model is attainable. Nevertheless, MINE suffers from high bias and variance. Thus, new mutual information estimators have been recently proposed to subvert such limitations \cite{Song2020, Poole2019a, LetiziaNIPS}. Among them, a promising approach is given by the discriminative mutual information estimator (DIME) \cite{LetiziaNIPS}, a neural network that directly estimates the density ratio
\begin{equation}
\label{eq:density_ratio}
R(\mathbf{x},\mathbf{y}) = \frac{p_{XY}(\mathbf{x},\mathbf{y})}{p_X(\mathbf{x})\cdot p_Y(\mathbf{y})}
\end{equation}
instead of the individual densities in \eqref{eq:density_ratio}. Concurrently, it is expected that novel and optimal channel coding techniques based on mutual information learning, estimation and maximization can significantly impact the development of beyond 5G communication technologies.

In this paper, inspired by the DIME estimator and the $f$-GAN training objectives \cite{Nowozin2016}, we firstly propose a new family of estimators referred to as $f$-DIME. We secondly propose $\gamma$-DIME, a family of estimators that can be used in the end-to-end autoencoder training process to target the channel capacity. We then include the developed estimators in the capacity-driven autoencoder proposed in \cite{Letizia2021} and evaluate their performance in terms of block-error-rate (BLER) and accuracy of the mutual information estimation w.r.t. to MINE.

The rest of the paper is organized as follows. Section \ref{sec:autoencoders} revisits the autoencoder-based communication systems and discusses the major advantages of using a mutual information regularization term in the loss function. Section \ref{sec:theory} reviews some variational lower bounds on the mutual information and the related estimators. Section \ref{sec:f-DIME} presents a new set of discriminative estimators. Section \ref{sec:results} compares the estimators and illustrates the results. Finally, conclusions are reported.


\section{Conclusions}
\label{sec:conclusions}
In this paper, we presented $f$-DIME, a set of novel mutual information estimators based on the $f$-divergence. When $f$ is the function generator of the KL divergence, we obtained as a special case $\gamma$-DIME, a parametric estimator that can be used both for mutual information estimation and design of optimal channel coding. We integrated $\gamma$-DIME as the mutual information regularization term inside the autoencoder loss function. A comparison in terms of estimation accuracy and BLER was done with MINE, an established neural mutual information estimator. Results show the advantage of adopting $\gamma$-DIME over other estimators in situations where the mutual information between the channel input and output is high. 
\subsection{Lemmas and Proofs}
\label{sec:omitted_proofs}

 \setcounter{lemma}{0}
 \setcounter{theorem}{0}

\begin{lemma}
Let $X\sim p_X(x)$ and $Y\sim p_{Y}(y|x)$ be the channel input and output, respectively. Let $D(\cdot)$ be a positive function. If $\mathcal{J}_{\gamma}(D)$, $\gamma>0$, is a value function defined as 
\begin{align}
\mathcal{J}_{\gamma}(D) = \; & \gamma \cdot \mathbb{E}_{(x,y) \sim p_{XY}(x,y)}\biggl[\log \biggl(D\bigl(x,y\bigr)\biggr)\biggr] \nonumber \\ 
& + \mathbb{E}_{(x,y) \sim p_{X}(x)p_{Y}(y)}\biggl[- D^{\gamma}\bigl(x,y\bigr)\biggr],
\end{align}
then
\begin{equation}
I(X;Y) \geq \tilde{I}_{\gamma DIME}(X;Y) = \mathcal{J}_{\gamma}(D^*)+1,
\end{equation}
where
\begin{equation}
D^*(x,y) = \biggl(\frac{p_{XY}(x,y)}{p_{X}(x)\cdot p_Y(y)}\biggr)^{1/\gamma} = \arg \max_D \mathcal{J}_{\gamma}(D).
\end{equation}
\end{lemma}

\begin{proof}
Consider a scaled generator $f(u) = \frac{u}{\gamma}\log u$ and for simplicity of notation, denote $p_{XY}$ and $p_Xp_Y$ with $p$ and $q$, respectively. Then
\begin{equation}
D_{KL}(p||q) = \gamma \int_{x}{q(x) \frac{p(x)}{\gamma q(x)}\log\biggl(\frac{p(x)}{q(x)}\biggr) \diff x},
\end{equation}
with conjugate $f^*(t)$, with $t \in \mathbb{R}$ given by
\begin{equation}
f^*(t) = \frac{e^{\gamma t -1}}{\gamma}.
\end{equation}
Substituting in \eqref{eq:f_bound} yields to
\begin{equation}
D_{KL}(p||q) \geq \sup_{T\in \mathbb{R}} \biggl\{ \gamma \mathbb{E}_{x \sim p(x)} \bigl[T(x)\bigr]-\mathbb{E}_{x\sim q(x)}\bigl[e^{\gamma T(x)-1} \bigr]\biggr\}.
\end{equation}
Using \eqref{eq:optimal_ratio} it is easy to verify that the optimal value of $T$ is the log-likelihood ratio rather than the density ratio. Indeed,
\begin{equation}
T^*(x) = \frac{1}{\gamma}\log\biggl(\frac{p(x)}{q(x)}\biggr) + \frac{1}{\gamma}.
\end{equation}
Finally, with the change of variable $T(x)=\log(D(x))+1/{\gamma}$, the optimal discriminator has form
\begin{equation}
D^*(x) = \biggl(\frac{p(x)}{q(x)}\biggr)^{1/\gamma}
\end{equation}
and
\small
\begin{equation}
D_{KL}(p||q) \geq 1 + \sup_{D\in \mathbb{R}_+} \gamma \mathbb{E}_{x \sim p(x)} \bigl[\log \bigl( D(x) \bigr) \bigr] \nonumber \\
 -\mathbb{E}_{x\sim q(x)}\bigl[D^{\gamma}(x) \bigr].
\end{equation}
Denoting the two terms inside the supremum with $J_{\gamma}(D)$ and replacing $p$ and $q$ with the joint and marginal distributions, concludes the proof.
\end{proof}

\section{Capacity-driven Autoencoders}
\label{sec:autoencoders}

The considered communication system comprise an encoder, a channel a decoder. The encoder maps $M=2^k$ messages of $k$ bits into $M$ encoded messages of $n$ symbols with $N-ary$ alphabet. The uncoded message is denoted with the vector $\mathbf{b}$, while the coded message with the vector $\mathbf{x}$. Thus, the code rate is $R=k/n$ bits per channel use. The channel conveys the coded message to the receiver. Without loss of generality, we can write that the received message $\mathbf{b}=H(\mathbf{x},\mathbf{h},\mathbf{n})$, where the channel transfer function $H$ implicitly takes into account the channel internal state $\mathbf{h}$ and the noise $\mathbf{n}$. 
For instance, for an LTI channel with additive noise, we obtain $\mathbf{y}=\mathbf{h}*\mathbf{x}+\mathbf{n}$. 

The aim of the decoder is to retrieve the transmitted message so that to minimize error, or equivalently to maximize the received information for each transmitted message. In fact, if we consider the mutual information of the pair $(\mathbf{y},\mathbf{x})$ this can be written as
 
\begin{equation}
I(\mathbf{y};\mathbf{x})=i(\mathbf{x})-i(\mathbf{x}|\mathbf{y})
\end{equation}
where $i(\mathbf{x})$ is the information of the transmitted message, $i(\mathbf{x}|\mathbf{y})$ is the a posterior information (residual uncertainty observing $\mathbf{y}$. It follows that for each transmitted message, the mutual information is maximized when the a posterior information is minimized. The a posteriori information is given by 

\begin{equation}
i(\mathbf{x}|\mathbf{y})=-log_2p(\mathbf{x}|\mathbf{y})
\end{equation}
so that the decoding criteria becomes

\begin{equation}
\mathbf{x}=\argmin{i(\mathbf{x}|\mathbf{y})}=\argmax{p(\mathbf{x}|\mathbf{y})
\end{equation}

which corresponds to the maximum a posteriori decoding principle. 

To compute the a posterior information we need to know the a posterior probability $p(\textbf{x}|\textbf{y}$. If the channel transfer function can be modeled, then the decoder can be realized according to known approaches. For instance, for an LTI channel with additive Gaussian noise, the decoder is realized by firstly estimating the channel impulse response and then implementing a MAP sequence estimator that uses as the main metric the Euclidean distance between the received and transmitted message. The task becomes more intriguing for channels that cannot be easily modelled. In such a case a learning strategy turn out to be an attractive solution. Herein, a supervised learning approach is considered. That is, firstly a training phase is implemented. It consists of transmitting repeatedly all coded messages. The a posterior information is estimated. 
Then, in a second phase, decoding can take place for each received message by minimizing the a posterior conditional information. 

The task is formidable because learning the posterior information is not trivial. To do so we can exploit the discriminator used in generative adverbial neural network architectures (GANs). In fact the discriminator     
  

 



\section{Conclusions}
\label{sec:conclusions}
In this paper, we presented $f$-DIME, a set of novel mutual information estimators based on the $f$-divergence. When $f$ is the function generator of the KL divergence, we obtained as a special case $\gamma$-DIME, a parametric estimator that can be used both for mutual information estimation and design of optimal channel coding. We integrated $\gamma$-DIME as the mutual information regularization term inside the autoencoder loss function. A comparison in terms of estimation accuracy and BLER was done with MINE, an established neural mutual information estimator. Results show the advantage of adopting $\gamma$-DIME over other estimators in situations where the mutual information between the channel input and output is high. 
\section{Introduction}
\label{sec:introduction}
In digital communication systems, data detection and decoding is a fundamental task at the receiver. The maximum a-posteriori (MAP) decoding approach is the optimal one to minimize the symbol or sequence error probability. If the channel model is known, the decoder can be realized according to known algorithms. For instance, for a linear time invariant (LTI) channel with additive Gaussian noise, the decoder comprises a first stage where the channel impulse response is estimated. Then, a max-log-MAP sequence estimator algorithm is implemented. The decoding metric is the Euclidean distance between the received and hypothesized transmitted message filtered by the channel impulse response. The task becomes more intriguing for channels that cannot be easily modelled or are unknown. In such a case, a learning strategy turns out to be an attractive solution. 
 


Recent advancements in machine learning for communications have fostered the design of systems with a data-driven paradigm. As a consequence, physical layer design has been reinterpreted using machine learning techniques to improve the coding and decoding scheme performance \cite{Oshea2017, Nachmani2018, Dorner2018, Alberge2019, Stark2019}. In particular, the work in \cite{Oshea2017} introduced the concept of autoencoder-based communication systems. In contrast to traditional bottom-up approaches, the encoder and decoder blocks can be jointly learned during the end-to-end autoencoder training phase \cite{Oshea2017}. The autoencoder is a deep neural network that maps a sequence of input bits $\mathbf{s}$ into a sequence of output bits $\hat{\mathbf{s}}$. The input sequence is generally transformed into data symbols $x$ that are then fed into an intermediate channel layer (or network), which introduces constraints, distortions and uncertainties. Given the intermediate channel model, the autoencoder is typically trained by minimizing the cross-entropy loss function, so that it essentially performs a classification task. However, it is known that training a classifier via cross-entropy often suffers from overfitting issues, especially for large networks \cite{Zhang2017}. Moreover, autoencoders for communications shall consider the channel capacity during the learning process in order to produce optimal channel input samples (latent codes). In this direction, the work in \cite{Letizia2021} proposed to use a mutual information regularizer to control and estimate the amount of information stored in the latent representation. In addition, it proposed to use label smoothing to improve the network decoding ability. In \cite{Letizia2021}, the mutual information was computed with the neural estimator MINE \cite{Mine2018} by inserting an appropriate block in the autoencoder architecture. MINE was also exploited in \cite{Wunder2019} to study optimal coding schemes when no channel model is attainable. Nevertheless, MINE suffers from high bias and variance. Thus, new mutual information estimators have been recently proposed to subvert such limitations \cite{Song2020, Poole2019a, LetiziaNIPS}. Among them, a promising approach is given by the discriminative mutual information estimator (DIME) \cite{LetiziaNIPS}, a neural network that directly estimates the density ratio
\begin{equation}
\label{eq:density_ratio}
R(\mathbf{x},\mathbf{y}) = \frac{p_{XY}(\mathbf{x},\mathbf{y})}{p_X(\mathbf{x})\cdot p_Y(\mathbf{y})}
\end{equation}
instead of the individual densities in \eqref{eq:density_ratio}. Concurrently, it is expected that novel and optimal channel coding techniques based on mutual information learning, estimation and maximization can significantly impact the development of beyond 5G communication technologies.

In this paper, inspired by the DIME estimator and the $f$-GAN training objectives \cite{Nowozin2016}, we firstly propose a new family of estimators referred to as $f$-DIME. We secondly propose $\gamma$-DIME, a family of estimators that can be used in the end-to-end autoencoder training process to target the channel capacity. We then include the developed estimators in the capacity-driven autoencoder proposed in \cite{Letizia2021} and evaluate their performance in terms of block-error-rate (BLER) and accuracy of the mutual information estimation w.r.t. to MINE.

The rest of the paper is organized as follows. Section \ref{sec:autoencoders} revisits the autoencoder-based communication systems and discusses the major advantages of using a mutual information regularization term in the loss function. Section \ref{sec:theory} reviews some variational lower bounds on the mutual information and the related estimators. Section \ref{sec:f-DIME} presents a new set of discriminative estimators. Section \ref{sec:results} compares the estimators and illustrates the results. Finally, conclusions are reported.


\section{Introduction}
\label{sec:introduction}
In digital communication systems, data detection and decoding is a fundamental task implemented at the receiver. The maximum a-posteriori (MAP) decoding approach is the optimal one to minimize the symbol or sequence error probability \cite{Proakis2001,Bahl1974}. If the channel model is known, the MAP decoder can be realized according to known algorithms. For instance, for a linear time invariant channel (LTI) with additive Gaussian noise, the decoder comprises a first stage where the channel impulse response is estimated. Then, a max-log-MAP sequence estimator algorithm is implemented. The decoding metric is essentially related to the Euclidean distance between the received and hypothesized transmitted message filtered by the channel impulse response \cite{Proakis2001}. The task becomes more intriguing for channels that cannot be easily modeled or that are unknown. In such a case, a learning strategy turns out to be an attractive solution and it can be enabled by recent advances in machine learning for communications \cite{Oshea2017, Nachmani2018, Dorner2018}.

In this paper, we derive a neural architecture for optimal decoding in an unknown channel. We propose to use the mutual information (MI) of the input-output channel pair as decoding metric. The computation of the MI is a formidable task and it is unknown in many instances. Therefore, the MI has to be learned, which can be done in principle with neural architectures \cite{Mine2018, Wunder2019, Letizia2021, Song2020, LetiziaNIPS}. 
Additionally, in the problem at hand, the decoder has a more specific task: it has to learn and minimize the a-posteriori information $-\log_2{(p_{X|Y}(\mathbf{x}|\mathbf{y}))}$ of the transmitted codeword $\mathbf{x}$ given the observed/received signal vector $\mathbf{y}$. The direct estimation of such a-posteriori information becomes then the objective. Such an estimation is possible by directly learning the conditional probability density function (pdf) $p_{X|Y}(\mathbf{x}|\mathbf{y})$ with a new estimator that exploits a discriminative model referred to as mutual information neural decoder (MIND). MIND allows not only to perform the decoding task, but also to return an estimate of the channel achievable rate and of the decoding error probability. 

This paper is organized as follows. Section \ref{sec:mid} presents the decoding approach principle. Section \ref{sec:ndec} describes the core idea about MIND. Section \ref{sec:implementation} discusses two practical neural architectures. Section \ref{sec:estimation} explains how MIND can be exploited to estimate both the achievable rate and the decoding error probability. Numerical results for a number of scenarios are reported in Section \ref{sec:results}. Finally, the conclusions are drawn.


\section{Mutual Information Decoding Principle}
\label{sec:mid}

The considered communication system comprises an encoder, a channel, and a decoder. The encoder maps $M=2^k$ messages of $k$ bits into $M$ encoded messages of $n$ (not necessarily binary) symbols. The uncoded message is denoted with the vector $\mathbf{b}$, while the coded message with the vector $\mathbf{x}$ belonging to the set $\mathcal{A}_x$. Thus, the code rate is $R=k/n$ bits per channel use. The channel conveys the coded message to the receiver. Without loss of generality, we can write that the received message is $\mathbf{y}=H(\mathbf{x},\mathbf{h},\mathbf{n})$, where the channel transfer function $H$ implicitly models the channel internal state $\mathbf{h}$ including stochastic variables, e.g. noise, $\mathbf{n}$. For instance, for an LTI channel with additive noise, we obtain $\mathbf{y}=\mathbf{h}*\mathbf{x}+\mathbf{n}$, where $*$ denotes convolution.

The aim of the decoder is to retrieve the transmitted message so that to minimize the decoding error, or equivalently to maximize the received information for each transmitted message. Indeed, if we consider the mutual information of the pair $(\mathbf{y},\mathbf{x})$, with joint pdf $p_{XY}(\mathbf{x},\mathbf{y})$, and marginals $p_X(\mathbf{x})$ and $p_Y(\mathbf{y})$, this can be written as

\begin{equation}
i(\mathbf{y};\mathbf{x})=i(\mathbf{x})-i(\mathbf{x}|\mathbf{y})
\end{equation}
where $i(\mathbf{x})$ is the information of the transmitted message, $i(\mathbf{x}|\mathbf{y})$ is the a-posteriori information (residual uncertainty observing $\mathbf{y}$) \cite{Gallager1968}. It follows that for each transmitted message, the mutual information is maximized when the a-posteriori information is minimized. The a-posteriori information is given by
\begin{equation}
\label{eq:PI}
i(\mathbf{x}|\mathbf{y})=-\log_2 p_{X|Y}(\mathbf{x}|\mathbf{y}) = -\log_2{\frac{p_{XY}(\mathbf{x},\mathbf{y})}{p_{Y}(\mathbf{y})}}
\end{equation}
so that the decoding criterion becomes
\begin{equation}
\mathbf{\hat{x}}=\argmin_{\mathbf{x}\in \mathcal{A}_x}{i(\mathbf{x}|\mathbf{y})} = \argmax_{\mathbf{x}\in \mathcal{A}_x}{\log_2 p_{X|Y}(\mathbf{x}|\mathbf{y})}
\end{equation}
which corresponds to the log-MAP decoding principle \cite{Proakis2001,Bahl1974}.

To compute the a-posteriori information, we need to know the a-posteriori probability $p_{X|Y}(\mathbf{x}|\mathbf{y})$. If the channel transfer function can be modeled, then the decoder is realized according to known approaches. In fact, the a-posteriori probability can be explicitly calculated as $p_{X|Y}(\mathbf{x}|\mathbf{y})=p_{Y|X}(\mathbf{y}|\mathbf{x})p_{X}(\mathbf{x})$, i.e., from the conditional channel pdf and the coded message a-priori probability. However, for unknown channels and sources, we have to resort to a learning strategy as explained in the next section.

\section{Discriminative Formulation of MIND}
\label{sec:ndec}

In the following, we show that it is possible to design a parametric discriminator whose output leads to the estimation of the a-posteriori information. The optimal discriminator is found by maximizing an appropriately defined loss function.
Before describing the proposed solution, it should be noted that in generative adversarial networks (GANs) \cite{Goodfellow2014} a discriminator $D$ is used to distinguish if a sample is an original one coming from real data rather than a fake one generated by the generator $G$. The adversarial training pushes the discriminator $D(\mathbf{a})$ towards the optimum value \cite{Goodfellow2014} 
\begin{equation}
D^*(\mathbf{a}) = \frac{p_{data}(\mathbf{a})}{p_{data}(\mathbf{a})+p_{gen}(\mathbf{a})}
\end{equation} 
where $p_{data}(\mathbf{a})$ is the collected data pdf and $p_{gen}(\mathbf{a})$ is the generated data pdf.
Hence, the output of the optimal discriminator can be used to estimate the probability density ratio 
\begin{equation}
\frac{p_{gen}(\mathbf{a})}{p_{data}(\mathbf{a})} = \frac{1-D^*(\mathbf{a})}{D^*(\mathbf{a})}.
\end{equation}
 
\begin{comment}
The training procedure for $G$ is to maximize the probability of $D$ making a mistake. GANs implement a min-max two-player game such that the Nash equilibrium is reached. Given an input noise vector $\mathbf{z}$ with joint pdf $p_{noise}(\mathbf{z})$, the mapping to the data space is achieved through $G(\mathbf{z};\theta_{gen})$. It has been proved that the generator implicitly learns the true distribution when the following value function $V(G,D)$ is used
\begin{align}
\label{eq:GAN}
V(G,D) &=  \mathbb{E}_{\mathbf{x} \sim p_{data}(\mathbf{x})}[\log D(\mathbf{x})] \\ \nonumber
&+ \mathbb{E}_{\mathbf{z} \sim p_{noise}(\mathbf{z})}[\log(1-D(G(\mathbf{z})))],
\end{align}
an the equilibrium is reached so that $p_{gen} = p_{data}$.

The adversarial training pushes the discriminator $D(\mathbf{x})$ towards the optimum value (see \cite{Goodfellow2014})
\begin{equation}
D^*(\mathbf{x}) = \frac{p_{data}(\mathbf{x})}{p_{data}(\mathbf{x})+p_{gen}(\mathbf{x})},
\end{equation}
thus, we notice that the output of the optimal discriminator indirectly estimates the probability density ratio $p_{gen}/p_{data}$ since we obtain
\begin{equation}
\frac{p_{gen}(\mathbf{x})}{p_{data}(\mathbf{x})} = \frac{1-D^*(\mathbf{x})}{D^*(\mathbf{x})}.
\end{equation}
\end{comment}
Now, the idea is to design a discriminator that estimates the probability density ratio $p_{XY}(\mathbf{x},\mathbf{y})/p_{Y}(\mathbf{y})$ and consequently the a-posteriori information as shown in \eqref{eq:PI}. 

The following Lemma provides an a-posteriori information estimator that exploits a discriminator trained to minimize a certain loss function. At the equilibrium, a transformation of the discriminator output yields the searched density ratio.   

\begin{comment}
take only the optimal discriminator in the formulation of \eqref{eq:GAN} and use it as indirect a-posteriori probability and mutual information estimator. More in detail, since $p_{X|Y}(\mathbf{x}|\mathbf{y})=p_{XY}(\mathbf{x},\mathbf{y})/p_{Y}(\mathbf{y})$, then $i(\mathbf{x}|\mathbf{y})$ can be indirectly estimated using the optimum discriminator $D^*$ when $p_{data}\equiv  p_Y(\mathbf{y})$ and $p_{gen} \equiv p_{XY}(\mathbf{x}, \mathbf{y})$. 
\end{comment}

\begin{lemma}
\label{lemma:Lemma1}
Let $\mathbf{x}\sim p_X(\mathbf{x})$ and $\mathbf{y}\sim p_Y(\mathbf{y})$ be the channel input and output, respectively. Let $H(\cdot)$ be the channel transfer function, in general stochastic, such that $\mathbf{y} = H(\mathbf{x}, \mathbf{h}, \mathbf{n} )$, with $\mathbf{h}$ and $\mathbf{n}$ being internal state and noise variables, respectively. If $\mathcal{J}_{MIND}(D)$ is the value function defined as
\begin{align}
\label{eq:discriminator_function}
\mathcal{J}_{MIND}(D) &= \; \mathbb{E}_{(\mathbf{x},\mathbf{y}) \sim p_{X}(\mathbf{x})p_{Y}(\mathbf{y})}\biggl[\log \biggl(D\bigl(K,\mathbf{y}\bigr)\biggr)\biggr] \nonumber \\ 
& + \mathbb{E}_{(\mathbf{x},\mathbf{y}) \sim p_{XY}(\mathbf{x},\mathbf{y})}\biggl[\log \biggl(1-D\bigl(\mathbf{x},\mathbf{y}\bigr)\biggr)\biggr],
\end{align}
where $K$ is a constant, then
\begin{equation}
\label{eq:optimal_discriminator_1}
D^*(\mathbf{x},\mathbf{y}) = \frac{p_{Y}(\mathbf{y})}{p_{Y}(\mathbf{y}) + p_{XY}(\mathbf{x},\mathbf{y})}= \arg \max_D \mathcal{J}_{MIND}(D),
\end{equation}
and
\begin{equation}
\label{eq:i-DMIE}
i(\mathbf{x}|\mathbf{y}) = -\log \biggl(\frac{1-D^*(\mathbf{x},\mathbf{y})}{D^*(\mathbf{x},\mathbf{y})}\biggr).
\end{equation}
\end{lemma}

\begin{proof}
From the hypothesis of the Lemma, the value function can be rewritten as
\begin{align}
\mathcal{J}_{MIND}(D) &= \int_{\mathcal{T}_x} \int_{\mathcal{T}_y}\biggl[p_{XY}(K(\mathbf{x}),\mathbf{y}) \log \biggl(D(\mathbf{x},\mathbf{y})\biggr) \nonumber \\ 
&+ p_{XY}(\mathbf{x},\mathbf{y}) \log \biggl(1-D(\mathbf{x},\mathbf{y})\biggr)\biggr] \diff \mathbf{x} \diff \mathbf{y}
\end{align}
where $\mathcal{T} = \mathcal{T}_x \times \mathcal{T}_y $ is the domain of definition of the joint density function and $K(\mathbf{x})$ is a function that maps any $\mathbf{x}$ into the constant value $K$ with probability one.

To maximize $\mathcal{J}_{MIND}(D)$, a necessary and sufficient condition requires to take the derivative of the integrand with respect to $D$ and to set it to $0$, yielding the following equation in $D$
\begin{equation}
\frac{p_{XY}(K(\mathbf{x}),\mathbf{y})}{D(\mathbf{x},\mathbf{y})} -\frac{p_{XY}(\mathbf{x},\mathbf{y})}{1-D(\mathbf{x},\mathbf{y})} =0.
\end{equation}
The solution of the above equation yields the optimum discriminator $D^*(\mathbf{x},\mathbf{y})$ since: $p_{XY}(K(\mathbf{x}),\mathbf{y})=p_Y(\mathbf{y})$ and $\mathcal{J}_{MIND}(D^*)$ is a maximum being the second derivative w.r.t. $D$ a non-positive function.

Finally, at the equilibrium,
\begin{equation}
p_{X|Y}(\mathbf{x}|\mathbf{y})=\frac{p_{XY}(\mathbf{x},\mathbf{y})}{p_{Y}(\mathbf{y})} = \frac{1-D^*(\mathbf{x},\mathbf{y})}{D^*(\mathbf{x},\mathbf{y})},
\end{equation}
therefore the thesis follows.
\end{proof}

\section{Parametric (Neural Network) Implementation}
\label{sec:implementation}
The practical implementation of MIND is realized through a neural network-based learning approach. That is, firstly a training phase is implemented. It consists of transmitting repeatedly all coded messages such that at equilibrium the a-posteriori information is estimated. Then, in the testing phase, decoding can take place for each received new message by minimizing the a-posteriori conditional information. However, when $D(\mathbf{x},\mathbf{y})$ is parameterized by a neural network, it may be difficult to maximize the value function in \eqref{eq:discriminator_function}. One of the reason resides in the type of learning problem formulated in \eqref{eq:discriminator_function}. In the following, we propose two different interpretations and discuss their differences.

\begin{comment}
In addition, when $D(\mathbf{x},\mathbf{y})$ is parameterized by a neural network, the activation function of the last layer of a GAN discriminator is often chosen as the sigmoid function.  Therefore,  at the equilibrium
\begin{equation}
 D^*(\mathbf{x},\mathbf{y}; \theta) = \frac{1}{1+\exp(-\mathbf{W}^*_L \cdot \mathbf{z}_{L-1}-b^*_L)},
\end{equation}
where $\theta$ represents the parameters (weights and biases) of the neural network and $ \mathbf{z}_{L-1}$ is the output of the $(L-1)$-th layer.  Substituting the last layer output into the a posterior MI estimator in \eqref{eq:i-DMIE} yields
\begin{equation}
i(\textbf{x}|\textbf{y}) =  \biggl[ \mathbf{W}^*_L \cdot \mathbf{z}_{L-1}+b^*_L  \biggr],
\end{equation}
thus, the second-to-last layer of a GAN discriminator provides the a-posteriori MI.
\end{comment}

\subsection{Unsupervised Approach}


The value function in \eqref{eq:discriminator_function}, as it stands, requires the discriminator to be a scalar parameterized function which takes as input both the transmitted and received samples $(\mathbf{x},\mathbf{y})$. The resulting architecture is presented in Fig.\ref{fig:unsupervised}. Such formulation is general and can be applied to any coding scheme. In fact, it is sufficient during training to alternate the input joint samples $(\mathbf{x},\mathbf{y})$ with marginal samples $(K(\mathbf{x}),\mathbf{y})$. Then, during the testing phase, decoding is accomplished by finding $\mathbf{x}$ that minimizes \eqref{eq:i-DMIE}, for all possible coded messages $\mathbf{x}$. This  method can be thought as an unsupervised learning approach since no known labels are exploited in the loss function. Hence, the objective is to implicitly estimate both the joint and marginal distributions, $p_{XY}(\mathbf{x},\mathbf{y})$ and $p_{Y}(\mathbf{y})$, respectively, as shown in \eqref{eq:optimal_discriminator_1}.
Nevertheless, it has been numerically verified that the unsupervised approach may lead to unstable training, especially for low signal-to-noise ratio (SNR) values. Therefore, another architecture is proposed below.

\begin{figure}[h]
	\centering
	\includegraphics[scale=0.43]{pics/ND_unsupervised_new.pdf}
	\caption{Unsupervised discriminator architecture and relative value function. Input is fed with paired samples from the joint distribution while the output consists of a single probability ratio value.}
	\label{fig:unsupervised}
\end{figure} 

\begin{comment}
the single output network, which has to represent a transformation of a probability mass function, tends to collapse the information from previous layers in a unique node and empirically can be verified that it can lead to training problems; lastly, the knowledge of the true label (the transmitted sample) does not explicitly appear in the loss function, thus, it does not guide the training towards optimality. In summary, learning the joint distribution $p_{XY}(\mathbf{x},\mathbf{y})$ may be not convenient for our purpose. All these empirical and heuristic considerations made us rethink about the architecture to use, as explained next. 
\end{comment}

\subsection{Supervised Approach}

Another idea consists of decomposing the value function in \eqref{eq:discriminator_function} as a sum of independent elements for every input code $\mathbf{x}_i$, for $i\in \{1,\dots,M\}$ as follows
\begin{equation}
\mathcal{J}_{MIND} = \sum_{\mathbf{x}_i \in \mathcal{A}_x}{\mathcal{J}_i},
\end{equation}
where 
\begin{align}
\mathcal{J}_{i}(D_i) &= \; \mathbb{E}_{\mathbf{y} \sim p_{Y}(\mathbf{y})}\biggl[\log \biggl(D_i\bigl(\mathbf{y}\bigr)\biggr)\biggr] \\ \nonumber
&+ \mathbb{E}_{(\mathbf{x}_i,\mathbf{y}) \sim p_{Y|X}(\mathbf{y}|\mathbf{x}_i)}\biggl[\log \biggl(1-D_i\bigl(\mathbf{y}\bigr)\biggr)\biggr].
\end{align}
Given the channel input $\mathbf{x}_i$, we can feed the discriminator $D_i$ with only the output channel $\mathbf{y}$. It is then simple to verify that the maximum of $\mathcal{J}_{MIND}$ is obtained when
\begin{equation}
\label{eq:metric}
\frac{1-D^*_i}{D^*_i} = \frac{p_{XY}(\mathbf{x}_i,\mathbf{y})}{p_Y(\mathbf{y})} = p_{X|Y}(\mathbf{x}_i|\mathbf{y}).
\end{equation}
From an implementation perspective, every $D_i$ corresponds to the $i$-th output of a vector discriminator (see Fig.\ref{fig:supervised}) $\mathbf{D}(\cdot)$. Therefore the number of outputs in this setup depends on $M$ and each output can be interpreted as a likelihood ratio.
Finally, the multiple output architecture value function can be rewritten in a vector form as follows
\begin{align}
\label{eq:new_discriminator_function}
\mathcal{J}_{MIND}(\mathbf{D}) &= \; \mathbb{E}_{\mathbf{y} \sim p_{Y}(\mathbf{y})}\biggl[\log \biggl(\mathbf{D}\bigl(\mathbf{y}\bigr)\biggr) \cdot \mathbf{1}_M\biggr] \\ \nonumber
&+ \mathbb{E}_{(\mathbf{x},\mathbf{y}) \sim p_{XY}(\mathbf{x},\mathbf{y})}\biggl[\log \biggl(1-\mathbf{D}\bigl(\mathbf{y}\bigr)\biggr)\cdot \mathbf{1}_{M,\mathbf{x}}\biggr],
\end{align}
where $\mathbf{1}_M = [1,1,\dots,1]$ is a vector of $M$ elements and $\mathbf{1}_{M,\mathbf{x}}$ is the one-hot code of $\mathbf{x}$.
The new value function in \eqref{eq:new_discriminator_function} possesses now the label information in the scalar product with the one-hot code and therefore enjoys the property of supervised learning. In addition, the discriminator does not need to implicitly learn the joint distribution but only the conditional $p_{X|Y}(\mathbf{x}|\mathbf{y})$, which is done by the $M$ outputs $D_i$.

\begin{figure}
	\centering
	\includegraphics[scale=0.43]{pics/ND_supervised_new.pdf}
	\caption{Supervised discriminator architecture and relative value function. Input is fed with the received channel samples while the output consists of multiple probability ratio values, one for each transmitted message.}
	\label{fig:supervised}
\end{figure} 

\begin{figure*}
\begin{minipage}[t]{0.49\textwidth}
	\centering
  \includegraphics[scale=0.2]{pics/non-uniform_source_ALL.pdf}
  \caption{a) Symbol error rate for a 4-PAM modulation with non-uniform source distribution over an AWGN channel. Comparison among the optimal MAP decoder, MaxL decoder, MIND decoder and the estimated probability of error provided by MIND. b) Estimated source and conditional entropy (top) and estimated average mutual information (bottom) using MIND.}
	\label{fig:non-uniform_source}
\end{minipage}%
\hfill
\begin{minipage}[t]{0.49\textwidth}
	\centering
   \includegraphics[scale=0.2]{pics/non-linear_channel_ALL.pdf}
  \caption{a) Symbol error rate for a 4-PAM modulation with uniform source distribution over a non-linear channel affected by additive Gaussian noise. Comparison among MaxL decoder with no CSI, MaxL decoder with perfect CSI, MIND decoder and the estimated probability of error provided by MIND. b) Estimated source and conditional entropy (top) and estimated average mutual information (bottom) using MIND.}
  \label{fig:non-linear_channel}
\end{minipage}%
\end{figure*}

Among the $M$ discriminator outputs \eqref{eq:metric}, and similarly to the unsupervised MAP scenario, the final decoding stage implements
\begin{equation}
\mathbf{\hat{x}}_i = \argmax_{i \in \{1,...M\}} \log_2 \frac{1-D^*_i}{D^*_i} = \argmin_{\mathbf{x}_i \in \mathcal{A}_x} i(\mathbf{x}_i|\mathbf{y}).
\end{equation}
From an information perspective, MIND aims at finding the codeword $\mathbf{x}_i$ which minimizes the residual uncertainty $i(\mathbf{x}_i|\mathbf{y})$ after the observation of $\mathbf{y}$.

\begin{comment}
\subsection{Remarks}
It should be noted that the neural decoder does not assume any a-priori knowledge of either the channel model or the encoded message distribution. The estimation of the a-posteriori information implicitly includes the estimation of both the channel transfer function $H$ and the encoded message distribution. The complexity grows with the message length. NON SO SE CANCELLARE O SPIEGARE MEGLIO.
\end{comment}

\section{Estimation of Achieved Information Rate and Decoding Error Probability}
\label{sec:estimation}
Dealing with an unknown channel, a relevant question is the estimation of the information rate achieved with the used coding scheme. MIND can be exploited for such a goal. In fact, the normalized average mutual information (in bits per channel use) is given by \cite{Gallager1968}
\begin{equation}
    I(X;Y) = \frac{1}{n}(H(X)-H(X|Y))
    \end{equation}
whose computation requires the entropy of the source $H(X)$ and the conditional entropy $H(X|Y)$:
\begin{equation}
\label{eq:source_entropy}
H(X) = \mathbb{E}_{\mathbf{x} \sim p_X(\mathbf{x})}[i(\mathbf{x})] = -\sum_{\mathbf{x}_i \in \mathcal{A}_x}{p_X(\mathbf{x}_i)\log_2 p_X(\mathbf{x}_i)}
\end{equation}
\begin{equation}
\label{eq:conditional_entropy}
H(X|Y) = -\mathbb{E}_{\mathbf{y} \sim p_Y(y)}\biggl[\sum_{\mathbf{x}_i \in \mathcal{A}_x}{p_{X|Y}(\mathbf{x}_i|\mathbf{y})\log_2 p_{X|Y}(\mathbf{x}_i|\mathbf{y})}\biggr].
\end{equation}
The discriminator in MIND, for a given coding scheme, returns an estimate of the a-posteriori probability $p_{X|Y}(\mathbf{x}_i|\mathbf{y})=(1-D_i)/D_i \; \forall i=\{1,\dots,M\}$, therefore \eqref{eq:conditional_entropy} can be directly computed. About \eqref{eq:source_entropy}, it is worth mentioning that the source distribution can be obtained using Monte Carlo integration from the a-posteriori probability given by MIND, so that with a number of realizations $N\to \infty$ the source entropy becomes  
\begin{equation}
\small
\label{eq:monte-carlo}
H(X) \approx -\frac{1}{N}\sum_{\mathbf{x}_i \in \mathcal{A}_x}{\sum_{j=1}^{N}{p_{X|Y}(\mathbf{x}_i|\mathbf{y}_j)}\log_2 \biggl(\frac{1}{N} \sum_{j=1}^{N}{p_{X|Y}(\mathbf{x}_i|\mathbf{y}_j)}\biggr)},
\end{equation}
and similarly
\begin{equation}
H(X|Y) \approx -\frac{1}{N} \sum_{j=1}^{N}{\sum_{\mathbf{x}_i \in \mathcal{A}_x}{p_{X|Y}(\mathbf{x}_i|\mathbf{y}_j)\log_2 p_{X|Y}(\mathbf{x}_i|\mathbf{y}_j)}}
\end{equation}
where again the "instantaneous" conditional pdf at time $j$ is given by the MIND discriminator output.


MIND can be used also to estimate the decoding error probability $P_e$. Indeed, 
\begin{equation}
    P_e = 1-P[\mathbf{x}=\hat{\mathbf{x}}] = 1-\mathbb{E}_{\mathbf{y} \sim p_Y(y)}[p_{X|Y}(\mathbf{x}=\mathbf{\hat{x}}|\mathbf{y})],
\end{equation}
where the probability $p_{X|Y}(\mathbf{x}=\mathbf{\hat{x}}|\mathbf{y})$ comes from the decision criterion $\max_{\mathbf{x}_i} p_{X|Y}(\mathbf{x}_i|\mathbf{y})$ which is the instantaneous probability of correct decoding. Hence, using again Monte Carlo integration
\begin{equation}
    P_e \stackrel{N\to \infty}{=} 1-\frac{1}{N} \sum_{j=1}^{N}{\max_{\mathbf{x}_i \in \mathcal{A}_x} p_{X|Y}(\mathbf{x}_i|\mathbf{y}_j)}.
\end{equation}

In the next section, we compare MIND to other well-known decoding criteria in scenarios for which an analytic solution is available. Furthermore, thanks to the network ability to estimate the a-posteriori probability, the average mutual information between pairs of input-output channel samples is also reported.


\section{Numerical Results}
\label{sec:results}

The assessment of the proposed decoding approach is done by considering the following three representative scenarios: a) Uncoded transmission of symbols that are produced by a non-uniform source over an additive Gaussian noise channel; b) Uncoded transmission in a non-linear channel with additive Gaussian noise; c) Short block coded transmission in an additive Middleton noise channel.

Details about the architecture and hyper parameters of the neural networks are reported in the GitHub repository \cite{MIND_github}.

\subsection{Non-uniform Source}
To show the ability of the decoder to learn and exploit data dependence at the source, we firstly consider 4-PAM transmission with symbols $\mathbf{x}_i \in \{-3,-1,1,3\}$ and mass function given by a non-uniform source $p(\mathbf{x}) = [(1-P)/2, P/2, (1-P)/2, P/2]$, with $P=0.05$. It should be noted that data dependence at the source can be either created with a coding stage from a source that emits i.i.d. symbols, or intrinsically by the source of data traffic having certain statistics, as in this case. In Fig. \ref{fig:non-uniform_source}a the probability of symbol error is shown. MIND is compared with the optimal MAP decoder that knows the a-priori distribution of the transmitted data symbols $p(\mathbf{x})$ and with the maximum-likelihood (MaxL) decoder that only knows the Gaussian nature of the channel and assumes a source with uniform distribution. Moreover, the estimated $P_e$ is also shown. In Fig. \ref{fig:non-uniform_source}b, the source entropy, the conditional entropy and the average achieved mutual information, as estimated by MIND, are reported. It is worth noticing that the average mutual information increases with the SNR, or alternatively, the residual uncertainty introduced by the channel decreases. In addition, the estimated source entropy tends to stabilize around its real value given by $P\log_2(P/2)-(1-P)\log_2((1-P)/2)$.
\begin{figure*}[h]
	\centering
	\includegraphics[scale=0.25]{pics/middleton_info.pdf}
	\caption{Bit error rate and mutual information of a short block coded transmission in an additive Middleton noise channel: a) Repetition code; b) Hamming code; c) Convolutional code.}
	\label{fig:middleton}
\end{figure*} 
\subsection{Non-linear Channel}
As a second example, we consider 4-PAM uncoded transmission with uniform source distribution over a non-linear channel with additive Gaussian noise. The objective of this experiment is to show the ability of MIND to discover such a channel non-linearity. In particular, the channel introduces a non-linearity (for instance because of the presence of non-linear amplifiers) modeled as $y_k=\text{sign}(x_k)\sqrt{|x_k|}+n_k$, where $k$ denotes the $k$-th time instant.
Fig.\ref{fig:non-linear_channel}a demonstrates how the MIND decoder has managed to implicitly learn the non-linear channel model during the training phase and effectively use such information during decoding. A comparison with the MaxL decoder with and without perfect channel state information (CSI) knowledge is also conducted. Results show that MIND exhibits performance close to optimality. Fig. \ref{fig:non-linear_channel}b illustrates both the behaviour of the entropies and the average mutual information as function of the SNR.

\subsection{Additive Non-gaussian Noise Channel}
Lastly, and perhaps as most interesting example, we consider a short block coded transmission over an additive non-gaussian channel. The aim is to assess the ability of MIND to learn and exploit the presence of non-gaussian noise. In particular, we suppose that the noise follows a truncated Middleton distribution \cite{Middleton1977}, also called Bernoulli-Gaussian noise model, so that at any given time instant $k$ it is obtained as $n_k = (1-\epsilon_k)n_{1,k}+\epsilon_k n_{2,k}$ where $n_{1,k} \sim \mathcal{N}(0,\sigma_b^2)$ is a zero-mean Gaussian random variable with variance $\sigma_b^2$ and $n_{2,k} \sim \mathcal{N}(0,B\sigma_b^2)$ is also a zero-mean Gaussian random variable but with variance $B$ times larger. Instead, $\epsilon_k$ is a Bernoulli random variable with probability of success $P$. The pdf of the noise samples is then given by
\begin{equation}
p_{N}(\mathbf{n}_k) = (1-P)\mathcal{N}(0,\sigma_b^2)+P\mathcal{N}(0,B\sigma_b^2).
\end{equation} 
Assuming that the noise model is known and that BPSK symbols are transmitted, two decoding strategies can be devised. The first, denoted with MaxL Middleton, uses maximum likelihood decoding with the known conditional pdf $p(\mathbf{y}|\mathbf{x})=p_{N}(\mathbf{y}-\mathbf{x})$. The second strategy is a genie decoder that knows the outcome of the Bernoulli event for every time instance. That is, it knows whether a received sample is hit by Gaussian noise with variance $\sigma_b^2$ or $B\sigma_b^2$. A third decoding strategy is offered by MIND that learns the channel statistics.
We distinguish among three different types of codes: a) a binary repetition code with length $5$; b) a $(7,4)$ Hamming code; c) a rate $1/2$ convolutional code with memory $2$ and block-length $18$. For each of them, the bit error rate (BER) obtained with the genie, the MaxL Middleton and the MIND decoders is reported. The parameter $B$ was set to $5$ in all the experiments involving Middleton noise, wherein we also set $P=0.05$. 
Fig.\ref{fig:middleton} shows the gain provided by the neural-based decoder scheme over the classical maximum likelihood one that exploits the Middleton distribution in (24). With MIND, the BER performance gets closer to the genie decoder. Furthermore, Fig.\ref{fig:middleton} also reports an estimate of the average mutual information provided by MIND.

\begin{comment}
\begin{enumerate}
\item Uncoded transmission in additive Gaussian noise with block transmission of binary PSK data symbols that are dependent; this is to show the ability of the decoder to learn and exploit data dependence at the source. The MI decoder is compared with the optimal MAP decoder that knows (a-priori) that joint distribution of the block of transmitted data symbols $p(\mathbf{x})$. In addition the ML decoder that assumes i.i.d. transmitted data symbols is considered. It should be noted that data dependency at the source can be either created with a coding stage from a source that emits i.i.d. symbols, or intrinsically by the source of data traffic having certain semantics.
\item Uncoded transmission in a non-linear channel with additive Gaussian noise; this is to show the ability of the MI decoder to discover such a channel non linearity. We still assume to transmit i.i.d. 4-PAM symbols. However, the channel introduces an non linearity (for instance because of the presence of non linear amplifiers). Thus, the considered channel model is $y_i=sign(\mathbf{x}_i)\sqrt{|\mathbf{x}_i|}+n_i$.
\item Uncoded transmission in an additive Middleton noise channel; this is to assess the ability of the MI decoder to learn and exploit the presence of non Gaussian noise. A comparison in terms of BER is made between the MI decoder, the genie optimal decoder that knows the noise statistics, and the one that assumes Gaussian noise. The considered modulation is binary PSK with equally probable symbols.

\item Short block coded transmission in additive white Gaussian noise. In particular, a rate $1/2$ convolutional code with memory 2 is assumed. The input block has length 20 bits. In addition, the distribution of the source of $2^{20}$ information blocks is not uniform but Poissonian. A comparison between the MI decoder and the Viterbi decoding algorithm is made.
\end{enumerate}
\end{comment}

\section{Conclusions}
\label{sec:conclusions}
In this paper, MIND, a neural decoder that uses the mutual information as decoding criterion, has been proposed. Two specific architectures have been described and they are capable of learning the conditional information $i(\mathbf{x}|\mathbf{y})$ of the codeword $\mathbf{x}$ given the channel output observation $\mathbf{y}$ in unknown channels. Additionally, MIND allows the estimation of the achieved information rate with the used coding scheme as well as the decoding error probability. Several numerical results obtained in illustrative channels show that MIND can achieve the performance of the genie MAP decoder that perfectly knows the channel model and source distribution. It outperforms the conventional max-log decoder that assumes the presence of Gaussian noise.        

\bibliographystyle{IEEEtran}

\section{Results}
\label{sec:results}
In the experimental results, we consider the transmission of $M$ messages at rate $R$ over an AWGN channel, for which we know the channel capacity under an
average power constraint. The choice of AWGN channel shall not be seen as a limitation since the formulation transcends the channel characteristics. Moreover, it provides a reference curve for the validation of the estimators and it allows comparisons across several application domains \cite{Mine2018, LetiziaNIPS}. 

We describe two scenarios: a) the transmission at a rate $R>1$ and short code-length $n$ and b) the transmission at a low rate with larger code-length. For both scenarios, we train a rate-driven autoencoder (see \cite{Letizia2021}) and compare its performance for different estimators in terms of BLER and accuracy of the mutual information estimation. 

For all the experiments and as a proof of concept, we use simple encoder, decoder and neural mutual information estimator architectures. In particular, the encoder is a shallow neural network with $M$ neurons in the input layer, $M$ in the hidden one and $2n$ neurons at the output while the decoder is its complementary.  
The discriminative estimator possesses instead a fixed architecture, a
two layer multilayer perceptron neural network with $200$
neurons in each layer with ReLU and sigmoid activation functions. The only
difference in the estimators resides in the final layer where we use a linear dense layer for MINE, while a softplus layer for the DIME based estimators, defined as
\begin{equation}
 D^*(x,y; \theta) = \log(1+\exp(W^*_L \cdot z_{L-1}+b^*_L)),
\end{equation}
where $W$ and $b$ are weights and biases, respectively, $z$ is the output of the previous layer and $L$ is the number of layers of the network.
\begin{figure}
\centering

  	\includegraphics[scale=0.255]{pics/ber_all_1}
  	\caption{Comparison of the BLER obtained by the autoencoder AE($3,18$) and AE($6,3$) with an AWGN intermediate layer, using different estimators with $\beta = 0.2$, $\epsilon = 0.2$.}
  	\label{fig:ber_all}
\end{figure}
The training of the autoencoder was conducted at a fixed $E_b/N_0$ ratio of $7$ dB. We minimized the loss function alternating encoder and decoder weights update with the discriminator weights update. The number of iterations was set to $10$k with a learning rate of $0.01$ for both the components. For more implementation details, we refer to the code publicly available. \footnote{\url{https://github.com/tonellolab/capacity-approaching-autoencoders}}

The first scenario describes a digital transmission scheme with input alphabet size $M=64$ over $3$ channel uses. We denote such system as AE($6,3$) since $k = \log_2(M)$. The communication rate is $R=2$. Fig.~\ref{fig:ber_all} shows the autoencoder performance in terms of BLER for different estimators. From the curves we can infer that the $\gamma$-DIME estimators, and in particular, higher values of $\gamma$ seem to provide improved performance in terms of BLER. The result of Fig.~\ref{fig:MI_6_6} is insightful: none of the estimators provides stable estimations, especially at high SNRs where it is expected that the mutual information saturates to the information rate $R=2$. The DIME based estimators are more accurate compared to the divergent MINE. Nevertheless, the greater is $\gamma$, the greater is the variance (see theorem 2 of \cite{Song2020}) of the estimation. It is also interesting to notice that the concavity of $J_{\gamma}$ (shown in Fig.~\ref{fig:gammaDIME}) influences the mutual information estimation as it appears in the results. Indeed, for low values of $\gamma$, fixed learning rate and training iterations, the estimators are trained by moving on the flat curves of Fig.~\ref{fig:gammaDIME}, which is equivalent as using a small learning rate to upgrade the gradients. For completeness and to show the advantage of the proposed estimators under fading, we report in Fig.~\ref{fig:MI_6_6_Ray} the achieved estimated information rate with Rayleigh fading. Also in this case, MINE diverges for high SNR values while the DIME estimators are much more accurate and tend to saturate to the rate $R=2$. 

\begin{figure}
\centering

  	\includegraphics[scale=0.255]{pics/MI_6_6_1}
  	\caption{Comparison of different mutual information estimators for the autoencoder AE($6,3$) with $\beta = 0.2$, $\epsilon = 0.2$ on the AWGN channel.}
  	\label{fig:MI_6_6}
\end{figure}

\begin{figure}
\centering

  	\includegraphics[scale=0.255]{pics/MI_6_6_Ray_1}
  	\caption{Comparison of different mutual information estimators for the autoencoder AE($6,3$) with $\beta = 0.2$, $\epsilon = 0.2$ on the Rayleigh fading channel.}
  	\label{fig:MI_6_6_Ray}
\end{figure}

The second scenario describes a digital transmission scheme with input alphabet size $M=8$ over $9$ channel uses. We denote such system as AE($3,9$). The communication rate is $R=1/3$. Fig.~\ref{fig:ber_all} illustrates the advantage of using MINE for low values of the mutual information. Under the same training iterations, MINE is a tighter bound compared to $\gamma$-DIME and therefore converges faster and provides stronger gradients for the encoder. However, Fig.~\ref{fig:MI_3_18} highlights the fact that for low SNRs and alphabet dimension and for relatively small $n$, the estimators perform similarly. We can conclude that the choice of the estimator to use varies according to input constraints such as alphabet dimension, code-length and SNR. We did not comment about the architecture to use although we remark that it is one of the main aspects in modern machine learning applications and therefore we expect significant improvements. 

\begin{figure}
\centering

  	\includegraphics[scale=0.255]{pics/MI_3_18_1}
  	\caption{Comparison of the achieved information rate using different mutual information estimators for the autoencoder AE($3,18$) with $\beta = 0.2$, $\epsilon = 0.2$ on the AWGN channel.}
  	\label{fig:MI_3_18}
\end{figure}
\section{Discriminative Mutual Information Estimators}
\label{sec:theory}
The mutual information $I(X;Y)$ quantifies the statistical dependence between two random variables, $X$ and $Y$, by measuring the amount of information of $X$ via observation of $Y$. Unlike the correlation, the mutual information measures linear and non-linear associations \cite{Letizia2020} since its estimation requires the knowledge of three distributions
\begin{equation}
I(X;Y) = \mathbb{E}_{(x,y)\sim p_{XY}(x,y)}\biggl[\log\frac{p_{XY}(x,y)}{p_X(x)p_Y(y)}\biggr],
\end{equation}
where $p_X(x)$ and $p_Y(y)$ are the distributions of $X$ and $Y$, respectively, and $p_{XY}(x,y)$ is the joint probability distribution. The presence of these three components renders the mutual information estimation a challenging task and several approaches have been proposed to tackle the problem. Standard and traditional approaches rely on binning, density and kernel estimation \cite{Moon1995} and $k$-nearest neighbours \cite{Kraskov2004}. However, for applications such as channel coding, high-dimensional data is often involved and classic approaches fail to scale. Moreover, they can not easily be integrated in an optimization framework that targets the channel capacity. To overcome such limitations, neural networks can be used to obtain a parametric representations of the variational lower bounds on the mutual information, where the objective of the network is to estimate the density ratio. 

One of the first examples of neural estimators is MINE \cite{Mine2018}. The idea behind MINE is to exploit the Donsker-Varadhan dual representation of the Kullback-Leibler (KL) divergence to produce the bound to optimize
\begin{align}
\label{eq:MINE}
I(X;Y) \geq & I_{MINE}(X;Y) \nonumber \\
= & \sup_{\theta \in \Theta} \mathbb{E}_{(x,y)\sim p_{XY}(x,y)}[T_{\theta}(x,y)]  \nonumber \\
& - \log(\mathbb{E}_{(x,y)\sim p_X(x) p_Y(y)}[e^{T_{\theta}(x,y)}]),
\end{align}
where $\theta \in \Theta$ parameterizes a family of functions $T_{\theta} : \mathcal{X}\times \mathcal{Y} \to \mathbb{R}$ via a deep neural network. MINE is a biased estimator, consequently, the authors suggested to replace the expectation in the denominator of the gradient with an exponential moving average. However, MINE suffers from high-variance estimations and a recent solution named SMILE \cite{Song2020} tries to alleviate such issue by clipping the density ratio between $-e^{-\tau}$ and $e^{\tau}$, with $\tau \geq 0$. SMILE is equivalent to MINE in the limit $\tau \rightarrow +\infty$.

A mutual information unbiased variational lower bound relies on the $f$-divergence representation proposed in \cite{Nguyen2010}
\begin{align}
\label{eq:NWJ}
I(X;Y) \geq & I_{NWJ}(X;Y) \nonumber \\
= & \sup_{\theta \in \Theta} \mathbb{E}_{(x,y)\sim p_{XY}(x,y)}[T_{\theta}(x,y)]  \nonumber \\
& -\mathbb{E}_{(x,y)\sim p_X(x) p_Y(y)}[e^{T_{\theta}(x,y)-1}].
\end{align}
where it is easy to verify that $I_{NWJ}\leq I_{MINE}$.

In \cite{LetiziaNIPS}, the discriminative mutual information estimator (DIME) was proposed. It is a neural estimator that mimics the task of the discriminator in GANs \cite{Goodfellow2014}. Indeed, the optimal discriminator in the GAN or $f$-GAN framework can be used either as a direct or indirect mutual information estimator when it is fed with paired and unpaired input data of distribution $p_{XY}$ and $p_X p_Y$, respectively. We refer to as direct estimator, the one where the optimal discriminator $D^*$ is proportional to the density ratio in \eqref{eq:density_ratio}. An example of this is the d-DIME estimator, a mutual information variational lower bound that reads as follows
\begin{equation}
\label{eq:d_DIME}
I(X;Y) \geq \tilde{I}_{dDIME}(X;Y) = \frac{\mathcal{J}_{\alpha}(D^*)}{\alpha}+1-\log(\alpha),
\end{equation}
where $\mathcal{J}_{\alpha}(D)$, $\alpha>0$, is a value function defined as
\begin{align}
\mathcal{J}_{\alpha}(D) = \; & \alpha \cdot \mathbb{E}_{(x,y) \sim p_{XY}(x,y)}\biggl[\log \biggl(D\bigl(x,y\bigr)\biggr)\biggr] \nonumber \\ 
& +\mathbb{E}_{(x,y) \sim p_{X}(x)p_{Y}(y)}\biggl[-D\bigl(x,y\bigr)\biggr],
\label{eq:discriminator_function}
\end{align}
and
\begin{equation}
\label{eq:optimal_discriminator_2}
D^*(x,y) = \alpha \cdot \frac{p_{XY}(x,y)}{p_{X}(x)\cdot p_Y(y)} = \arg \max_D \mathcal{J}_{\alpha}(D).
\end{equation}
When the positive function $D(\cdot)$ is parameterized by a neural network, the NWJ estimator can be thought as a special case of d-DIME with $\alpha=1$ and $T_{\theta}(x,y) = \log(D_{\theta}(x,y))+1$.

In the following section we extend DIME, in particular, we introduce $f$-DIME, a discriminative mutual information estimator based on the $f$-divergence measure.   

\section{$f$-DIME}
\label{sec:f-DIME}
In this section, we describe a general methodology to estimate the mutual information by applying the variational representation of $f$-divergence functionals $D_f(P||Q)$. In detail, if the measures $P$ and $Q$ are absolutely continuous w.r.t. $\diff x$ and possess densities $p$ and $q$, then the $f$-divergence reads as follows
\begin{equation}
D_f(P||Q) = \int_{\mathcal{X}}{q(x)f\biggl(\frac{p(x)}{q(x)}\biggr)\diff x},
\end{equation}
where $\mathcal{X}$ is a compact domain and the function $f:\mathbb{R}_+ \to \mathbb{R}$ is convex, lower semicontinuous and satisfies $f(1)=0$. Under these conditions, the authors of \cite{Nguyen2010} (and \cite{Nowozin2016} for the $f$-GAN) exploited the Fenchel convex duality to derive a lower bound on $D_f$
\begin{equation}
\label{eq:f_bound}
D_f(P||Q) \geq \sup_{T\in \mathbb{R}} \biggl\{ \mathbb{E}_{x \sim p(x)} \bigl[T(x)\bigr]-\mathbb{E}_{x\sim q(x)}\bigl[f^*\bigl(T(x)\bigr)\bigr]\biggr\},
\end{equation}
where $T: \mathcal{X} \to \mathbb{R}$ and $f^*$ is the Fenchel conjugate of $f$ defined as
\begin{equation}
f^*(t) := \sup_{u\in \mathbb{R}} \{ ut -f(u)\}.
\end{equation}
Therein, it was shown that the bound in \eqref{eq:f_bound} is tight for optimal values of $T(x)$ and it takes the following form
\begin{equation}
\label{eq:optimal_ratio}
T^*(x) = f^{\prime} \biggl(\frac{p(x)}{q(x)}\biggr),
\end{equation}
where $f^{\prime}$ is the derivative of $f$.

To create a novel neural estimator based on the representation in \eqref{eq:f_bound}, we firstly substitute the mutual information $I(X;Y)$ with its KL divergence representation 
\begin{equation}
I(X;Y) = D_{KL}(p_{XY}||p_X p_Y),
\end{equation}
and by defining 
\begin{align}
\mathcal{J}_{f}(T) = \; & \mathbb{E}_{(x,y) \sim p_{XY}(x,y)}\bigl[T\bigl(x,y\bigr)\bigr] \nonumber \\ 
& -\mathbb{E}_{(x,y) \sim p_{X}(x)p_{Y}(y)}\bigl[f^*\bigl(T\bigl(x,y\bigr)\bigr)\bigr],
\label{eq:value_function_f}
\end{align}
we propose to parametrize $T(x,y)$ with a deep neural network $T_{\theta}$ of parameters $\theta$ and solve with gradient ascent and back-propagation
\begin{equation}
\theta^* = \arg \max_{\theta} \mathcal{J}_f(T_{\theta}).
\end{equation}
The main idea is the following: \textit{it is possible to exploit the relation \eqref{eq:optimal_ratio} to estimate the mutual information}. Indeed, if the derivative of $f$ is invertible, 
\begin{equation}
\bigl(f^{\prime}\bigr)^{-1}\bigl(T_{\theta^*}(x,y)\bigr) = \frac{p_{XY}(x,y)}{p_X(x)p_Y(y)},
\end{equation}
thus, having access to $T_{\theta^*}$ and using the density ratio estimated above, we obtain $f$-DIME as
\begin{equation}
\label{eq:f_dime}
I_{fDIME}(X;Y) = \mathbb{E}_{(x,y) \sim p_{XY}(x,y)}\biggl[ \log \biggl(\bigl(f^{\prime}\bigr)^{-1}\bigl(T_{\theta^{*}}(x,y)\bigr) \biggr) \biggr].
\end{equation}

We argue that any value function $\mathcal{J}_f$ of the form in \eqref{eq:value_function_f}, seen as dual representation of a certain $f$-divergence $D_f$, can be maximized to estimate the mutual information using \eqref{eq:f_dime}. As a notable example, the GAN-based i-DIME estimator in \cite{LetiziaNIPS} can be derived using $f$-DIME with generator 
\begin{equation}
f(u) = u\log u-(u+1)\log(u+1)+\log4
\end{equation}
and $T(x)=\log(D(x))$.
However, only a subset of $\mathcal{J}_f$ can be used inside the capacity-driven autoencoder loss function since a maximization of the mutual information is also required. In particular, only the generator $f$ of the KL divergence is suitable for the mutual information maximization, since in such case
\begin{equation}
\max_{p_X(x)} D_f = \max_{p_X(x)} D_{KL} = \max_{p_X(x)} I(X;Y).
\end{equation}

The generator function of the KL divergence is $f(u) = u\log u$ with conjugate $f^*(t) = e^{t-1}$. It easy to notice that such choice of $f$ coincides with the NWJ estimator or d-DIME when $T(x) = \log(D(x))$ and $\alpha=1$ (see \eqref{eq:d_DIME}). Nevertheless, as shown in \cite{LSGAN}, the value function plays a fundamental role during the training process of a discriminator. Consequently, we combined the approach described for the derivation of $f$-DIME with d-DIME to obtain a family of lower bounds on the mutual information that may result in more robust estimators according to specific case studies.

\begin{lemma}
\label{lemma:Lemma1}
Let $X\sim p_X(x)$ and $Y\sim p_{Y}(y|x)$ be the channel input and output, respectively. Let $D(\cdot)$ be a positive function. If $\mathcal{J}_{\gamma}(D)$, $\gamma>0$, is a value function defined as 
\begin{align}
\mathcal{J}_{\gamma}(D) = \; & \gamma \cdot \mathbb{E}_{(x,y) \sim p_{XY}(x,y)}\biggl[\log \biggl(D\bigl(x,y\bigr)\biggr)\biggr] \nonumber \\ 
& + \mathbb{E}_{(x,y) \sim p_{X}(x)p_{Y}(y)}\biggl[- D^{\gamma}\bigl(x,y\bigr)\biggr],
\end{align}
then
\begin{equation}
I(X;Y) \geq \tilde{I}_{\gamma DIME}(X;Y) = \mathcal{J}_{\gamma}(D^*)+1,
\end{equation}
where
\begin{equation}
\label{eq:optimal_discriminator_gamma}
D^*(x,y) = \biggl(\frac{p_{XY}(x,y)}{p_{X}(x)\cdot p_Y(y)}\biggr)^{1/\gamma} = \arg \max_D \mathcal{J}_{\gamma}(D).
\end{equation}
\end{lemma}
\begin{comment}
\begin{proof}
Consider a scaled generator $f(u) = \frac{u}{\gamma}\log u$ and for simplicity of notation, denote $p_{XY}$ and $p_Xp_Y$ with $p$ and $q$, respectively. Then
\begin{equation}
D_{KL}(p||q) = \gamma \int_{x}{q(x) \frac{p(x)}{\gamma q(x)}\log\biggl(\frac{p(x)}{q(x)}\biggr) \diff x},
\end{equation}
with conjugate $f^*(t)$, with $t \in \mathbb{R}$ given by
\begin{equation}
f^*(t) = \frac{e^{\gamma t -1}}{\gamma}.
\end{equation}
Substituting in \eqref{eq:f_bound} yields to
\begin{equation}
D_{KL}(p||q) \geq \gamma \sup_{T\in \mathbb{R}} \biggl\{ \mathbb{E}_{x \sim p(x)} \bigl[T(x)\bigr]-\frac{1}{\gamma}\mathbb{E}_{x\sim q(x)}\bigl[e^{\gamma T(x)-1} \bigr]\biggr\}.
\end{equation}
Using \eqref{eq:optimal_ratio} it is easy to verify that the optimal value of $T$ is the log-likelihood ratio rather than the density ratio itself. Indeed,
\begin{equation}
T^*(x) = \frac{1}{\gamma}\log\biggl(\frac{p(x)}{q(x)}\biggr) + \frac{1}{\gamma}.
\end{equation}
Finally, with the change of variable $T(x)=\log(D(x))+1/{\gamma}$, the optimal discriminator has form
\begin{equation}
D^*(x) = \biggl(\frac{p(x)}{q(x)}\biggr)^{1/\gamma}
\end{equation}
and
\begin{align}
D_{KL}(p||q) \geq \; & \sup_{D\in \mathbb{R}_+} \biggl\{ \gamma \mathbb{E}_{x \sim p(x)} \bigl[\log \bigl( D(x) \bigr) \bigr] \nonumber \\
& -\mathbb{E}_{x\sim q(x)}\bigl[D^{\gamma}(x) \bigr]\biggr\} + 1.
\end{align}
Denoting the two terms inside the supremum with $J_{\gamma}(D)$ and replacing $p$ and $q$ with the joint and marginal distributions, concludes the proof.
\end{proof}
\end{comment}
\begin{figure}
\centering
  	\includegraphics[scale = 0.255]{pics/gammaDIME}
  	\caption{Behaviour of the functional $J_{\gamma}$ varying the parameter $\gamma$. Concavity of the neighbourhood of the maximal value depends on $\gamma$.}
  	\label{fig:gammaDIME}
\end{figure}

Fig.~\ref{fig:gammaDIME} illustrates how the value function in \eqref{eq:value_function_f} (divided by $\gamma$) behaves according to the choice of the positive parameter $\gamma$. The position of the maximum depends on the density ratio as described in \eqref{eq:optimal_discriminator_gamma} ($R = 1$ in the figure). The role of $\gamma$ is to modify the concavity of the value function to improve the convergence rate of the gradient ascent algorithm. Moreover, when $D$ is parameterized by a neural network, the choice of the learning rate and activation functions may be related with $\gamma$. Indeed, notice that $\gamma$-DIME can also be derived from d-DIME using the substitution $D(x,y) \rightarrow D^{\gamma}(x,y)$.