DUNE Software and High Performance Computing

DUNE, like other HEP experiments, faces a challenge related to matching execution patterns of our production simulation and data processing software to the limitations imposed by modern high-performance computing facilities. In order to efficiently exploit these new architectures, particularly those with high CPU core counts and GPU accelerators, our existing software execution models require adaptation. In addition, the large size of individual units of raw data from the far detector modules pose an additional challenge somewhat unique to DUNE. Here we describe some of these problems and how we begin to solve them today with existing software frameworks and toolkits. We also describe ways we may leverage these existing software architectures to attack remaining problems going forward. This whitepaper is a contribution to the Computational Frontier of Snowmass21.

Introduction
DUNE [1] software faces many software development and data processing challenges shared with other large HEP experiments. Moore's law has morphed from delivering ever faster CPUs to delivering ever more numerous and sometimes even slower cores. This is epitomized by GPUs which provide O(1k) cores per card but which are clocked far slower than even low-end CPUs. We must develop software that can efficiently exploit the highly parallel architectures of high-performance computers (HPC) while still making efficient use of the more familiar high-throughput computing (HTC) facilities that offer modest parallelism of O(10) CPU cores per compute node.
Contemporaneously, the available RAM capacity associated with CPU and GPU is not growing as fast as their core counts and RAM/core ratio has become a limiting parameter for many types of jobs. Where possible we must reimplement and rearchitect to reduce memory waste. On the other hand, many classes of jobs have an unavoidable memory overhead. Where it can make it safely shared across multiple threads we may exploit the typical HTC policy whereby RAM allocation is tied to CPU core allocation. Thus, we must convert our software from single-threaded (ST) to multi-threaded (MT) to alleviate memory pressure as a side-effect.
Introduce a MT execution model then brings new problems related to variations in the number of cores that can be engaged at any one point in the process life cycle. Such dips in CPU utilization matter for HTC if not for HPC. In worse case, we must allocate N HTC cores to gain enough CPU RAM to solve the above memory problems yet run ST wasting N − 1 cores. Some algorithms are amenable to running in the highly SIMD-parallel architecture of GPUs. AI/ML inference, FFTs and other large but "simple" algorithms in particular are greatly accelerated and require fairly simple software development to provide. While GPU offloading is welcome, it brings another problem similar to the N − 1 core problem but which strikes both CPU and GPU hardware. For GPU to be highly utilized, a requirement for gaining access to HPC, current software must bring to bear O(100) CPU cores for each GPU device. HPC allocations tend to be on a per-compute node where 4, 8 or 18 GPU devices are provided with O(100) CPU cores or less. Thus we must find ways to increase the GPU/CPU utilization ratio in our software. Turning this around, in some HPC or some dedicated HTC+GPU facilities we may have a small GPU/CPU hardware ratio. In these cases, high CPU utilization is again an important metric. With simple software architectures, a CPU thread that submits a GPU task will cause a core to go idle while it waits for a return. In order to absorb fluctuations in the rate of GPU task submission some task queue is required in order to avoid exhausting GPU RAM and the wait time for servicing will be come longer at the exact moment that CPUs are most overloaded. To combat this, an advanced execution model is required which is more asynchronous in nature than currently.
Given the complexity of our algorithms and diversity of our developers experience and coordination a third problem naturally arises. Different developers are free to create MT solutions based on different threading libraries. For example, one may use std::thread and another may use TBB. The variety of thread pools are in general not cohesive. That is, an allocation of N cores can not be simply shared between the pools. This requires the process to partition N between the pools which decreases core sharing and increases idleness. Alternatively, the process may place no internal partition and this risks usage beyond the allocation. Depending on facility policy this may be acceptable or may lead to the process being killed.
The above issues can largely be ignored if jobs were not judged against full utilization of the computing resources. Only when we attempt to obtains high resource utilization of CPU and GPU while also fitting our jobs into available RAM do we face this otherwise unwanted challenges. But, face them we must if we wish to exploit current HTC and HPC systems.
2 Current DUNE execution model DUNE offline software uses the art framework [2] which enacts an event-based execution model. Code units (art "modules") may put and get data from the art "event" throughout the "event loop" over which they are executed. The scope of one "event" is application-dependent. That is, it is a choice to be made by the collaboration using art.
Traditionally, the choice is to associate an "event" with a detector trigger or with one simulated particle interaction. This is appropriate for collider experiments but ends up being very harmful when applied to LArTPC detectors, particularly those composed of multiple, independent sub-detectors. Treating a full detector readout as an "event" produces prohibitively large memory resident data even in "small" detectors such as those from ProtoDUNE, SBND and ICARUS which "only" have O(10 − 50)k channels.
Typical DUNE jobs today are composed as a linear sequence of modules, each executed serially by art and internally utilizing a single CPU thread. Parallel execution of a number of these linear "paths" of modules is also supported. The TBB library provides the basis for these MT operation and thus any module that internally also utilizes TBB for any MT operations avoids the "thread pool" problem described. Almost no DUNE modules employ internal MT code. One notable exception is described in the next section.
The single art "event" data structure is shared by all modules in a given job with thread-safe access assured. In addition, over the job's life, art provides points of full synchronization among the modules at the boundaries of "sub run" groups of events. The span of a sub run is part of the application data model. art also supports parallel execution over multiple events within a sub run. A single instance of one module may run concurrently on different threads or multiple instances of one module may each run individually on their given thread. Modules run in these ways must assure they access any shared data in a thread-safe manner.
The DUNE Computing Consortium is currently in the process to define requirements for its ultimate offline software framework. Some of these issues and other described below are being woven into these requirements. The final decision may to continue to use art or something else. At the same time, a green field development is recognized as a major undertaking and would come with a host of risks.

More parallelism with Wire-Cell Toolkit
The Wire-Cell toolkit [3] (WCT) provides a different execution model which is orthogonal to that enacted by art and is also somewhat synergistic. It is based on defining code units (called "components" instead of "modules" as with art) as executable nodes in a data-flow graph.
One such graph is rendered from a job configuration in Figure 1. Each node in the graph is one instance of a particular type which defines the number of input and/or output "ports" and the types of data they pass. Graph nodes are connected by edges between ports and every port must have exactly one edge for the graph to be considered complete and executable. Special code units which are shown as connected through gray lines are "services" which do not participate in the data flow execution but rather provide methods called by the nodes. Many of these calls return static data and the few that manage mutable data are made thread-safe with mutex or semaphore patterns.  A graph execution engine provides a policy to govern the scheduling of node executions. As part of any execution policy, a data object produced on an output port is transferred across the edge and made ready for subsequent input by the other connected node. The engine is otherwise free to apply any strategy toward deciding when to execute any given node. An engine implements an abstract interface and this allows novel engines to be engaged simply by naming it in the user configuration. WCT provides two engine implementations: an ST engine with a policy that minimizes memory utilization and an MT engine based on TBB flow_graph library.
The MT engine will execute nodes in parallel using up to a configured maximum number of threads. This allows flexibility to utilize otherwise wasted cores allocated for their associated RAM or to assure enough CPU processing is engaged in order to provide tasks fast enough to assure a GPU is well utilized. In the example of Figure 1, we see six parallel pipelines, each associated to one sub-detector unit. When a single "event" worth of data is allowed into the graph at one time then at most six threads may be fully utilized.
In addition to this "transverse" parallelism, a certain "longitudinal" or pipelined-parallelism may be allowed to develop in the MT flow graph which can allow higher thread utilization. This form of parallelism will develop when multiple "events" of data are allowed into the graph at any given moment.
The potential complexity of object lifetime and memory management is well controlled. Any object transferred across a graph edge is passed via shared pointer and otherwise immutable. This removes explicit memory management duties from the node implementations. It also allows thread-safe shared data access by the nodes. Guards such as std::mutex are not used and their overhead and potential to block progress is simply avoided.

Combining art and WCT
As a toolkit, Wire-Cell is designed to run as part of a larger application. The "larwirecell" package has been developed as a part of the LArSoft project, used by most LArTCP based experiments, to provide an interface between art and WCT. It provides an art "tool" which is used by an art module to run arbitrary WCT flow graphs. It also includes a set of WCT components to bridge data between art and WCT worlds. Finally, it provides a conduit from the art configuration mechanism based on the FHiCL language to that of WCT based on objects typically defined in Jsonnet.
In an accidental but very happy synergy, the WCT MT flow graph execution engine is based on the same TBB thread library as used by art.
Though not yet exploited in practice, it enables art path-level MT to be employed in processes that include WCT while avoiding the "multiple thread pool" problem. However, when a MT WCT flow graph is run with ST art modules, we must face the N − 1 problem. We may run the WCT ST flow graph engine or the MT with a single thread in order to fit the HTC allocation. Or, in the case where multiple cores must be allocated in order to gain their associated memory, we may run the MT engine with a matching number of threads to make good use of otherwise idle cores.
In other jobs where the WCT stage dominates the processing time we become free to allocate a higher number of cores without fear of them becoming idle during execution of ST modules. art support for multiple events in flight and the throttling provided by sub run synchronization points may be a perfect combination in such jobs and their use should be explored. Together they may provide a means to maximize thread usage by enabling pipelined-parallelism while also regulating the pipelined depth and thus the memory usage.

Sub-event parallelism and I/O
DUNE far detector modules, like other detectors, are structured as an array of independent and functionally identical sub-detector units (eg, the 150 DUNE APAs in each 10 kT far detector module). WCT reflects this structure into its graphs as we see in the pipelines of 1. When data from "small" detectors like ProtoDUNE-SP are run in current art jobs, these graphs have a singular upstream subgraph which accesses the art event store and then fans out to the graph. This leads to a ST choke point and to unwanted memory usage to hold the entire input data while only some portion needs be in memory for immediate consumption. The use of abstract interface classes for WCT data objects (as well as components) allows some amelioration of this problem. For example, special lazy-loading data objects were developed to allow ProtoDUNE-SP data processing to avoid the ruinous memory consumption that comes with the standard, monolithic data loading policy.
At the time, the lazy loading trick was required as only the ST engine was in regular use. Today we may solve the problem by allocating multiple cores and their RAM and using the MT engine to assure their full utilization. But, scaling this strategy to the DUNE far detector with its 25 times more APAs than ProtoDUNE will require yet novel solutions.
These solutions are expected to be based on rejecting the monolithic data loading policy entirely and exploiting the sub-detector granularity of the DUNE far detector data. And, we must support the fact that DUNE far detector events may be sparse across the array of sub-detectors. For example, advanced DAQ trigger and readout is expected to deliver at most empty marker objects for otherwise empty readouts from APAs that lack enough "interesting" activity to pass DAQ criteria.
Thus we seek some mechanism to enable sub-event level loading that can enable MT parallelism both "across" and "along" the data-flow graph. One way to achieve this pattern is to load data from file directly via a WCT "source" node as illustrated in Figure 2. This may have particular synergy with the parallel file structure expected from the DUNE DAQ as described above. We may consider developing jobs with O(100) input nodes in the WCT flow graph, each reading in data from one APA. The current leading candidate for the raw data format is HDF5 and so may exploit both inter-file and intra-file parallel I/O methods.  Figure 2: A high-level cartoon illustration of a possible Wire-Cell Toolkit data flow graph allowing for independent input of multiple raw detector data files followed by their signal processing and a single synchronized output.
One of the first tasks applied to far detector data is "signal processing" which is implemented as a per-APA WCT component. Figure 2 shows one possible context for this task. The flow graph utilize parallel input streams, is composed of parallel signal processing pipelines and produce parallel output streams. Without a single output node to synchronize the streams, such graph can be decomposed into an array of "embarrassingly parallel" processes, each running one pipeline between its independent input and output files.
However, the output data volume of signal processing is greatly reduced compared to the raw data input and further processing soon comes after (if not directly) which must operate on data from multiple or even all APAs together. Thus, the cost of placing a synchronization node at the output of the array of pipelines must be born at some point. Placing it just after the signal processing stage is well motivated.
To keep such a graph as a single executable image requires the host computer to supply enough CPU cores, potentially O(100) for one DUNE far detector module. We will discuss strategies next for how we may still utilize computing facilities where this density of CPU cores is not provided. The ability to factor such large graphs in flexible ways is expected to be particularly important in processing supernova neutrino burst candidate data from DUNE far detector modules. A single such "event" from one 10 kT module is expected to span 100 seconds of readout and produce 115 TB of raw, packed data. This enormous amount of data will be split both along APA sub-detector lines and in chunks of time. For example if split into 100 files of 1 second each and for each APA produces a bit under 8 GB files. Likely each second worth of data would be read in time slices of O(1) millisecond.
Thus, we may map the data to a variety of processing patterns in flexible ways at either the data-flow graph node or the process level. This mapping can be tuned to best meet limitations imposed by a given computing facility allocation.

Enter GPUs
Introduced above are flow-graph nodes for performing "signal processing". Another critical set of WCT nodes are those that perform tasks related to ionization electron drift and field response and noise simulations. Execution of these tasks is dominated by a few algorithms including FFT, scatter-add and AI/ML inference.
These code "hot spots" were initially provided as CPU-based implementations. Work performed in the context of HEP-CCE/PPS [4] have ported these simulation hot spots to GPU. Initial implementations were based on CUDA. Subsequently, more full-featured and portable implementations are based on Kokkos [5]. AI/ML inference is performed via libtorch where we again have the option to configure to utilize CPU or GPU.
Through simple configuration changes we may now select implementations which may best fit the resources provided by the hosting computers. The current benchmarks show that GPU based simulation runs 18× faster [6], [7] than CPU, bringing the time to simulate one APA readout of 5ms of data to less than one second. The signal processing is still largely implemented with CPU code while GPU can accelerate the FFTs and AI/ML inference [8]. There, GPU running is still a more modest O(2×) faster than CPU running. Plans are in place to reimplement more of the signal processing on GPU so we may expect to eventually achive performance that has some parity with that achieved for simulation.
While the simulation can be made GPU-heavy, it still requires many CPU cores to provide GPU tasks fast enough to keep a GPU fully utilized. This means the job is subject to the problem of load balancing many CPU cores generating tasks just fast enough for the GPU to absorb those tasks at near full utilization. This balancing must also assure that the GPU is not overloaded resulting in out-of-memory errors.
Load balancing is provided today in a rather unsophisticated way. A semaphore pattern is implemented in WCT which is limited in scope to a single process. It works by allowing only a limited number of GPU tasks to be submitted at any given time. Once the limit is reached, a subsequent task must wait for "its turn". Currently, this waiting is implemented by allowing the thread and thus the CPU core that issued the task to go idle. This idleness will tend to strike and remove cores from making progress precisely when CPU tasks are most needed. To break this bottleneck we plan to implement a so far ignored type of node provided by the TBB flow_graph library called async_node. Such a node is allowed a free thread under the assumption it will be lightly used. This allows an asynchronous query protocol with the graph execution engine. Essentially, the node may be periodically checked for completion. This can cause delay in the GPU task return but does not lead to CPU nor GPU idleness.
Another imperfection with the current semaphore based GPU load balancing is that the semaphore is implemented at the C++ language level. As such it produces a tight code binding between all flow-graph node implementations that wish to issue GPU tasks. This goes against the general wellfactored, interface-based code architecture of the toolkit. As the semaphore is process-local it is useless for moderation of GPU access which is shared across processes.
To solve this we expect to develop a distributed application composed of multiple processes based on art or WCT or in some cases both. The flow-graphs will have nodes that provide communication between processes which potentially may be executing on different compute nodes.
Initial plans for this distributed architecture were based on the ZIO [9] project. More recently, investigation into ADIOS2 [10] show it as a promising and likely preferred alternative. Regardless of the implementation, we require methods for high-throughput data flow and for low-latency round-trip query-response task servicing patterns. The exact throughput and latency requirements still must to be quantitatively estimated but they are not expected to particularly stress current technology.
Data flow may be used to connect WCT graphs in different processes an to enact a distributed asynchronous storage I/O path. Task servicing may be used to implement shared GPU load balancing. One possible configuration is illustrated in Figure 3.

Summary
DUNE software faces new challenges in order to run efficiently at HTC and HPC facilities. To respond, new parallel processing patterns are being and must continue to be employed at many levels of scale: loop level, task level, data flow level, between processes and between host computers. Jobs of a certain scale must become a distributed collection of heterogeneous parts. Distribution paths include file input/output, inter-graph data flow and GPU task servicing. Though introducing MT execution solves problems, new ones arise (eg N − 1, competing thread pools). The pattern of these problems are seen both in use of CPU-only resources and with mixed CPU-GPU. Flexible ways to configure jobs to match the limitations of any given hardware allocation are needed.
From its start, the Wire-Cell Toolkit was designed with some of these problems in mind to solve. It allows for highly multi-threaded operation, works in a synergistic manner with any TBB based application such as those based on the art event processing framework. In particular, the art MT support which has yet to be widely used provides potentially well matched solutions when coupled with WCT.
Much work is still required to realize the remaining solutions described here. When achieved, we expect to have software for DUNE and other adjacent experiment that is flexible enough to efficiently exploit both HTC and HPC resources.


extfieldsofstudy:
Computer Science
Engineering
Mathematics

text:
A Perspective on Neural Capacity Estimation: Viability and Reliability

Recently, several methods have been proposed for estimating the mutual information from sample data using deep neural networks. These estimators ar referred to as neural mutual information estimation (NMIE)s. NMIEs differ from other approaches as they are data-driven estimators. As such, they have the potential to perform well on a large class of capacity problems. In order to test the performance across various NMIEs, it is desirable to establish a benchmark encompassing the different challenges of capacity estimation. This is the objective of this paper. In particular, we consider three scenarios for benchmarking:i the classic AWGN channel, ii channels continuous inputs optical intensity and peak-power constrained AWGN channel iii channels with a discrete output, i.e., Poisson channel. We also consider the extension to the multi-terminal case with iv the AWGN and optical MAC models. We argue that benchmarking a certain NMIE across these four scenarios provides a substantive test of performance. In this paper we study the performance of mutual information neural estimator (MINE), smoothed mutual information lower-bound estimator (SMILE), and directed information neural estimator (DINE). and provide insights on the performance of other methods as well. To summarize our benchmarking results, MINE provides the most reliable performance.

I. INTRODUCTION
Determining the capacity of a channel and the optimal input distribution are problems of fundamental importance in communication scenarios of practical relevance. The value of capacity and the optimal input distribution provide critical insight on the choice of coding rates and input constellation shape, respectively. The traditionally, the literature on capacity computation has straddled two tendencies: that of (i) considering channel models that have an elegant formulation for which capacity is evaluated in closed-form through mathematical manipulations, versus that of (ii) considering channel models obtained from a practical communication setting for which capacity is determined through numerical optimization techniques. Neural mutual information estimators (NMIE) represent a novel approach to the analysis and design of communication systems that does not incur in the above limitations, since (i) they do not require the human mathematical modelling of the channel behaviours, as well as (ii) can efficiently tackle the inherent complexity of the capacity estimation problem. Instead, NMIES only require the ability to simulate the channel input/output relationship and use coupled DNN to estimate the capacity of the model: (i) one DNN that generates samples from the estimated optimal input distribution that, and (ii) a DNN that estimate the mutual information (MI) between the channel input and the channel output. In this sense, NMIEs can be consider as "universal" estimators in the sense that they only require data samples from the input/output relationship.
Since NMIEs are numerical methods, any question regarding their performance can only be investigated by benchmarking. A natural question that then arises is then What constitutes a substantive and yet minimal set of channels to test the performance of an NMIE?
In this paper we make a first attempt in defining such a benchmark for NMIE. More specifically, We consider the following models: AWGN channel, optical intensity, peak-power constrained AWGN, and the Poisson channel. For each channel, we extensively test the performance of three NMIEs over different SNR ranges: MINE, [1], SMILE, [2], and DINE. [3]. We give quantitative as well as qualitative assessment of their performance in the various regimes, as well as provide insight on the training process. In particular, we test the ability of these NMIEs to estimate optimal input distributions with large support, with discrete mass points, and with mass points that grow large as the SNR increases.
We argue that this rigorous benchmarking is necessary not only to discover the relative advantages of the various approaches, but also for the development of new NMIEs which are able to fill the performance gaps of existing solutions.

A. Relevant Literature
The computation of capacity using numerical methods most dates back to the well-known works of Blahut and Arimoto [4], [5]. This algorithm consists of an iterative alternating maximization method using the primal and dual formulation of the capacity problem [6,Sec. 4]. Various authors have considered extension and refinements to this original algorithm these include [7] [8] [9].
Similarly, computing closed-form expressions for the capacity region of general multi-user channels, except in few cases, has proved to be challenging. The capacity region of a two and three user memoryless multiple access channel (MAC) was formulated in [10]. Later, [11] considered the capacity region of MAC with correlated sources and derived its achievable rates region. In [12], it was shown For this class of capacity estimators, we benchmark performance over (i) the classic AWGN channel, (ii) channels continuous inputs -optical intensity and peak-power constrained AWGN channel (iii) channels with a discrete output-i.e., Poisson channel. We also consider the extension to the multi terminal case with (iv) the AWGN and optical MAC models.
The reminder of the paper is organized as follows.
• Sec. II: We start the paper by introducing the different NMIE proposed in the literature. Generally speaking, there are two general classes of estimators. A first class of estimator performs a direct estimation of the mutual information: among such estimators we study MINE [1] and SMILE [2].
Another class of estimators estimates the entropy based on reference random variables [28]: among such estimators we consider DINE [3] and InfoNCE [29].
Note that these entropy-based methods can be used for estimation of other information theoretic measures besides mutual information such as directed information and conditional mutual information.
• Sec. III: We discuss how the NMIE presented in Sec. II are employed for channel capacity estimation. This can be accomplished through an alternating optimization: first (i) a neural distribution transform (NDT) network trained to produce an estimate of the capacity-approaching distribution; successively (ii) the NMIE is trained to find a tight bound on the mutual information between channel input and output. By alternating these two optimization steps, convergence to a tight estimation of both the capacity and the capacity achieving distribution can be obtained. In this section, we shall also discuss the role of the initialization of both the NDT and the NMIE is discussed.
• Sec. IV: In this section, we introduce the channel models that we select for benchamrking, that is: the 1) AWGN channel with an average power constraint, 2) the optical intensity channel, 3) the AWGN with a peak power constraint, and 4) the Poisson channel. We also consider the extension of the proposed approach to the multiple access channel (MAC) by considering 1) the AWGN MAC and the 2) the optical intensity MAC. We present different capacity results in each case.
• Sec. V: We begin our benchmarking study by testing the performance of the various NMIE in Sec.
II for the case of the classic AWGN channel. Since the capacity and the optimal input distribution for this classic channel are well-known, we can perform a detailed comparison of different methods.
Various insights on the precision, numerical stability, and computational complexity of the different algorithms are provided. Overall, MINE emerges as a reliably and efficient NMIE estimator for capacity estimation.
• Sec. VI: The insights gained on the performance of the various NMIE estimators on AWGN channel are extended to estimating the capacity of the other channels presented in Sec. IV. Based on our numerical evaluations we observe that for the point-to-point channels, NMIE based on direct estimation of the mutual information give more accurate estimates of the capacity and estimated the optimal input distribution more closely. We also show that neural capacity estimation can give tighter bounds for channels with only upper and lower bounds on capacity such as the optical intensity channel. In the latter part of the section, in Sec. VI-D, we consider the problem of estimating the capacity region of the Gaussian and optical MAC channel.
Before proceeding further, let us introduce the notation employed in the paper. The L-norm of the vector x ∈ R d is indicated as x L . Random variables (RVs) are indicated with capital roman letters, i.e. X, their support with the same calligraphic letter, i.e. X , and their pdf as P X (x) or Q X (x), and when the context is clear as P (x) or Q(x). The normal distribution is indicated as N , the tail probability function for the standard normal distribution in x is indicated Q(x). In the remainder of the paper, we consider four measures with respect to RVs: the Kullback-Leibler (KL) divergence, entropy, mutual information (MI), and the χ 2 divergence, respectively defined as I(X, Y ) = D(P (x, y)||P (x)P (y)) (3) respectively, where we have agreed that 0 log 0/0 = 0. The binary entropy is defined as When this set of random variables is used to construct a random vector, we employ the notation X n m = [X m , X m+1 , . . . , X n ] with n ≥ m. Again, when m = 1, the subscript is omitted, i.e. X n 1 = X n . Given the scalar P , the value in dB is indicated as P db , i.e. P db = 10 log 10 (P ).
The code used in the remainder of the paper can be found online here: https://github.com/Farhad-Mrkm/Benchmarking-Neural-Capacity-Estimators

II. NEURAL ESTIMATION OF MUTUAL INFORMATION
In recent years, there have been many different methods proposed for estimating mutual information using neural networks [1], [2], [20], [28], [30]. We refer to these approaches as NMIE. We review these methods in this section and then in the next section, we discuss how they can be employed to estimate capacity and the optimal input distribution.
The NMIE methods can be divided into two groups: NMIE that evaluate the mutual information in (3) directly by using bounds on KL divergence [1], [2], [30], and (ii) NMIE that break the mutual information into entropy terms and estimate each entropy term separately using bounds on KL divergence [3], [20], [28]. We start this section by presenting the bounds on KL divergence, which are used by these NMIEs.

A. Bounds on Information-theoretic Measures
In this section, we provide lower and upper bounds on KL-divergence. Then in the next subsections, we use these bounds to describe the various methods that have been proposed in the literature for estimating mutual information using neural networks.
The first bound considered is the Donsker-Varadhan [31] bound rewritten in the theorem below for convenience.
Theorem 1 (Donsker-Varadhan representation [31]): Consider the RV X and two probability measures over X , P and Q, then the KL-divergence can be rewritten as where T is the set of all functions with finite expectations in (5).
Another lower bound for KL-divergence is obtained in [20], [30], and is presented here for convenience.
Theorem 2 (Tractable Unnormalized Barber and Agakov (TUBA) lower bound [20], [30]): Consider the RV X and two probability measures over X , P and Q, then the KL-divergence is lower-bounded where α is a positive constant.
An upper bound on KL-divergence using χ 2 -divergence can be obtained using [32,Theorem 23], which is rewritten in the theorem below for convenience.
Theorem 3 (KL divergence upper bound [32]): Consider a RV X and two probability measures over X , P and Q. The KL divergence is then upper bounded as D(P ||Q) ≤ log(1 + χ 2 (P ||Q)) (7) − 1.5 (χ 2 (P ||Q)) 2 log e (1 + χ 2 (Q||P ))(1 + χ 2 (P ||Q)) 2 − 1 In order to make this bound amenable to computation using data driven methods we use a variational form for χ 2 from [33]: As a result we have following variational form for the upper bound in (8) Where we have used monocity of log. In next section we use this bound to find our proposed explicit variational lower bound on mutual information, and we prove its consistency. The next section, Sec.
II-B, is devoted to present different neural mutual estimation approaches that utilize the results in Thm.

B. Direct Mutual Information Estimation
In this section, we review the NMIEs that estimate the mutual information in (3) through bounds on the KL divergence term on the right hand side of (3).

1) MINE:
One of the earliest attempt to NMIE estimation can be found in [1], where the MINE is proposed. The MINE is obtained from the the Donsker-Varadhan (DV) representation in Th. 1 when a DNN is used to represent T (X, Y ) in (5). Note that here the KL divergence is between the joint distribution and the product of the marginals as in (3), and hence the function T (.) is a function of X and Y . It is noted in [1] that the optimization of T (X, Y ) using a DNN through stochastic gradient descent (SGD) is generally challenging. This is because using a naive gradient estimate over the samples of the mini-batch leads to a biased estimate of the full gradient. Exponential moving average is proposed for mitigating this problem. However, this leads to an MI estimator with large variance, as noted in [30]. In [1], the MINE is shown to be effective in preventing the mode collapse in generative adversarial networks (GANs).

2) SMILE:
To address the large variance of the MINE, the authors of [2] propose the SMILE. In this method, the DV representation in Theorem 1 is rewritten as in the following theorem.
Theorem 4 (SMILE [2]): Consider the RV X and two probability measures over X , P and Q. The KL-divergence can then be approximated as where clip(x, τ ) min{max{x, e −τ } , e τ }.
SMILE leverages the bound in (11) as the clipping function clip is equivalent to clipping the log density ratio estimator in the interval [−τ, +τ ]. Similar to MINE mutual information can be estimated by using a neural network to represent T (.), and optimizing the neural network using gradient descent.
The choice of τ crucially affects the bias-variance trade-off. With smaller τ , the variance of the estimator is reduced at the cost of potential increase to its bias.
3) Reverse Jensen Inequality: Another approach for addressing the high variance in estimation of the mutual information caused by the second expectation in the DV representation uses a reverse Jensen inequality [34]. In this case, the mutual information is lower-bounded as where ζ b (a) is defined in [34,Lemma 1], and c is a scalar parameter. One of the challenges of this approach is that the parameters on the right hand side must be tuned. Another is that the proposed reverse Jensen inequality in [34] may not be tight for all settings. [29]: This bound provides a method for computing MI directly using K samples of joint distribution of X, Y . This bound is given by:

4) InfoNCE Contrastive Lower Bound
where p K (x, y) is the distribution of K independent samples from the joint distribution: j p(x j , y j ). [35]: This bound has been extensively used in the literature of f -divergence estimation [35], and is used as loss function for generative adversarial networks [36]. Its representation is based on the Legendre-Fenchel duality and is expressed as:

C. Entropy-based Mutual Information Estimation
Another technique for estimating MI, rewrites MI in terms of a number of entropy terms and uses the variational bounds on KL divergence to estimate the entropy terms from data. Specifically, using (2), and known arbitrary i.i.d. reference random variables X and Y , the mutual information can be expressed as Note that since we can choose the reference distributions to be i.i.d., the cross-entropy terms in (2) will cancel out leaving only the KL-divergence terms. We now present the methods that use this approach for estimating mutual information.

1) DINE:
The DINE method estimates the directed information, which is used to compute the capacity of channels with feedback [3]. It can also be used to estimate the MI and the capacity of memoryless channels. Instead of using (13)

to decompose mutual information, DINE uses h(Y )
and h(Y |X) and hence has two KL-divergence terms to estimate instead of the three in (13). Each KL-divergence term is then computed using (5) adopted from [28]. Since the KL divergence term corresponding to h(Y |X) is positive while the KL divergence term corresponding to h(Y ) is negative, and (5) is a maximization, the estimated MI is neither a lower nor an upper bound. To support feedback and channels with memory, DINE uses a recurrent neural network to estimate the directed information and hence it is slower to train for memoryless channels without feedback. This can be easily resolved however by using fully connected layers to represent the function T as is done in [20].

D. Chi-Square Bound
In this subsection, We use the upper bound on the chi-square (χ 2 ) divergence to find an explicit lower bound on MI. A variant of this method was employed in [20] to compute the boundaries of the capacity region of the multiple access channels. This method uses the decomposition in (13). For the positive KL divergence terms in (13), (6) is used where T is represented and optimized using a neural network. For the negative KL divergence terms in (13), the upper bound χ 2 UP (P ||Q) from Theorem 5 is used, where these terms are replaced by a variational form of χ 2 divergence. This is replaced the histogram based plug-in estimator which initially was used in [20], and allows for flow of gradient in the upper bound as well. A detailed discussion about consistency of this estimator is provided in the appendix. The following theorem summarizes this approach: The MI between the RVs X and Y is lower bounded as and for χ 2 UP in (10). We refer to this method as χ 2 bound in the paper.

E. Extension to Conditional Mutual Information
In this subsection we extend the NMIE methods to the estimation of the conditional mutual information (CMI), which is needed for estimating the capacity region of the MAC. Just like the MI estimation, the CMI can be estimated using several different approaches.

1) Direct Conditional Mutual Information Estimation:
The CMI can be evaluated using the DV bound in Theorem 1 directly using The main challenge in using the DV bound for estimating CMI from data samples is creating training samples from Q ≡ P X,Z P Y |Z . In estimating MI, the samples for the product of the marginals are easily obtained using random permutations on the samples of x and y drawn i.i.d. from the joint distribution, i.e., (x, y) ∼ P X,Y . Several methods have been considered for drawing samples from P X,Z P Y |Z , given a data from the joint distribution using conditional generative adversarial network (CGAN), conditional variational autoencoder (CVAE), and nearest neighbors sampling [37], [38].

2) Mutual Information Based Conditional Mutual Information Estimation:
Another approach for estimating CMI relies on the the decomposition of the CMI into difference of MI terms as: Each MI term in (16) could be estimated using one of the NMIE methods such as MINE and SMILE method. Note that since the consistency of these estimators has been proved in [1], [2] for MI estimation, the CMI estimation using this approach will also be consistent. 3) Entropy Based Conditional Mutual Information Estimation: Finally, CMI can be decomposed into a summation of entropy terms, which can then be evaluated by using arbitrary reference random variables X , Y , and Z . Specifically, CMI can be decomposed as: where the last equality follows from the independence of X , Y , and Z , which cancels out the crossentropy terms from (2). The four KL-divergence terms can then be estimated using any of the prior methods based on the DV bound such as MINE or SMILE.

III. NEURAL CAPACITY ESTIMATION
In the previous section, we presented different NMIEs techniques. These methods can be used to estimate the mutual information by jointly sampling the channel input and the channel output, without knowing the underlying channel models. In this section, we use these NMIEs as part of an optimization process for estimating the channel capacity from data samples.
The block diagram for neural capacity estimation approach is shown in Fig. 1. In this approach, which was also used in [3], the channel input X is generated from a distribution that is learned by a neural network called the NDT, parameterized by θ. Specifically, the NDT take random samples S, generated from a known distribution P S , as input and outputs X (i.e., samples of the channel input). Given a sufficiently large neural network, which can act as a general function approximator, the distribution of P S can be transformed into an arbitrary channel input distribution P X , through parameters θ. Note that the activation function at the final layer of the NDT should be chosen to match the support of X. Then, given a channel and the channel input X, the channel output Y is generated.
Samples from the channel input and output are then used to estimate the mutual information using NMIEs, parameterized by φ. Therefore, the capacity is estimated using a double maximization where S are samples drawn from a known distribution, f θ is the NDT, and I φ is the NMIE.
The NMIE network and NDT are trained iteratively to estimate the mutual information for a given P X and updating P X to maximize mutual information. Hence, a single training iteration has two phases.
In phase one, the NMIE is trained exclusively to estimate the mutual information. This is achieved by keeping the weights of the NDT network constant and only updating the weights of the NMIE during backpropagation. In the second phase, the weights of the NMIE is kept constant and the weights of the NDT network are updated using backpropagation. The estimate of the capacity is obtained when this iterative optimization converges or certain number of training steps are reached.
Note that this method is very similar in sprite to the generative adversarial networks [39], where a generator network transforms random samples from a known distribution to random image samples, and a discriminator network classifies the generated images as real or fake. An iterative optimization process is then employed to improve the distribution over which random images are generated from, and the classification accuracy of the discriminator. At convergence, the output of the discriminator is random as it is unable to identify an image as real or fake.
Algorithm 1 summarizes this procedure. Several factors affect the performance of the neural capacity estimator, namely the choice of NMIE technique from Section II, and the random distribution used for generating the samples s i . When the optimal input distribution is known to be continuous, the normal Gaussian distribution can be a good choice for P S . When the channel has discrete outputs (like Poisson channel), and optimal input is known to have discrete mass points, we use Algorithm 2. The reason behind this change in the algorithm for computing capacity is that we have found experimentally that using continuous normal Gaussian distribution for channels with discrete output and discrete optimal distributions have a degraded performance with Algorithm 1 compared to using discrete distributions as initialization distributions.
This algorithm starts by evaluating the capacity using Algorithm 1 with a uniform binary distribution for P S . Then a new mass point is added to the distribution of S and the capacity is re-evaluated using Algorithm 1. This process of adding new mass points to the distribution of S and then re-evaluating capacity is continued until the estimated capacity no longer improves due to the newly added mass point. As we show in our numerical results, this proposed approach can achieve better estimate of the using NDT and the channel Calculate the MI loss in (18) Train the NMIE I φ using Adam optimizer using NDT and the channel Calculate the MI loss in (18) with respect to θ Train NDT network f θ using Adam optimizer Return:Ĉ = I φ (X; Y ) andp θ using large samples of S drawn at random. Despite this advantage, neural capacity estimation can be unstable based on the choice of distribution for S, and the choice of NMIE technique. Therefore, in this work, we evaluate the performance of neural capacity estimation using different NMIE methods, and different distributions for S for variety of communication channels. To the best of our knowledge this is the most comprehensive empirical study of neural capacity estimators. In the next section, we describe the channel models we will consider as part of our empirical study, followed by our numerical results.
1) Implementation details: Let us briefly introduce here the NDT and NMIE architecture, training, and simulations. All DNN consists of multi-layer fully-connected neural networks, with up to four layers. The batch size for training is 256, and we use unit variance Gaussian distribution as initialization, unless otherwise stated. For obtaining a histogram of the input distribution, we use 64000 samples.
The ADAM optimizer is used in all cases with a learning rate=0.0001.

IV. CHANNEL MODELS
In the remainder of the paper, we investigate the performance of the neural capacity estimators and the quality of the learned input distribution. For this purpose, we compare the performance of the proposed neural estimator with the capacity results derived in the literature. Before we present our numerical results, in this section, we present the channel models studied in the reminder of the paper and present relevant bounds on capacity, and characterizations of the corresponding capacity-achieving input distribution.

A. Point-to-Point (P2P) Memoryless Channels
The point-to-point (P2P) channel describes the communication between one transmitter and one receiver as described by the conditional probability P Y |X . For the P2P class of channels, it has been show [40], [41] that the capacity is obtain by the supremum of the mutual information between input and output over all feasible input distributions, i.e.
Let us briefly introduce the models we shall investigate in the remainder of the paper.
where N is the block-length, Z N is a vector of i.i.d standard normal distribution, and the input is subject to the average power constraint The capacity of the AWGN channel is well-known and given by: Note that the LHS of (22) explicitly captures the dependency on the second moment from the constraint in (21). Additionally, it is well-known that the capacity achieving input distribution is the Gaussian distribution with zero mean and variance P .
2) Optical Intensity (OI) Channel: For the OI channel, the input/output relationship is again given by (20), but the input is subject to the constraints for all n ∈ [N ]. As in [42], we denote the ratio between the average power and the peak power as Note that, while the input is constrained to be positive, there is no restriction on output of channel, which could be either positive or negative.
The capacity of the OI channel is indicated as C OI (A, E), where again we have explicitly indicated the dependency of the constraints in (23). An analytical expression for the capacity of the OI channel C OI (A, E) is not known, however, prior works have found upper and lower bounds on capacity [42]- [45]. We present the best known bounds on the capacity of the OI channel, used in this work for the performance evaluation in the theorem below for convenience.
Theorem 6 (Bounds on the capacity of the OI channel [42]): The capacity of the OI channel for and upper-bounded as one equations in [42, eq.19] and [42, eq.20] (the smallest determines the upper bound).

3) Peak Power-Constrained AWGN (PPC-AWGN) Channel:
The PPC-AWGN is again obtained from the Gaussian additive P2P channel in (20) but under the additional peak-power constraint A closed-form expression for the capacity for the PPC-AWGN channel, C(P, A), and the corresponding optimal input distribution are not know in general. Therefore, in this work, we use the lower and upper bounds on the capacity derived in [46] for the performance evaluation of the neural capacity estimators. The specific bounds used for evaluation are presented in the theorems below for convenience.
Theorem 7 (Bounds on the capacity of PPC-AWGN [46]): The capacity of the PPC-AWGN channel and lower-bounded as where f (y) is the mixture density resulting from the sum of the discrete M -point uniform input distribution with independent unit-variance Gaussian noise.
A refinement of the upper bound from [46] is derived in [47] for some specific intervals of P .

4) Poisson Channel:
The Poisson channel has a formulation that is similar to the OI channel in which the channel input X n is continuous over the interval [0, A]. Unlike the OI channel, the channel output Y n , however, is discrete and positive. More specifically, the Poisson channel is described by the conditional probability for y n ∈ N and where λ 0 is noise due to dark current. The channel input has a peak constraint A, and an average power constraint We denote the capacity of this channel as C PC (E, A). An analytical expression for C PC (E, A) is not known in general. In [48], it was shown that the capacity-achieving input distribution has finite number of mass points and the number of mass points goes to infinity as the constraints are relaxed. Upper and lower bounds on the capacity for different input constraints were derived in [49]. Finally, numerical methods such as deterministic annealing can be used to numerically evaluate capacity and the optimal input [7].
In the numerical results, we use the lower bounds and upper bounds from [49] to evaluate the performance of neural capacity estimators on the Poisson channel. The bounds are restated as theorems below for convenience.

Theorem 9 (Th. 3-[49]):
When Another theorem in [49] bounds the capacity in the 1/3 ≤ E/A ≤ 1 regime. It should be noted that upper bound and lower bound above are asymptotically tight but are not tight over the whole SNR range. The authors of [50] propose the Deterministic Annealing (DA) algorithm for the evaluation of the capacity of the Poisson channel. We will also compare the neural capacity estimators with this approach.

B. Multiple Access Channel (MAC)
The 2-user MAC channel is the multi-terminal channel model in which two users simultaneously and synchronously transmit toward a common receiver. The relationship between the input at the two transmitters and the output at the receiver is described by the transition probability P Y |X 1 ,X 2 . For this model, the ultimate communication performance under which reliable communication is possible is described by the capacity region, that is a tuple of transmission rates (R 1 , R 2 ) for which error probability can be made vanishing small for large block lengths.
In [51], [52] it was shown that the capacity of the MAC is obtained as the convex hull of rate points satisfying for some P X 1 ,X 2 = P X 1 P X 2 .

1) AWGN MAC:
Among the most studied MAC models, AWGN MAC is perhaps the most wellknown model. In this channel, the input/output relationship can be expressed as where X N 1 and X N 2 are the channel inputs from users 1 and 2, respectively, and Z N is an i.i.d. Gaussian sequence with zero mean and unitary variance. The input is further subject to a power constraint analogous to that in (21) as for i ∈ {1, 2}.
The convex hull of the tuple in (33) can be shown to be equivalent to the following region For the model in (34), it is further known that the capacity achieving input distribution is Gaussian, that is, any tuple on the boundary of (36) can be achieved through Gaussian distributed inputs at both transmitters.
2) Optical Intensity (OI) MAC: Similarly to the relationship between the AWGN and the OI channel, in Sec. IV-A1 and Sec. IV-A2 respectively, the OI MAC [53] is described by the input/output relationship in (34) but the channel inputs are subject to the constraint When only the constraint in (37b) is considered, the capacity region becomes [14], [15] The channels described in this section results will be used in the next sections to evaluate the performance of the neural capacity estimators as well as the quality of the learned input distribution.

V. AWGN CHANNEL: BENCH-MARKING NMIES
In this section, we benchmark the MI estimators discussed in Sec. II by comparing the estimated value of capacity and the learned input distribution to the optimal ones for the AWGN channel described in Sec. IV-A1. Since a closed-form expression for both the capacity and optimal input distribution is available for the AWGN channel model, this simple example provides a number of insights on the choice of NMIE.
Results for this channel are obtained using 10 separate estimation trials, where for each trial the NDT and the NMIE networks are initialized randomly. In Table I, we provide a comparison between the true capacity and its neural estimate for the 2dB, 20dB, and 40dB SNRs. Specifically, we provide the average estimate across the 10 trials, and the same number of trials for variance of the capacity estimation. We note that the relative accuracy decreases as the SNR grows but the loss of accuracy does not affect all methods equally. The MINE method gives a more accurate estimate of the capacity compared to other methods, while DINE seems to diverge from the true capacity at higher SNRs.
Moreover, the standard deviation of MINE is lower than all other methods suggesting that it is more robust to random initialization and has a lower variance in estimating capacity. Note that this is contrary to estimation of mutual information, where methods such as SMILE show lower variance. Fig. 2 shows the histogram of the NDT learned input distribution for the DINE, MINE, and SMILE.
We observe that at higher SNRs the MINE achieves a more accurate approximation of the optimal input distribution, i.e. Gaussian. While DINE can also learn accurate input distributions at low SNRs, at higher SNRs, it seems to deviate from the optimal input. Note that all three methods are trained for the same number of iterations here for a fair comparison.
During testing, we note that the clipping operation used in SMILE, significantly effects the quality of the learned input distribution. After experimenting with the range of values in the range τ ∈ (10, 0) for clip in (11), we have found clipping value τ = 0.2 gives the best estimates of capacity in terms of stability of results.
Remark 1: During testing, we note that the clipping operation used in SMILE, significantly effects the quality of the learned input distribution. After experimenting with the range of values in the range τ ∈ (10, 0) for clip in (11)

VI. NMIE FOR CAPACITY INVESTIGATION
In Sec. V we validated the neural capacity estimation described in Sec. II by considering a wellunderstood channel model. In this section, we benchmark the NMIE for channel models in which no tight capacity bounds exists and for which the optimal input distribution is not well understood.
From a high level perspective, an overview of the results in this section is as follows. Although all methods in Sec. II can be used for obtaining capacity bounds, the MINE method has the most reliable performance. For this reason, most of the results in the section will rely on this method and focus on different aspects of this approach, such as initialization, stability, variance of the estimates. To this end, in the first two subsection, we study the OI channel and PPC-AWGN channel and investigate whether and how accurately NMIE is capable of identifying the discrete masses in the optimal input distribution.
In Sec. VI-C, we study the Poisson channel and propose a discrete initialization procedure that promotes sparsity in the optimal input distribution estimates.

A. The Optical Intensity (OI) Channel
In this section, let us further explore the merits of the various approaches in Sec. II-B by considering the OI channel in Sec. IV-A2. The capacity of optical intensity is generally unknown [14], [15]. Since the exact capacity with closed-form solution is not available for this channel we compare our results with the bounds in [14], which are some of the best known bounds in the literature which we have presented in Th. 6. Table. III shows the estimated value of the capacity and achievable rates using   Table. III. The upper and lower bounds in Table. III are evaluated as in Th. 6.
A significant observation for OI channel is that, unlike other methods, MINE provide accurate estimates of optimal input distributions for OI channel. It is well known that optimal input distribution for OI channel is discrete [15], [42] which agrees by our finding using NDT and NMIE.
As we discussed before since χ 2 , SMILE and DINE method only provides estimates of achievable rates, for the approximation of optimal input distribution for the channel, we focus on MINE method.
In Fig. 4 we present the estimated optimal input distribution through MINE for the case = A, that is only peak constraint. For this simulations, the algorithm was initialized using standard Gaussian random variable as initialization input distribution.   [47], simulations consider the case A = √ P dB. In this subsection, we apply the proposed approach for computing capacity and optimal input distribution using MINE method to the channel PPC-AWGN, introduced in Sec. IV-A3. When comparing the numerical results with the bounds in Th. 7 and Th. 8, we observe that our approach either match or fill the gap between upper and lower bounds in each case. This can be verified from in Tab. IV, which reports the capacity approximation evaluated through the various NDT networks, together with upper and lower bounds in Th. 7 and Th. 8.
Based on the numerical evaluations, we conclude that using NMIE methods that estimate the MI directly can result in more accurate estimates of the channel capacity and the optimal input distribution.
While more investigation is needed to gain a deeper understanding of this observation, we suspect that this is because the error terms in entropy-based estimations compound through the learning process, resulting in less accurate estimates of MI and, consequently, capacity. Moreover, entropybased estimators tend to be less robust to the random initialization of the networks.

C. The Poisson Channel
Next, let us investigate the capacity of the Poisson channel in Sec. IV-A4 to further understand the role of discrete initialization of proposed deep estimation approach as detailed in Algorithm 2.
As we mentioned earlier, the motivation for the use of discrete initialization is that for the Poisson channel, where the output of the channel and the optimal distribution are discrete, discrete initialization distributions achieve higher information rates compared to using normal distribution as initialization distributions which are closer to the capacity. This has been illustrated in the simulation summarized in the Table V for the Poisson channel.
For this model it is known [48] that the capacity-achieving distributing is discrete and with a number of points that increases with the SNR. In order to identify the cardinality of the optimal input distribution we proceed as in Algorithm 2, that is by choosing a discrete initialization of this distribution with uniformly spaced mass points between zero and the peak power constraint (if considered). An increasing number of mass points is considered in the initialization until the MI is no longer increasing. The various stages of Algorithm 2 are represented in Fig. 6.
Having established the capacity optimization approach, let us validate the usefulness of such approach by considering the general Poisson channel in which both average and peak power constraints are present.
Since for the Poisson channel, a closed-form expression of the capacity and the optimal input distribution is not available, we consider those regimes in which tight lower and upper bounds on capacity are known in the literature. More specifically it is known that [48], in the case of zero dark current (that is λ = 0) and only peak constraint, the optimal input, is binary in the regime 0 ≤ A ≤ 3.3679 and has larger cardinality when A ≥ 3.3679. Furthermore in [50] it has been proved that, for the Poisson channel with both peak and average constraint, the optimal input always contain a mass point at zero. In addition [50] show that, under only peak power constraint, the optimal input distribution contains a mass point at the peak value A. In our simulations, we consider the case of a Poisson channel under average power constraint only. In this regime, [48] and [50] conjecture that the optimal input distribution is exponential.
Is Fig. 7   Note that the transition between two to three mass points in optimal input distribution occurs for A = 3.36 [48].   mass points. This is correctly recovered by the proposed approach. For this scenario, the capacity estimated in [7] is 0.91nats, while our approach yields an estimated capacity equal to 0.82nats.
In each scenario, we compare the performance of MINE with the upper and lower bound in Th. 9, originally from [49].
In Table V we provide a similar comparison of the estimated capacity for the Poisson channel with A = 100 and E ∈ {5, 10, 15, 20}. Here we compare the MINE, the lower bound in Th. 9 as well as the deterministic annealing (DA) approach from [7]. Note that, in Table V, we also provide the performance of the NDT in estimating capacity for both the Gaussian initialization as well as the discrete initialization devised in Algorithm 2. It can be observed that, for this model, the performance of the discrete initialization far outperforms that of the Gaussian one.
Finally, in Fig. 8, we plot the Poisson channel with only average power constraint. In [48], the authors conjecture that the optimal input distribution has unbounded support and takes shape of exponential distribution. We notice that Fig. 8

D. Multiple Access Channel (MAC)
In this section we apply MINE method to compute the capacity of MAC introduced in Sec. IV-B.
To compute the sum capacity using MINE method we produce two sequence of i.i.d random variables and after concatenating them we feed them to NDT neural networks. Then, each column of the output matrix is used as two-input distribution of the channel which satisfy the channel constraints. After producing the output of channel using these inputs from NDT network, we use NMIE networks to estimate the mutual information and train the joint network of NDT and NMIE based on the Algorithm 1. The corner point of capacity region are computed by the conditional mutual information I(X 1 ; Y | X 2 ) and I(X 2 ; Y | X 1 ). To compute these points we simply decompose the mutual information as I(X 1 ; Y | X 2 ) = I(X 1 , X 2 ; Y ) − I(X 2 ; Y ) and use previous approach for each term to compute the conditional term. For both the AWGN MAC and the OI MAC, we observe that MINE provides relatively accurate estimates of sum-capacity. Furthermore, we observe that for the OI MAC, our approach provides estimates of achievable rates which fall between the tightest inner and outer bounds for this model.

VII. CONCLUSION
In this paper, we make a first attempt in establishing a benchmark for neural capacity estimation.
We consider the problem of capacity estimation on a number of point-to-point channel, as well as multiple access channel (MAC). More precisely, we consider point-to-point channels (i) with a known capacity -AWGN channel as well as (ii) channels with unknown capacity and continuous inputs -i.e., optical intensity and peak-power constrained AWGN channel (iii) channels with unknown capacity and discrete number of mass points -i.e., Poisson channel. We also extend the analysis to the AWGN MAC -case (i) above-and the optical intensity MAC -case (ii). For the models above, we study the performance of mutual information neural estimator (MINE), smoothed mutual information lowerbound estimator (SMILE), and directed information neural estimator (DINE). and provide insights on the performance of other methods as well. Our numerical evaluations show that all methods perform rather well at low SNR, while only some methods perform well at high SNR. In particular, MINE emerges as the most versatile method across the different channel model. From our numerous numerical experiments, we observe that the network initialization plays a crucial role attaining accurate capacity estimation. For this reason, we propose a specific initialization strategy for channel which the optimal input distribution has a discrete number of masses, such as the Poisson channel. Overall, we notice that, although very promising, accurate neural capacity estimation crucially depends an the ability of the underlying network to approximate functions with a large support: future research will focus on the development of such DNN.

VIII. APPENDIX
Consistency proof for χ 2 bound : We first discuss that expectation in the bound can be replaced by emperical averages over batches of data where the optimal function T * is T * = log dP dQ + k where k is a constant defined such that dP is relate to Gibbs density by dP = 1 k e T * dQ. For the second term we have: to see that above inequality is true, it sufficies to note that x i , y i and their transformation through exponential function are iid random variables, and invoke AEP theorem [55]. by choosing a large n, could be made arbitrary small. Now we turn to χ 2 UP (P ||Q) terms.We provide the argument for one term, and similar arguments hold for the last term. We have The estimation of first term considered in (39), the estimation of second term using averages over mini-batches is also evident from weak law of large numbers and noting that x i , x i are i.i.d.
Now we turn to our main consistency result. For the first term in (14) the consistency of estimators using neural networks as universal approximators have been established in [1] [56]. Here we provide consistency result for the variational form χ 2 and show that it can provide consistent estimates when feed-forward neural network are used for its representation. In fact we show the following: