# Machine Learning-powered Compact Modeling of Stochastic Electronic Devices
using Mixture Density Networks

Jack Hutchins  Department of Electrical Eng. & Computer Sci., University of
Tennessee, Knoxville, Tennessee 37996, USA  Shamiul Alam  Department of
Electrical Eng. & Computer Sci., University of Tennessee, Knoxville, Tennessee
37996, USA  Dana S. Rampini  National Institute of Standards and Technology,
Boulder, Colorado 80305, USA  Bakhrom G. Oripov  National Institute of
Standards and Technology, Boulder, Colorado 80305, USA  Adam N. McCaughan
National Institute of Standards and Technology, Boulder, Colorado 80305, USA
Ahmedullah Aziz  Department of Electrical Eng. & Computer Sci., University of
Tennessee, Knoxville, Tennessee 37996, USA  Corresponding Author. Email:
<EMAIL_ADDRESS>
###### Abstract

The relentless pursuit of miniaturization and performance enhancement in
electronic devices has led to a fundamental challenge in the field of circuit
design and simulation—how to accurately account for the inherent stochastic
nature of certain devices. While conventional deterministic models have served
as indispensable tools for circuit designers, they fall short when it comes to
capturing the subtle yet critical variability exhibited by many electronic
components. In this paper, we present an innovative approach that transcends
the limitations of traditional modeling techniques by harnessing the power of
machine learning, specifically Mixture Density Networks (MDNs), to faithfully
represent and simulate the stochastic behavior of electronic devices. We
demonstrate our approach to model heater cryotrons, where the model is able to
capture the stochastic switching dynamics observed in the experiment. Our
model shows 0.82% mean absolute error for switching probability. This paper
marks a significant step forward in the quest for accurate and versatile
compact models, poised to drive innovation in the realm of electronic
circuits.

## Introduction

In the ever-evolving landscape of electronic devices, the pursuit of smaller,
faster, and more efficient technology has led to groundbreaking innovations,
unlocking new possibilities for various applications across industries [1].
However, as devices continue to shrink and push the boundaries of Moore’s law
[2], the inherent stochastic nature of certain electronic components has
become increasingly significant [3]. These stochastic electronic devices
exhibit inherent variability in their behavior, posing a formidable challenge
for conventional deterministic modeling approaches [4].

In the pursuit of greater accuracy and ease of implementation in electronic
circuit simulations, a significant leap has been taken by incorporating
machine learning methodologies into compact modeling [5, 6, 7, 8, 9, 10].
While various machine learning-based compact models have emerged, these
models, thus far, have failed to address the implications of stochastic
behavior in electronic devices. And while traditional physics-based compact
models have been created with stochastic properties[11, 12], these models lack
the benefits of machine learning-based methods such as decreased turnaround
time and little required device knowledge. The dynamics of electronic devices,
traditionally considered deterministic, are increasingly proving to be
intrinsically stochastic at the nanoscale. Consequently, conventional compact
models that ignore these inherent stochastic processes fail to capture the
true behavior of the devices, leading to inaccuracies in circuit simulation.

This paper presents a new approach that addresses this gap in the realm of
compact modeling. We introduce the use of Mixture Density Networks (MDNs)[13]
to capture the stochastic nature of electronic devices accurately. In addition
to this neural network approach, we introduce a novel sampling methodology
that uses inverse transform sampling [14] to make our model more accurate to
the stochastic nature of devices. By doing these, we create compact models
that account for the inherent stochasticity of certain devices, allowing the
development of new circuits that are both reliable and innovative. Our
approach represents a paradigm shift in electronic device modeling, as it for
the first time encompasses the full spectrum of electronic device behavior,
from deterministic to stochastic.

To demonstrate the effectiveness and versatility of our approach, we focus on
the modeling of heater cryotron [15], a three-terminal device that shows gate-
current controlled switching between superconducting and resistive states.
Heater cryotrons, with their stochastic switching behavior, are a particularly
compelling case study. Our methodology accurately captures the nuances of this
stochasticity, providing a foundation for the development of novel circuits
and the optimization of existing designs.

The contributions of this paper are threefold:

  * •

Mixture Density Networks: By employing MDNs, we equip our modeling framework
with the capacity to capture complex probability distributions, thereby
enabling a more faithful representation of device variability. This empowers
designers to explore device behavior beyond traditional deterministic bounds.

  * •

Inverse Transform Sampling: By leveraging inverse transform sampling, we can
use our MDN to its full potential. This approach allows the model to output
smooth outputs in transient simulations while maintaining stochasticity. This
in turn makes our model more realistic to device-to-device or cycle-to-cycle
variation.

  * •

Demonstration with Heater Cryotrons: To showcase the effectiveness of our
approach, we focus on modeling heater cryotrons—a class of devices known for
their stochastic switching behavior. Through this demonstration, we illustrate
the practical applicability of our methodology in real-world electronic
systems.

As we delve into the details of our approach, we will elucidate the
intricacies of MDN-based modeling for stochastic electronic devices, providing
a comprehensive understanding of the benefits it offers to circuit simulation
and electronic design.

## Results

The following section will describe our neural network architecture, sampling
methodology, and the results of our approach using experimentally derived
heater cryotron data.

### Mixture Density Network

Mixture density networks provide the perfect approach for modeling stochastic
devices. Mixture density networks work similarly to standard multilayer
perceptions (MLP)[16], but with three output layers connected to the last
hidden layer instead of just one. Each output layer has the same number of
neurons $N$, and we label the output layers as $\mu$, $\sigma$, and $\alpha$.
The output of these layers is used as the parameters of a Gaussian mixture
distribution probability density function (PDF) as defined in equation (1).

$p(x)=\sum_{k=1}^{N}\alpha_{k}\cdot\frac{1}{\sqrt{2\pi\sigma_{k}^{2}}}\cdot\exp\left(-\frac{(x-\mu_{k})^{2}}{2\sigma_{k}^{2}}\right)$
(1)

Using this approach allows the model to learn the unique output distribution
for any given input. The PDF is a combination of $N$ different Gaussian
distributions, which are multiplied by a scaling factor $\alpha$. In order to
maintain a valid probability distribution, we need to ensure that the sum of
$\alpha$ is always equal to 1. To do that, we apply a softmax function to the
output of the $\alpha$ layer, as defined in equation (2).

$s(\alpha_{i})=\frac{e^{\alpha_{i}}}{\sum_{j=1}^{N}e^{\alpha_{j}}}$ (2)

Additionally, we need to ensure that the output of the $\sigma$ layer needs to
be positive and non-zero since they represent the standard deviation of the
Gaussian distributions in the mixture density PDF. In order to do this, we use
the exponential linear unit (ELU) activation function plus 1 plus $\epsilon$,
where $\epsilon$ is the smallest value before loss of precision in floating
point calculations. The activation function can be seen in equation (3).

$ELU(\sigma_{i})=\left\\{\begin{array}[]{lr}\sigma_{i}+1+\epsilon,&\text{if
}\sigma\geq 0\\\ 0.5(e^{\sigma_{i}}-1)+1+\epsilon,&\text{if
}\sigma<0\end{array}\right\\}$ (3)

Finally, we allow $\mu$ to be any value, so there is no need to use an
activation function to restrict its output. For the hidden layers, we chose to
use the rectified linear unit activation function, which is defined in
equation (4).

$ReLU(x)=\left\\{\begin{array}[]{lr}x,&\text{if }x\geq 0\\\ 0,&\text{if
}x<0\end{array}\right\\}$ (4)

In order to train the model, we need an appropriate loss function to measure
the performance of the model. We choose to use Gaussian Negative Log-
Likelihood (GNLL), since it effectively captures the performance with respect
to probability. The GNLL for a single data point $x$ is the negative logarithm
of this PDF:

$\text{NLL}(x)=-\log\left(\sum_{k=1}^{N}\alpha_{k}\cdot\frac{1}{\sqrt{2\pi\sigma_{k}^{2}}}\cdot\exp\left(-\frac{(x-\mu_{k})^{2}}{2\sigma_{k}^{2}}\right)\right)$
(5)

The goal during training is to minimize the GNLL across all training data
points. Therefore, the overall GNLL loss is the average GNLL over all data
points:

$L=-\frac{1}{N}\sum_{i=1}^{N}\log\left(\sum_{k=1}^{K}\alpha_{k}\cdot\frac{1}{\sqrt{2\pi\sigma_{k}^{2}}}\cdot\exp\left(-\frac{(x_{i}-\mu_{k})^{2}}{2\sigma_{k}^{2}}\right)\right)$
(6)

During training, we adjust the network’s parameters to minimize this GNLL
loss. We can use backpropagation and the AdamW[17] optimizer to update the
network’s weights and biases by finding the gradients of GNLL with respect to
$\mu$, $\sigma$, and $\alpha$ as shown in figure 1.

Figure 1:  Overview of the machine learning-powered modeling approach for the
stochastic behavior of heater cryotron. After training the mixture density
network with the device characteristics obtained from experiments, the weights
are extracted which are eventually used to create the Verilog-A-based compact
model for circuit/system-level simulations.

### Model Sampling Methodology

Algorithm 1 Standard Sampling from a Gaussian Mixture Density Distribution

Input:

$K$ \- Number of Gaussian components

$\\{\mu_{1},\mu_{2},\ldots,\mu_{K}\\}$ \- Means of the Gaussians

$\\{\sigma_{1},\sigma_{2},\ldots,\sigma_{K}\\}$ \- Stand deviation of the
Gaussians

$\\{\alpha_{1},\alpha_{1},\ldots,\alpha_{1}\\}$ \- Mixing coefficients,
$\sum_{i=1}^{K}\pi_{i}=1$

Output:

Sample $x$ from the mixture distribution

Sample Function:

sum = 0

Sample $q$ from a uniform distribution [0, 1]

for $k\in\\{1,\ K\\}$ do

sum = sum + $\alpha_{k}$

if sum $\geq\ q$ then

break

end

Sample $x$ from the Gaussian distribution $\mathcal{N}(\mu_{k},\sigma_{k})$

Return:

$x$ \- Sample from the Gaussian mixture density

In order to properly utilize our model, it is critical that we use an
appropriate sampling methodology for the output distributions. The obvious
approach for this would be to randomly sample from the distribution, however,
this approach presents several issues. In order to sample in this way, we can
use the approach outlined in Algorithm 1, where we randomly select a
distribution $k$ for the mixture density distribution and sample from a
standard normal distribution $\mathcal{N}(\mu_{k},\sigma_{k})$. This approach
allows us to utilize uniform and standard Gaussian sampling, which is built
into most modeling languages including Verilog-A. The issue with this approach
is that it results in sporadic currents in transient simulations since the
current will jump around the distribution every time step. In addition, this
approach will cause multi-state devices to switch sporadically between states
near the critical value if the switching of the device is stochastic. While
this approach may be valuable for some devices, for most devices we will need
another sampling methodology that more accurately reflects the variation we
see in devices. We propose using inverse transform sampling[14] to accomplish
this.

Inverse transform sampling works by first obtaining the cumulative
distribution function (CDF) of the desired distribution and then finding its
inverse. A cumulative distribution function (CDF) is a function that gives the
probability that a random variable is less than or equal to a specific value.
In the context of inverse transform sampling, the input (a uniform random
number between 0 and 1) represents a quantile of the distribution because it
corresponds to a specific probability level, and the inverse CDF maps it to
the corresponding value in the distribution. To find the CDF, we can start
with the probability density function in equation (7).

$p(x)=\sum_{k=1}^{K}\alpha_{k}\cdot\frac{1}{\sqrt{2\pi\sigma_{k}^{2}}}\cdot\exp\left(-\frac{(x-\mu_{k})^{2}}{2\sigma_{k}^{2}}\right)$
(7)

Our CDF will be equal to the integral from $-\infty$ to $x$ as shown in
equation (8).

$F(x)=\int_{-\infty}^{x}\sum_{k=1}^{K}\alpha_{k}\cdot\frac{1}{\sqrt{2\pi\sigma_{k}^{2}}}\cdot\exp\left(-\frac{(t-\mu_{k})^{2}}{2\sigma_{k}^{2}}\right)dt$
(8)

While this is the CDF, it would be beneficial for us to manipulate this into a
more calculable form. By bringing the integral inside the summation and
leaving the scaling factor $\alpha$ outside the integral, we are left with
equation 9.

$F(x)=\sum_{k=1}^{K}\alpha_{k}\int_{-\infty}^{x}\frac{1}{\sqrt{2\pi\sigma_{k}^{2}}}\cdot\exp\left(-\frac{(t-\mu_{k})^{2}}{2\sigma_{k}^{2}}\right)dt$
(9)

This leaves us with the weighted sum of the CDF of a standard normal
distribution, which can be defined in terms of the error function as shown in
equation (10).

$F(x)=\sum_{k=1}^{K}\frac{\alpha_{k}}{2}\left[1+erf\left(\frac{x-\mu_{k}}{\sigma_{k}\sqrt{2}}\right)\right]$
(10)

Defining this in terms of the error function makes calculation much easier
since the error function is well-defined and easy to calculate numerically.
The error function is defined in equation (11).

$erf(x)=\frac{2}{\sqrt{\pi}}\int_{0}^{x}e^{-t^{2}}\,dt$ (11)

The issue is that there is no closed-form solution to the inverse of this CDF.
As such, we use a numerical root finder to find what value of x leads to
$F(x)-q=0$. We choose to use Brent’s method[18], however, most numerical root
finders should work here. For derivative-based root finders, the derivative of
the CDF is equal to the PDF.

Now that we have established a way to perform inverse transform sampling, we
can take full advantage of its properties to improve our model. To solve our
issue with transient simulations, we can keep the same value of $q$ for the
entirety of a sweep. This will result in a continuous output that is more in
line with the device’s properties of cycle-to-cycle variation than traditional
sampling. In this case, we can generate a new $q$ every time the derivative of
an input with respect to time crosses 0, which can easily be done in
Verilog-A. Another benefit of this sampling approach is the ability to clip
the probability distribution. For example, we could generate $q$ only between
0.05 and 0.95 so that the model doesn’t predict values more than 2 standard
deviations away (or in any range). Another use could be modeling device-to-
device variation by weighting the distribution when sampling for $q$, though
this work doesn’t explore this option.

### Device Characteristics of Heater Cryotron

Figure 2:  Prediction of Switching characteristics of heater cryotron. (a) A
false-colored scanning electron microscope (SEM) image of heater cryotron
device consisting two superconducting nanowires of tungsten silicide ($WSi$)
and a dielectric spacer of $SiO_{2}$. (b)-(c) Illustration of the gate-
current-controlled switching mechanism. An enough gate current ($I_{G}$)
switches the gate nanowire to its resistive state and generates thermal phonos
which eventually causes the switching of the channel from its superconducting
to resistive state. (d) Switching model for critical current evaluated at the
mean ($\mu$) and plus or minus 2 standard deviations ($\sigma$) of the
resultant distribution and (e) retrapping current of heater cryotron switching
model evaluated at the same points.

The heater cryotron is a superconducting device designed to exploit the unique
properties of superconducting nanowires [15]. This device consists of two
superconducting nanowires separated by a dielectric, forming the gate and the
channel of the device. Figure 2 (a) shows a false-colored scanning electron
microscope (SEM) image of the fabricated heater cryotron device, with a
channel width of 1$\mu m$ and a gate width 450$nm$. The gate and channel
nanowires are separated by a SiO2 dielectric spacer of 25$nm$ thickness.
Initially, for a given channel current ($I_{B}$), both nanowires (gate and
channel), when maintained at cryogenic temperatures, remain superconducting
exhibiting zero resistance, and remain so until the gate nanowire becomes
resistive (Figure 2 (b)). However, when a sufficient amount of current
($I_{G}$) is passed through the gate nanowire to make it resistive, it starts
generating thermal phonons which transport to the channel nanowire with the
help of the spacer. These thermal phonons suppress the superconductivity and
reduce the critical current of the channel nanowire ($I_{Ch}^{C}$).
Eventually, when the increase of gate current, reduces the channel critical
current below the applied channel current, the channel transitions from its
superconducting state to a resistive state as shown in Figure 2 (c).
Conversely, removing the current from the heater enables the channel nanowire
to revert to its superconducting state. The specific current levels at which
superconducting to resistive and resistive to superconducting transitions
occur are referred to as the critical current and retrapping current,
respectively. The gate current-controlled switching of heater cryotrons has
been used in a multitude of applications, including designing logic circuits
[15, 19, 20, 21], memory systems [22, 23, 24, 25, 26], and neuromorphic
systems [27, 28, 29, 30, 31] for cryogenic environment, and in interfacing
superconducting circuits with semiconducting technology [32].

Notably, the point at which the device switches between its superconducting
and resistive states is characterized by stochastic behavior, which introduces
an inherent element of randomness into its operation. This stochasticity makes
this device ideal for testing our MDN-based compact modeling approach. The
critical current value is not a fixed constant but instead depends on the
applied channel current. For example, a higher bias current leads to a lower
critical current for the gate, thus influencing the switching behavior of the
device. Conversely, a higher bias current leads to a lower retrapping current.
Additionally, a much lower bias current is required to switch from a resistive
to a superconducting state since the nanowire produces its own heat when
resistive among other factors. Understanding this dependence on the channel
current is pivotal in characterizing and modeling the behavior of the heater
cryotron.

To sample the characteristics of the device, we sweep the gate at different
bias currents. By performing each of the sweeps multiple times and recording
the resulting current-voltage characteristics of the device, we can have
enough data for the MDN to learn the stochastic switching behavior of the
device. This experimental approach allows us to explore the device’s response
under varying conditions and analyze the interplay between the gate and
channel currents.

### Heater Cryotron Model

For our first iteration of the model, we chose to use our MDN to model the
critical gate current for a given bias current. This approach allows the use
of a simpler neural network structure for our MDN when compared to learning
the I-V characteristics directly. To accomplish this, we use bias current and
current state of the device as input with the MDN trained to predict gate
critical current (if the heater cryotron is in its superconducting state) and
retrapping current if the heater cryotron is in its resistive state. The
results of this model can be seen in Figure 2 (d) and (e), where we show the
experimental data as well as our model’s output distribution at the mean
($\mu$), and with a variation of twice the standard deviation ($\sigma$) in
either side ($\mu-2\sigma$ and $\mu+2\sigma$) of the mean. We can see that the
model is able to capture the critical and recapture current at different bias
currents as well as the variance in these switching points. Another benefit to
this approach is that in addition to reduced model complexity, we don’t have
to run the model every time step as long as the bias current is constant,
making the model even more efficient in certain applications.

Figure 3:  Model validation for load voltage ($V_{L}$) against the experiment.
The predicted vs experimental distributions of $V_{L}$ (dropped across a 1
$k\Omega$ load resistor) at a bias current of 23.5 $\mu A$ and gate current of
(a) 1.35 $\mu A$, (b) 1.4 $\mu A$, (c) 1.45 $\mu A$, (d) 1.5 $\mu A$, (e) 1.55
$\mu A$, and (f) 1.6 $\mu A$.

While the switching model could work well for some devices, it comes with some
major drawbacks. Most notably, it requires separate models for the I-V
characteristics of the device in its different states in addition to the
switching model. This means that this approach is unable to capture the
variance in the device’s I-V characteristics outside of switching. Going
forward, we are going to use the MDN to directly learn the I-V characteristics
of the heater cryotron. This means that we will be using $I_{G}$, $I_{B}$, and
state as input to the model, and directly predicting $V_{L}$ as the output.

One challenge with MDNs is that there is no good way to provide easily
interpretable numerical metrics for performance. Log-likelihood can be used to
compare different models, but it is not a useful metric on its own. Since
there are not any existing approaches to compare against, this metric will not
be useful for us. Because of this, we will need to rely on graphical
comparisons between the model and the experimental data. First, we can look at
a comparison between the distribution from the model vs a histogram of the
experimental data for different values of $I_{G}$ and $I_{B}$. This is shown
in Figure 3, where we are using $I_{B}$ of 23.5$\mu A$ and values of $I_{G}$
from 1.35$\mu A$ to 1.6$\mu A$. Using this range of $I_{G}$ allows us to see
how the switching probability distributions change around the critical
current. Figure 3 shows that our probability distributions closely match the
distributions obtained from the experiment.

Next, we can compare the switching probability of the model. To do this, we
use a sweep of $I_{G}$ from 0$\mu A$ to 3$\mu A$ for $I_{B}$ values of 14$\mu
A$, 16.5, 23.5$\mu A$, 28$\mu A$, and 33$\mu A$. To evaluate the switching
probability of the model, we can use the CDF evaluated at the voltage midway
between the load voltage when the H-Tron is in superconducting state and when
it is in resistive state. Figure 4 (left) shows the results of switching
probability observed in the experimental data compared to the MDN model.
Framing the model in this way allows us to use traditional regression
performance metrics, which can be seen in Table 1. On average our model is off
by 0.82% for switching probability with an $R^{2}$ of 0.9891. Here, the
$R^{2}$ value represents the percentage of variance in the switching
probability explained by the model.

Figure 4:  Switching probabilities at varying bias currents and transient simulation results. The predicted vs experimental switching probability for gate current ($I_{G}$) between 0 $\mu A$ and 3 $\mu A$ at a bias current ($I_{B}$) of (a) 14 $\mu A$, (b) 16.5 $\mu A$, (c) 23.5 $\mu A$, (d) 28 $\mu A$, and (e) 33 $\mu A$. The increase in $I_{B}$ leads to a reduction in the gate switching current. Transient dynamics of (f) bias current and (g) gate current which cause switching of the channel after exceeding the corresponding thresholds. (h) Switching is shown and validated with the experiment by the time dynamics of the load voltage ($V_{L}$) at mean ($\mu$) and plus or minus 2 standard deviations ($\sigma$). (i) A zoomed-in view of $V_{L}$ shown in (h). Table 1:  Model Performance of Switching Probability  $I_{B}(\mu A)$ |  Mean Absolute Error (MAE) |  $R^{2}$
---|---|---
14 |  1.149% |  0.9941
16.5 |  0.534% |  0.9986
23.5 |  1.344% |  0.9572
28 |  0.611% |  0.9986
33 |  0.662% |  0.9971

Finally, we can evaluate the performance of the model in transient simulation.
So that we can match the output of the model with experimental results, we
mirror $I_{G}$ and $I_{B}$ from the experimental data as shown in Figure 6
(a). We use the inverse transform sampling technique as described in section
II to achieve the smooth, continuous curves as seen in Figure 4 (right). To
show the distribution of the model in the transient simulation, we report
$\mu$, $\mu-2\sigma$, and $\mu+2\sigma$ conditions. We can do this by setting
the quantile value ($q$) to 0.5 for the mean and 0.05 and 0.95 for the inverse
transform sampling. Figure 4 (right) shows that our MDN model accurately
captures the variance in the I-V characteristics of the devices, including the
switching dynamics.

## Discussion

Our approach utilizing mixture density networks has demonstrated its
remarkable ability to accurately model the variance of I-V characteristics and
stochastic switching behavior of heater cryotrons. With a 0.82% mean absolute
error in the switching probability and an $R^{2}$ value of 0.9891, our method
has proven its efficacy in capturing the true variance of device behavior.
While we’ve only shown the model using heater cryotrons, the generalizing
power of neural networks means this approach will be easily adaptable to other
devices. This achievement marks a significant stride towards more precise and
reliable electronic circuit simulations, offering unprecedented opportunities
for the development of cutting-edge technologies in an era where stochasticity
plays an increasingly pivotal role.

## Methods

### Heater Cryotron Fabrication

The fabrication process began with the deposition of a 4 $nm$ layer of
superconducting tungsten silicide (WSi) onto a silicon wafer through magnetron
sputtering. Following this, contact pads consisting of 5 $nm$ Ti and 30 $nm$
Au were applied using optical lithography and a liftoff procedure.
Subsequently, electron beam lithography was employed to define and etch the
WSi layer in an RIE process with SF6 chemistry to create the device’s channel.
A 25 $nm$ layer of SiO2 was then sputter-deposited across the entire wafer to
serve as the dielectric spacer separating the heater from the channel. Another
4 $nm$ layer of WSi was subsequently sputter-deposited, patterned, and etched
to form the heater input gate, with contact pads added to this upper layer
through the same liftoff process as the initial layer. Finally, openings were
etched through the SiO2 dielectric layer using an RIE process with CHF3:O2
chemistry to establish electrical contact with the underlying channel layer.

### Heater Cryotron Characterization

The measurements in this study employed an arbitrary waveform generator (AWG)
equipped with two channels and incorporated 10 $k\Omega$ resistors in series
with each channel. In the data acquisition process, one channel of the AWG was
configured to maintain a constant voltage, thereby delivering a fixed bias
current, while the other channel was programmed to incrementally ramp up from
zero current to 3 $\mu A$. Simultaneously, the voltage of the device channel
was meticulously recorded on an oscilloscope. The experimental conditions
covered six distinct bias current levels, spanning from 14 $\mu A$ to 33 $\mu
A$. Following the ramp-up of gate current, a subsequent phase involved the
concurrent reduction of gate and bias currents to get a larger breadth of
data. This entire procedure was repeated 1000 times for each bias current
setting, ensuring the capture of the stochastic nature of the device.

### Model Creation

We use Python with TensorFlow to create the mixture density network. We train
the model using the data derived in the previous section. We implement a
custom loss function to optimize the model for minimizing gaussian negative
log-likelihood. In addition, we implement a custom activation function for
exponential linear unit plus 1 plus epsilon as described in equation (3). In
order to use this model in circuit simulation, we implement the neural network
architecture and the inverse transform sampling methodology in Verilog-A. The
trained model weights are then exported and utilized in the Verilog-A model.
This process is derived from Hutchins et. al. [10].

## References

  * [1] Salahuddin, S., Ni, K. & Datta, S.  The era of hyper-scaling in electronics.  _Nature Electronics_ 1, 442–450 (2018).
  * [2] Lundstrom, M.  Moore’s law forever?  _Science_ 299, 210–211 (2003).
  * [3] Alawad, M. & Lin, M.  Survey of stochastic-based computation paradigms.  _IEEE Transactions on Emerging Topics in Computing_ 7, 98–114 (2016).
  * [4] Hamilton, T. J., Afshar, S., van Schaik, A. & Tapson, J.  Stochastic electronics: A neuro-inspired design paradigm for integrated circuits.  _Proceedings of the IEEE_ 102, 843–859 (2014).
  * [5] Li, M., Irsoy, O., Cardie, C. & Xing, H. G.  Physics-inspired neural networks for efficient device compact modeling.  _IEEE Journal on Exploratory Solid-State Computational Devices and Circuits_ 2, 44–49, DOI: 10.1109/JXCDC.2016.2636161 (2016).
  * [6] Zhang, L. & Chan, M.  Artificial neural network design for compact modeling of generic transistors.  _Journal of Computational Electronics_ 16, 825–832, DOI: 10.1007/s10825-017-0984-9 (2017).
  * [7] Wang, J. _et al._ Artificial neural network-based compact modeling methodology for advanced transistors.  _IEEE Transactions on Electron Devices_ 68, 1318–1325, DOI: 10.1109/TED.2020.3048918 (2021).
  * [8] Abouelyazid, M. S., Hammouda, S. & Ismail, Y.  Fast and accurate machine learning compact models for interconnect parasitic capacitances considering systematic process variations.  _IEEE Access_ 10, 7533–7553, DOI: 10.1109/ACCESS.2022.3142330 (2022).
  * [9] Lin, A. S. _et al._ Rram compact modeling using physics and machine learning hybridization.  _IEEE Transactions on Electron Devices_ 69, 1835–1841, DOI: 10.1109/TED.2022.3152978 (2022).
  * [10] Hutchins, J. _et al._ A generalized workflow for creating machine learning-powered compact models for multi-state devices.  _IEEE Access_ DOI: 10.1109/ACCESS.2022.3218333 (2022).
  * [11] Wang, Y. _et al._ Compact model of magnetic tunnel junction with stochastic spin transfer torque switching for reliability analyses.  _Microelectronics Reliability_ 54, 1774–1778, DOI: 10.1016/j.microrel.2014.07.019 (2014).
  * [12] Becle, E., Talatchian, P., Prenat, G., Anghel, L. & Prejbeanu, I.-L.  Fast behavioral veriloga compact model for stochastic mtj.  259–262, DOI: 10.1109/ESSDERC53440.2021.9631829 (IEEE, 2021).
  * [13] Bishop, C.  Mixture density networks.  WorkingPaper, Aston University (1994).
  * [14] Devroye, L.  _Non-Uniform Random Variate Generation_ (Springer New York, 1986).
  * [15] McCaughan, A. N. & Berggren, K. K.  A superconducting-nanowire three-terminal electrothermal device.  _Nano letters_ 14, 5748–5753 (2014).
  * [16] Hornik, K., Stinchcombe, M. & White, H.  Multilayer feedforward networks are universal approximators.  _Neural Networks_ 2, 359–366, DOI: 10.1016/0893-6080(89)90020-8 (1989).
  * [17] Loshchilov, I. & Hutter, F.  Decoupled weight decay regularization (2019).  1711.05101.
  * [18] Brent, R. P.  _Chapter 4: An Algorithm with Guaranteed Convergence for Finding a Zero of a Function_ (Prentice-Hall, 1973).
  * [19] Alam, S., Hossain, M. S., Ni, K., Narayanan, V. & Aziz, A.  Voltage-controlled cryogenic boolean logic family based on ferroelectric squid.  _arXiv preprint arXiv:2212.08202_ (2022).
  * [20] Alam, S., Rampini, D. S., Oripov, B. G., McCaughan, A. N. & Aziz, A.  Cryogenic reconfigurable logic with superconducting heater cryotron: Enhancing area efficiency and enabling camouflaged processors.  _Applied Physics Letters_ 123 (2023).
  * [21] Alam, S., McCaughan, A. N. & Aziz, A.  Reconfigurable superconducting logic using multi-gate switching of a nano-cryotron.  In _2023 Device Research Conference (DRC)_ , 1–2 (IEEE, 2023).
  * [22] Alam, S., Hossain, M. S., Srinivasa, S. R. & Aziz, A.  Cryogenic memory technologies.  _Nature Electronics_ 6, 185–198 (2023).
  * [23] Alam, S., Hossain, M. S. & Aziz, A.  A cryogenic memory array based on superconducting memristors.  _Applied Physics Letters_ 119 (2021).
  * [24] Alam, S. _et al._ Cryogenic memory array based on ferroelectric squid and heater cryotron.  In _2022 Device Research Conference (DRC)_ , 1–2 (IEEE, 2022).
  * [25] Alam, S., Islam, M. M., Hossain, M. S. & Aziz, A.  Superconducting josephson junction fet-based cryogenic voltage sense amplifier.  In _2022 Device Research Conference (DRC)_ , 1–2 (IEEE, 2022).
  * [26] Alam, S. _et al._ Cryogenic in-memory matrix-vector multiplication using ferroelectric superconducting quantum interference device (fe-squid).  In _2023 60th ACM/IEEE Design Automation Conference (DAC)_ , 1–6 (IEEE, 2023).
  * [27] Islam, M. M., Alam, S., Hossain, M. S., Roy, K. & Aziz, A.  A review of cryogenic neuromorphic hardware.  _Journal of Applied Physics_ 133 (2023).
  * [28] Islam, M. M., Alam, S., Hossain, M. S. & Aziz, A.  Dynamically reconfigurable cryogenic spiking neuron based on superconducting memristor.  In _2022 IEEE 22nd International Conference on Nanotechnology (NANO)_ , 307–310 (IEEE, 2022).
  * [29] Islam, M. M., Alam, S., Schuman, C. D., Hossain, M. S. & Aziz, A.  A deep dive into the design space of a dynamically reconfigurable cryogenic spiking neuron.  _IEEE Transactions on Nanotechnology_ (2023).
  * [30] Islam, M. M., Alam, S., Shukla, N. & Aziz, A.  Design space analysis of superconducting nanowire-based cryogenic oscillators.  In _2022 Device Research Conference (DRC)_ , 1–2 (IEEE, 2022).
  * [31] Islam, M. M., Alam, S., Udoy, M. R. I., Hossain, M. S. & Aziz, A.  A cryogenic artificial synapse based on superconducting memristor.  In _Proceedings of the Great Lakes Symposium on VLSI 2023_ , 143–148 (2023).
  * [32] McCaughan, A. N. _et al._ A superconducting thermal switch with ultrahigh impedance for interfacing superconductors to semiconductors.  _Nature electronics_ 2, 451–456 (2019).

## Acknowledgements

This research was funded by the University of Tennessee Knoxville
(https://ror.org/020f3ap87) and NIST (https://ror.org/ 05xpvk416). The U.S.
Government is authorized to reproduce and distribute reprints for governmental
purposes notwithstanding any copyright annotation thereon.

S. A. was supported with funds provided by the Science Alliance, a Tennessee
Higher Education Commission center of excellence administered by The
University of Tennessee-Oak Ridge Innovation Institute on behalf of The
University of Tennessee, Knoxville.

## Data Availability Statement

The data that support the plots within this paper and other findings of this
study are available from the corresponding author upon reasonable request.

## Author Contributions Statement

J.H. conceived the idea and developed the modeling framework. D.S.R and B.G.O.
fabricated and characterized the heater cryotron device. A.A. supervised the
project. All authors analyzed the results and contributed to writing the
manuscript.


$|\bar{a}_{\varepsilon}^{\prime}(\tau)|=-a_{\varepsilon}^{\prime}(\tau)\geqslant
C\tau$, so

$\displaystyle\bigg{|}\frac{a_{\varepsilon}^{\prime}(\tau)}{\bar{a}_{\varepsilon}^{\prime}(\tau)}\bigg{|}$
$\displaystyle=\frac{a_{\varepsilon}^{\prime}(\tau)}{|\bar{a}_{\varepsilon}^{\prime}(\tau)|}\leqslant
C\varepsilon,\qquad\text{for all }\tau\in(0,1),$

which, along with (7.16), implies the inequality for $a_{\varepsilon}$ in
(7.15).

Gathering these pieces of information, we obtain that

(7.18)
$\begin{split}\frac{\|A_{\varepsilon}(|\tilde{r}|)-A_{\varepsilon}(|r|)\|_{\mathrm{OP}}}{|\phi_{\varepsilon}(r)-\phi_{\varepsilon}(\tilde{r})|}\leqslant\;&\frac{C\sqrt{\big{(}(a_{\varepsilon}(|r|)-a_{\varepsilon}(|\tilde{r}|)\big{)}^{2}+\big{(}b_{\varepsilon}(|r|)-b_{\varepsilon}(|\tilde{r}|)\big{)}^{2}}}{\sqrt{\Big{(}a_{\varepsilon}(|r|)\sqrt{1-|r|^{2}}-a_{\varepsilon}(|\tilde{r}|)\sqrt{1-|\tilde{r}|^{2}}\Big{)}^{2}+\big{(}b_{\varepsilon}(|r|)|r|-b_{\varepsilon}(|\tilde{r}|)|\tilde{r}|\big{)}^{2}}}\\\
\leqslant\;&C\varepsilon.\end{split}$

Finally, since $G_{\varepsilon}\subset\Omega_{\varepsilon}\subset
B_{1+\varepsilon}$, we have that

(7.19) $\displaystyle|\phi_{\varepsilon}(\tilde{r})|$ $\displaystyle\leqslant
1+\varepsilon\leqslant C.$

Thus, (7.10) follows from (7.13), (7.14), (7.18), and (7.19). ∎

######  Lemma 7.6.

F3lTRgfZ Let $\varepsilon\in[0,1)$ and $a_{\varepsilon}$, $b_{\varepsilon}$,
and $\phi_{\varepsilon}$ be as in M2fW7tQX.

Then, for all $r,\tilde{r}\in B_{1}^{n-1}$ with $r\neq\tilde{r}$,

$\displaystyle\bigg{|}\frac{|(u_{\varepsilon}\circ\phi_{\varepsilon})(r)-(u_{\varepsilon}\circ\phi_{\varepsilon})(\tilde{r})|}{\varepsilon|\phi_{0}(r)-\phi_{0}(\tilde{r})|}-\frac{1}{2}s\gamma_{n,s}\left(\frac{3}{4}\right)^{s-1}\frac{\big{|}|r|^{2}-|\tilde{r}|^{2}\big{|}}{|\phi_{0}(r)-\phi_{0}(\tilde{r})|}\bigg{|}\leqslant
C\varepsilon.$

The constant $C>0$ depends only on $n$ and $s$.

###### Proof.

Let

$\displaystyle v_{\varepsilon}(x)$
$\displaystyle:=\frac{1}{\gamma_{n,s,\varepsilon}}u_{\varepsilon}(x)=\left(1-\frac{x_{1}^{2}}{(1+\varepsilon)^{2}}-|x^{\prime}|^{2}\right)^{s}$

and define

$\psi_{\varepsilon}(\tau):=\frac{\displaystyle\left(1-\frac{(a_{\varepsilon}(\tau))^{2}(1-\tau^{2})}{(1+\varepsilon)^{2}}-(b_{\varepsilon}(\tau))^{2}\tau^{2}\right)^{s}-\left(\frac{3}{4}\right)^{s}}{\varepsilon}+\frac{1}{2}s\left(\frac{3}{4}\right)^{s-1}(1-\tau^{2}).$

Exploiting the expressions of $a_{\varepsilon}$ and $b_{\varepsilon}$ in
(7.7), we compute the derivative of $\psi_{\varepsilon}$ with respect to
$\tau$ as

$\displaystyle\frac{d}{d\tau}\psi_{\varepsilon}(\tau)$ $\displaystyle=$
$\displaystyle
s\left[-\frac{\tau^{3}(1+\varepsilon)\left(2\varepsilon+\varepsilon^{2}\right)\left(1-\frac{1+\varepsilon}{2\sqrt{1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)}}\right)}{\left(1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)\right)^{3/2}}+\frac{2\tau\left(1+\varepsilon-\frac{1}{2\sqrt{1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)}}\right)^{2}}{(1+\varepsilon)^{2}}\right.$
$\displaystyle\quad\left.-2\tau\left(1-\frac{1+\varepsilon}{2\sqrt{1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)}}\right)^{2}-\frac{\left(1-\tau^{2}\right)\tau\left(2\varepsilon+\varepsilon^{2}\right)\left(1+\varepsilon-\frac{1}{2\sqrt{1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)}}\right)}{(1+\varepsilon)^{2}\left(1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)\right)^{3/2}}\right]$
$\displaystyle\times\left[1-\frac{\left(1-\tau^{2}\right)\left(1+\varepsilon-\frac{1}{2\sqrt{1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)}}\right)^{2}}{(1+\varepsilon)^{2}}-\tau^{2}\left(1-\frac{1+\varepsilon}{2\sqrt{1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)}}\right)^{2}\right]^{s-1}.$

We claim that, for $\varepsilon$ sufficiently small,

(7.20)
$\left\|\frac{d}{d\tau}\psi_{\varepsilon}\right\|_{L^{\infty}((0,1))}\leqslant
C\varepsilon,$

for some $C>0$ depending only on $s$. Not to interrupt the flow of the
argument, we defer the proof of (7.20) to Appendix A.

Now, we notice that

$\frac{(v_{\varepsilon}\circ\phi_{\varepsilon})(r)-\left(\frac{3}{4}\right)^{s}}{\varepsilon}+\frac{1}{2}s\left(\frac{3}{4}\right)^{s-1}(1-|r|^{2})=\psi_{\varepsilon}(|r|).$

Therefore,

$\displaystyle\bigg{|}\frac{|(v_{\varepsilon}\circ\phi_{\varepsilon})(r)-(v_{\varepsilon}\circ\phi_{\varepsilon})(\tilde{r})|}{\varepsilon|\phi_{0}(r)-\phi_{0}(\tilde{r})|}-\frac{1}{2}s\left(\frac{3}{4}\right)^{s-1}\frac{\big{|}|r|^{2}-|\tilde{r}|^{2}\big{|}}{|\phi_{0}(r)-\phi_{0}(\tilde{r})|}\bigg{|}$
$\displaystyle\leqslant$
$\displaystyle\bigg{|}\frac{\varepsilon^{-1}\left[(v_{\varepsilon}\circ\phi_{\varepsilon})(r)-\left(\frac{3}{4}\right)^{s}\right]-\frac{1}{2}s\left(\frac{3}{4}\right)^{s-1}|r|^{2}}{|\phi_{0}(r)-\phi_{0}(\tilde{r})|}-\frac{\varepsilon^{-1}\left[(v_{\varepsilon}\circ\phi_{\varepsilon})(\tilde{r})-\left(\frac{3}{4}\right)^{s}\right]-\frac{1}{2}s\left(\frac{3}{4}\right)^{s-1}|\tilde{r}|^{2}}{|\phi_{0}(r)-\phi_{0}(\tilde{r})|}\bigg{|}$
$\displaystyle=$
$\displaystyle\bigg{|}\frac{\varepsilon^{-1}\left[(v_{\varepsilon}\circ\phi_{\varepsilon})(r)-\left(\frac{3}{4}\right)^{s}\right]+\frac{1}{2}s\left(\frac{3}{4}\right)^{s-1}(1-|r|^{2})}{|\phi_{0}(r)-\phi_{0}(\tilde{r})|}-\frac{\varepsilon^{-1}\left[(v_{\varepsilon}\circ\phi_{\varepsilon})(\tilde{r})-\left(\frac{3}{4}\right)^{s}\right]+\frac{1}{2}s\left(\frac{3}{4}\right)^{s-1}(1-|\tilde{r}|^{2})}{|\phi_{0}(r)-\phi_{0}(\tilde{r})|}\bigg{|}$
$\displaystyle=$
$\displaystyle\frac{|\psi_{\varepsilon}(|r|)-\psi_{\varepsilon}(|\tilde{r}|)|}{|\phi_{0}(r)-\phi_{0}(\tilde{r})|}.$

Moreover,

$\displaystyle|\phi_{0}(r)-\phi_{0}(\tilde{r})|^{2}$
$\displaystyle=\frac{1}{4}\Big{(}\big{(}\sqrt{1-|r|^{2}}-\sqrt{1-|\tilde{r}|^{2}}\big{)}^{2}+|r-\tilde{r}|^{2}\Big{)}\geqslant
C\big{|}|r|-|\tilde{r}|\big{|}^{2},$

so, recalling (7.20), we conclude that

(7.21)
$\begin{split}&\bigg{|}\frac{|(v_{\varepsilon}\circ\phi_{\varepsilon})(r)-(v_{\varepsilon}\circ\phi_{\varepsilon})(\tilde{r})|}{\varepsilon|\phi_{0}(r)-\phi_{0}(\tilde{r})|}-\frac{1}{2}s\left(\frac{3}{4}\right)^{s-1}\frac{\big{|}|r|^{2}-|\tilde{r}|^{2}\big{|}}{|\phi_{0}(r)-\phi_{0}(\tilde{r})|}\bigg{|}\leqslant
C\frac{|\psi_{\varepsilon}(|r|)-\psi_{\varepsilon}(|\tilde{r}|)|}{\big{|}|r|-|\tilde{r}|\big{|}}\\\
&\qquad\qquad\qquad\leqslant
C\left\|\frac{d}{d\tau}\psi_{\varepsilon}\right\|_{L^{\infty}((0,1))}\leqslant
C\varepsilon.\end{split}$

Finally, note that $\gamma_{n,s,\varepsilon}$ depends smoothly on
$\varepsilon$ for $\varepsilon$ small and
$\gamma_{n,s,\varepsilon}\to\gamma_{n,s}$ as $\varepsilon\to 0^{+}$, which
implies that

(7.22) $\displaystyle|\gamma_{n,s,\varepsilon}-\gamma_{n,s}|$
$\displaystyle\leqslant C\varepsilon.$

Combing (7.22) with (7.21) gives the final result. ∎

We can now prove M2fW7tQX.

###### Proof of M2fW7tQX.

By Lemmata LABEL:c3Ge7XSV and LABEL:F3lTRgfZ, for all $r,\tilde{r}\in
B_{1}^{n-1}$ with $r\neq\tilde{r}$,

$\displaystyle\bigg{|}$
$\displaystyle\frac{|(u_{\varepsilon}\circ\phi_{\varepsilon})(r)-(u_{\varepsilon}\circ\phi_{\varepsilon})(\tilde{r})|}{\varepsilon|\phi_{\varepsilon}(r)-\phi_{\varepsilon}(\tilde{r})|}-\frac{1}{2}s\gamma_{n,s}\left(\frac{3}{4}\right)^{s-1}\frac{\big{|}|r|^{2}-|\tilde{r}|^{2}\big{|}}{|\phi_{0}(r)-\phi_{0}(\tilde{r})|}\bigg{|}$
$\displaystyle\leqslant\bigg{|}\frac{|(u_{\varepsilon}\circ\phi_{\varepsilon})(r)-(u_{\varepsilon}\circ\phi_{\varepsilon})(\tilde{r})|}{\varepsilon|\phi_{0}(r)-\phi_{0}(\tilde{r})|}\bigg{|}\cdot\bigg{|}\frac{|\phi_{0}(r)-\phi_{0}(\tilde{r})|}{|\phi_{\varepsilon}(r)-\phi_{\varepsilon}(\tilde{r})|}-1\bigg{|}$
$\displaystyle\qquad+\bigg{|}\frac{|(u_{\varepsilon}\circ\phi_{\varepsilon})(r)-(u_{\varepsilon}\circ\phi_{\varepsilon})(\tilde{r})|}{\varepsilon|\phi_{0}(r)-\phi_{0}(\tilde{r})|}-\frac{1}{2}s\gamma_{n,s}\left(\frac{3}{4}\right)^{s-1}\frac{\big{|}|r|^{2}-|\tilde{r}|^{2}\big{|}}{|\phi_{0}(r)-\phi_{0}(\tilde{r})|}\bigg{|}$
$\displaystyle\leqslant
C\varepsilon\bigg{(}\bigg{|}\frac{|(u_{\varepsilon}\circ\phi_{\varepsilon})(r)-(u_{\varepsilon}\circ\phi_{\varepsilon})(\tilde{r})|}{\varepsilon|\phi_{0}(r)-\phi_{0}(\tilde{r})|}\bigg{|}+1\bigg{)}.$

Moreover, by F3lTRgfZ, for $\varepsilon$ small,

$\displaystyle\bigg{|}\frac{|(u_{\varepsilon}\circ\phi_{\varepsilon})(r)-(u_{\varepsilon}\circ\phi_{\varepsilon})(\tilde{r})|}{\varepsilon|\phi_{0}(r)-\phi_{0}(\tilde{r})|}\bigg{|}$
$\displaystyle\leqslant
C\bigg{(}\frac{\big{|}|r|^{2}-|\tilde{r}|^{2}\big{|}}{|\phi_{0}(r)-\phi_{0}(\tilde{r})|}+1\bigg{)}\leqslant
C$

using (7.9). This completes the proof. ∎

##  Appendix A Proof of the claim in (7.20)

We write

(A.1)
$\frac{d}{d\tau}\psi_{\varepsilon}(\tau)=sg_{\varepsilon}(\tau)(h_{\varepsilon}(\tau))^{s-1},$

where

$\displaystyle g_{\varepsilon}(\tau)$ $\displaystyle:=$
$\displaystyle-\frac{\tau^{3}(1+\varepsilon)\left(2\varepsilon+\varepsilon^{2}\right)\left(1-\frac{1+\varepsilon}{2\sqrt{1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)}}\right)}{\left(1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)\right)^{3/2}}+\frac{2\tau\left(1+\varepsilon-\frac{1}{2\sqrt{1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)}}\right)^{2}}{(1+\varepsilon)^{2}}$
$\displaystyle\quad-2\tau\left(1-\frac{1+\varepsilon}{2\sqrt{1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)}}\right)^{2}-\frac{\left(1-\tau^{2}\right)\tau\left(2\varepsilon+\varepsilon^{2}\right)\left(1+\varepsilon-\frac{1}{2\sqrt{1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)}}\right)}{(1+\varepsilon)^{2}\left(1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)\right)^{3/2}}$
$\displaystyle{\mbox{and }}\qquad h_{\varepsilon}(\tau)$ $\displaystyle:=$
$\displaystyle
1-\frac{\left(1-\tau^{2}\right)\left(1+\varepsilon-\frac{1}{2\sqrt{1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)}}\right)^{2}}{(1+\varepsilon)^{2}}-\tau^{2}\left(1-\frac{1+\varepsilon}{2\sqrt{1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)}}\right)^{2}.$

We first prove that

(A.2) $|g_{\varepsilon}(\tau)|\leqslant C\varepsilon,$

for all $\tau\in(0,1)$, for some constant $C>0$.

For this, we observe that

(A.3)
$\begin{split}&\left|-\frac{\tau^{3}(1+\varepsilon)\left(2\varepsilon+\varepsilon^{2}\right)\left(1-\frac{1+\varepsilon}{2\sqrt{1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)}}\right)}{\left(1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)\right)^{3/2}}\right|\\\
&\qquad\qquad\leqslant\left|\tau^{3}(1+\varepsilon)\left(2\varepsilon+\varepsilon^{2}\right)\left(1-\frac{1+\varepsilon}{2\sqrt{1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)}}\right)\right|\\\
&\qquad\qquad\leqslant
6\varepsilon\left|1-\frac{1+\varepsilon}{2\sqrt{1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)}}\right|\leqslant
12\varepsilon.\end{split}$

Similarly,

(A.4)
$\begin{split}&\left|-\frac{\left(1-\tau^{2}\right)\tau\left(2\varepsilon+\varepsilon^{2}\right)\left(1+\varepsilon-\frac{1}{2\sqrt{1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)}}\right)}{(1+\varepsilon)^{2}\left(1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)\right)^{3/2}}\right|\\\
&\qquad\leqslant\left|\left(1-\tau^{2}\right)\tau\left(2\varepsilon+\varepsilon^{2}\right)\left(1+\varepsilon-\frac{1}{2\sqrt{1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)}}\right)\right|\\\
&\qquad\leqslant
3\varepsilon\left|1+\varepsilon-\frac{1}{2\sqrt{1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)}}\right|\leqslant
9\varepsilon.\end{split}$

Additionally,

$\begin{split}&2\tau\left[\frac{\left(1+\varepsilon-\frac{1}{2\sqrt{1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)}}\right)^{2}}{(1+\varepsilon)^{2}}-\left(1-\frac{1+\varepsilon}{2\sqrt{1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)}}\right)^{2}\right]\\\
=\;&2\tau\left[\left(1+\varepsilon-\frac{1}{2\sqrt{1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)}}\right)^{2}-\left(1-\frac{1+\varepsilon}{2\sqrt{1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)}}\right)^{2}\right]\\\
&\qquad+2\tau\left(1+\varepsilon-\frac{1}{2\sqrt{1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)}}\right)^{2}\left[\frac{1}{(1+\varepsilon)^{2}}-1\right]\\\
=\;&2\tau\left[\left(1+\varepsilon-\frac{1}{2\sqrt{1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)}}\right)^{2}-\left(1-\frac{1+\varepsilon}{2\sqrt{1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)}}\right)^{2}\right]\\\
&\qquad-2\tau\left(1+\varepsilon-\frac{1}{2\sqrt{1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)}}\right)^{2}\frac{2\varepsilon+\varepsilon^{2}}{(1+\varepsilon)^{2}}.\end{split}$

We observe that

$\begin{split}&\left(1+\varepsilon-\frac{1}{2\sqrt{1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)}}\right)^{2}-\left(1-\frac{1+\varepsilon}{2\sqrt{1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)}}\right)^{2}\\\
=\;&\left[\left(1+\varepsilon-\frac{1}{2\sqrt{1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)}}\right)+\left(1-\frac{1+\varepsilon}{2\sqrt{1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)}}\right)\right]\\\
&\qquad\times\left[\left(1+\varepsilon-\frac{1}{2\sqrt{1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)}}\right)-\left(1-\frac{1+\varepsilon}{2\sqrt{1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)}}\right)\right]\\\
=\;&\left(2+\varepsilon-\frac{2+\varepsilon}{2\sqrt{1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)}}\right)\left(\varepsilon+\frac{\varepsilon}{2\sqrt{1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)}}\right)\\\
=\;&\varepsilon(2+\varepsilon)\left(1-\frac{1}{2\sqrt{1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)}}\right)\left(1+\frac{1}{2\sqrt{1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)}}\right)\\\
=\;&\varepsilon(2+\varepsilon)\left(1-\frac{1}{4\big{(}1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)\big{)}}\right),\end{split}$

and therefore

$\begin{split}&\left|2\tau\left[\left(1+\varepsilon-\frac{1}{2\sqrt{1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)}}\right)^{2}-\left(1-\frac{1+\varepsilon}{2\sqrt{1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)}}\right)^{2}\right]\right|\\\
&\qquad\leqslant
6\varepsilon\left|1-\frac{1}{4\big{(}1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)\big{)}}\right|\leqslant
12\varepsilon.\end{split}$

Furthermore,

$\displaystyle\left|-2\tau\left(1+\varepsilon-\frac{1}{2\sqrt{1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)}}\right)^{2}\frac{2\varepsilon+\varepsilon^{2}}{(1+\varepsilon)^{2}}\right|\leqslant
6\varepsilon\left|\left(1+\varepsilon-\frac{1}{2\sqrt{1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)}}\right)^{2}\right|\leqslant
54\varepsilon.$

Gathering these pieces of information, we obtain that

$\displaystyle\left|2\tau\left[\frac{\left(1+\varepsilon-\frac{1}{2\sqrt{1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)}}\right)^{2}}{(1+\varepsilon)^{2}}-\left(1-\frac{1+\varepsilon}{2\sqrt{1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)}}\right)^{2}\right]\right|\leqslant
66\varepsilon.$

From this, (A.3) and (A.4) we obtain the claim in (A.2).

We now show that, for $\varepsilon$ sufficiently small,

(A.5) $h_{\varepsilon}(\tau)\geqslant c,$

for all $\tau\in(0,1)$, for some $c>0$.

For this, we observe that

$\begin{split}&1-\frac{\left(1-\tau^{2}\right)\left(1+\varepsilon-\frac{1}{2\sqrt{1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)}}\right)^{2}}{(1+\varepsilon)^{2}}-\tau^{2}\left(1-\frac{1+\varepsilon}{2\sqrt{1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)}}\right)^{2}\\\
=\;&1-\frac{\left(1+\varepsilon-\frac{1}{2\sqrt{1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)}}\right)^{2}}{(1+\varepsilon)^{2}}+\tau^{2}\left[\frac{\left(1+\varepsilon-\frac{1}{2\sqrt{1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)}}\right)^{2}}{(1+\varepsilon)^{2}}-\left(1-\frac{1+\varepsilon}{2\sqrt{1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)}}\right)^{2}\right]\\\
=\;&\frac{1+\varepsilon}{(1+\varepsilon)^{2}\sqrt{1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)}}-\frac{1}{4(1+\varepsilon)^{2}\big{(}{1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)}\big{)}}\\\
&\qquad+\tau^{2}\left[\frac{\left(1+\varepsilon-\frac{1}{2\sqrt{1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)}}\right)^{2}}{(1+\varepsilon)^{2}}-\left(1-\frac{1+\varepsilon}{2\sqrt{1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)}}\right)^{2}\right]\\\
=\;&\frac{1}{(1+\varepsilon)\sqrt{1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)}}-\frac{1}{4(1+\varepsilon)^{2}\big{(}{1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)}\big{)}}\\\
&\qquad+\tau^{2}\varepsilon(2+\varepsilon)\left[1-\frac{1}{4\big{(}1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)\big{)}}-\left(1+\varepsilon-\frac{1}{2\sqrt{1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)}}\right)^{2}\frac{1}{(1+\varepsilon)^{2}}\right]\\\
=\;&\frac{1}{(1+\varepsilon)\sqrt{1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)}}-\frac{1}{4(1+\varepsilon)^{2}\big{(}{1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)}\big{)}}\\\
&\qquad+\tau^{2}\frac{\varepsilon(2+\varepsilon)}{(1+\varepsilon)^{2}}\left[\frac{1+\varepsilon}{\sqrt{1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)}}-\frac{1+(1+\varepsilon)^{2}}{4\big{(}{1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)}\big{)}}\right]\\\
\geqslant\;&\frac{1}{(1+\varepsilon)^{2}}-\frac{1}{4(1+\varepsilon)^{2}}-\frac{\varepsilon(2+\varepsilon)}{(1+\varepsilon)^{2}}\left|\frac{1+\varepsilon}{\sqrt{1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)}}-\frac{1+(1+\varepsilon)^{2}}{4\big{(}{1+\tau^{2}\left(2\varepsilon+\varepsilon^{2}\right)}\big{)}}\right|\\\
\geqslant\;&\frac{1}{(1+\varepsilon)^{2}}\left(1-\frac{1}{4}-\frac{\varepsilon(2+\varepsilon)(4+4\varepsilon+\varepsilon^{2})}{2}\right)\\\
\geqslant\;&\frac{1}{2},\end{split}$

for $\varepsilon$ sufficiently small. This establishes (A.5).

From (A.1), (A.2) and (A.5), we obtain the desired claim in (7.20).

## References

  * [AJS21] Nicola Abatangelo, Sven Jarohs and Alberto Saldaña  “Fractional Laplacians on ellipsoids”  In _Math. Eng._ 3.5, 2021, pp. Paper No. 03834  DOI: 10.3934/mine.2021038
  * [Ale62] A.. Aleksandrov  “Uniqueness theorems for surfaces in the large. V”  In _Amer. Math. Soc. Transl. (2)_ 21, 1962, pp. 412–416
  * [AS92] “Handbook of mathematical functions with formulas, graphs, and mathematical tables” Reprint of the 1972 edition  Dover Publications, Inc., New York, 1992, pp. xiv+1046
  * [BD13] I. Birindelli and F. Demengel  “Overdetermined problems for some fully non linear operators”  In _Comm. Partial Differential Equations_ 38.4, 2013, pp. 608–628  DOI: 10.1080/03605302.2012.756521
  * [BHS14] Chiara Bianchini, Antoine Henrot and Paolo Salani  “An overdetermined problem with non-constant boundary condition”  In _Interfaces Free Bound._ 16.2, 2014, pp. 215–241  DOI: 10.4171/IFB/318
  * [BK11] Giuseppe Buttazzo and Bernd Kawohl  “Overdetermined boundary value problems for the $\infty$-Laplacian”  In _Int. Math. Res. Not. IMRN_ , 2011, pp. 237–247  DOI: 10.1093/imrn/rnq071
  * [Bra+08] B. Brandolini, C. Nitsch, P. Salani and C. Trombetti  “Serrin-type overdetermined problems: an alternative proof”  In _Arch. Ration. Mech. Anal._ 190.2, 2008, pp. 267–280  DOI: 10.1007/s00205-008-0119-3
  * [Cab+18] Xavier Cabré, Mouhamed Moustapha Fall, Joan Solà-Morales and Tobias Weth  “Curves and surfaces with constant nonlocal mean curvature: meeting Alexandrov and Delaunay”  In _J. Reine Angew. Math._ 745, 2018, pp. 253–280  DOI: 10.1515/crelle-2015-0117
  * [Cir+18] Giulio Ciraolo, Alessio Figalli, Francesco Maggi and Matteo Novaga  “Rigidity and sharp stability estimates for hypersurfaces with constant and almost-constant nonlocal mean curvature”  In _J. Reine Angew. Math._ 741, 2018, pp. 275–294  DOI: 10.1515/crelle-2015-0088
  * [Cir+23] Giulio Ciraolo, Serena Dipierro, Giorgio Poggesi, Luigi Pollastro and Enrico Valdinoci  “Symmetry and quantitative stability for the parallel surface fractional torsion problem”  In _Trans. Amer. Math. Soc._ 376.5, 2023, pp. 3515–3540  DOI: 10.1090/tran/8837
  * [CMS15] Giulio Ciraolo, Rolando Magnanini and Shigeru Sakaguchi  “Symmetry of minimizers with a level surface parallel to the boundary”  In _J. Eur. Math. Soc. (JEMS)_ 17.11, 2015, pp. 2789–2804  DOI: 10.4171/JEMS/571
  * [CMS16] Giulio Ciraolo, Rolando Magnanini and Shigeru Sakaguchi  “Solutions of elliptic equations with a level surface parallel to the boundary: stability of the radial configuration”  In _J. Anal. Math._ 128, 2016, pp. 337–353  DOI: 10.1007/s11854-016-0011-2
  * [CS09] Andrea Cianchi and Paolo Salani  “Overdetermined anisotropic elliptic problems”  In _Math. Ann._ 345 Springer, 2009, pp. 859–881
  * [CV19] Giulio Ciraolo and Luigi Vezzoni  “On Serrin’s overdetermined problem in space forms”  In _Manuscripta Math._ 159.3-4, 2019, pp. 445–452  DOI: 10.1007/s00229-018-1079-z
  * [Dip+22] Serena Dipierro, Giorgio Poggesi, Jack Thompson and Enrico Valdinoci  “The role of antisymmetric functions in nonlocal equations”  In _Trans. Amer. Math. Soc. (to appear)_ , 2022  DOI: 10.48550/ARXIV.2203.11468
  * [Dip+23] Serena Dipierro, Mateusz Kwaśnicki, Jack Thompson and Enrico Valdinoci  “The nonlocal Harnack inequality for antisymmetric functions: an approach via Bochner’s relation and harmonic analysis”  In _In preparation_ , 2023
  * [DPV21] Serena Dipierro, Giorgio Poggesi and Enrico Valdinoci  “A Serrin-type problem with partial knowledge of the domain”  In _Nonlinear Anal._ 208, 2021, pp. Paper No. 11233044  DOI: 10.1016/j.na.2021.112330
  * [DPV22] Serena Dipierro, Giorgio Poggesi and Enrico Valdinoci  “Radial symmetry of solutions to anisotropic and weighted diffusion equations with discontinuous nonlinearities”  In _Calc. Var. Partial Differential Equations_ 61.2, 2022, pp. Paper No. 7231  DOI: 10.1007/s00526-021-02157-5
  * [DS07] Francesca Da Lio and Boyan Sirakov  “Symmetry results for viscosity solutions of fully nonlinear uniformly elliptic equations”  In _J. Eur. Math. Soc. (JEMS)_ 9.2, 2007, pp. 317–330  DOI: 10.4171/JEMS/81
  * [FG08] Ilaria Fragalà and Filippo Gazzola  “Partially overdetermined elliptic boundary value problems”  In _J. Differential Equations_ 245.5, 2008, pp. 1299–1322  DOI: 10.1016/j.jde.2008.06.014
  * [FJ15] Mouhamed Moustapha Fall and Sven Jarohs  “Overdetermined problems with fractional Laplacian”  In _ESAIM Control Optim. Calc. Var._ 21.4, 2015, pp. 924–938  DOI: 10.1051/cocv/2014048
  * [FK08] A. Farina and B. Kawohl  “Remarks on an overdetermined boundary value problem”  In _Calc. Var. Partial Differential Equations_ 31.3, 2008, pp. 351–357  DOI: 10.1007/s00526-007-0115-8
  * [GL89] Nicola Garofalo and John L. Lewis  “A symmetry result related to some overdetermined boundary value problems”  In _Amer. J. Math._ 111.1, 1989, pp. 9–33  DOI: 10.2307/2374477
  * [GNN79] B. Gidas, Wei Ming Ni and L. Nirenberg  “Symmetry and related properties via the maximum principle”  In _Comm. Math. Phys._ 68.3, 1979, pp. 209–243  URL: http://projecteuclid.org/euclid.cmp/1103905359
  * [GNN81] B. Gidas, Wei Ming Ni and L. Nirenberg  “Symmetry of positive solutions of nonlinear elliptic equations in ${\bf R}^{n}$”  In _Mathematical analysis and applications, Part A_ 7, Adv. in Math. Suppl. Stud.  Academic Press, New York-London, 1981, pp. 369–402
  * [GT01] David Gilbarg and Neil S. Trudinger  “Elliptic partial differential equations of second order” Reprint of the 1998 edition, Classics in Mathematics  Springer-Verlag, Berlin, 2001, pp. xiv+517
  * [HP18] Antoine Henrot and Michel Pierre  “Shape variation and optimization” A geometrical analysis, English version of the French publication [ MR2512810] with additions and updates 28, EMS Tracts in Mathematics  European Mathematical Society (EMS), Zürich, 2018, pp. xi+365  DOI: 10.4171/178
  * [HP98] Antoine Henrot and Gérard A. Philippin  “Some overdetermined boundary value problems with elliptical free boundaries”  In _SIAM J. Math. Anal._ 29.2, 1998, pp. 309–320  DOI: 10.1137/S0036141096307217
  * [KP94] S. Kesavan and Filomena Pacella  “Symmetry of positive solutions of a quasilinear elliptic equation via isoperimetric inequalities”  In _Appl. Anal._ 54.1-2, 1994, pp. 27–37  DOI: 10.1080/00036819408840266
  * [Lio81] P.-L. Lions  “Two geometrical properties of solutions of semilinear problems”  In _Applicable Anal._ 12.4, 1981, pp. 267–272  DOI: 10.1080/00036818108839367
  * [MP20] Rolando Magnanini and Giorgio Poggesi  “Nearly optimal stability for Serrin’s problem and the soap bubble theorem”  In _Calc. Var. Partial Differential Equations_ 59.1, 2020, pp. Paper No. 3523  DOI: 10.1007/s00526-019-1689-7
  * [MP20a] Rolando Magnanini and Giorgio Poggesi  “Serrin’s problem and Alexandrov’s soap bubble theorem: enhanced stability via integral identities”  In _Indiana Univ. Math. J._ 69.4, 2020, pp. 1181–1205  DOI: 10.1512/iumj.2020.69.7925
  * [MS10] Rolando Magnanini and Shigeru Sakaguchi  “Nonlinear diffusion with a bounded stationary level surface”  In _Ann. Inst. H. Poincaré Anal. Non Linéaire_ 27.3, 2010, pp. 937–952  DOI: 10.1016/j.anihpc.2009.12.001
  * [PS89] LE Payne and Philip W Schaefer  “Duality theorems in some overdetermined boundary value problems”  In _Math. Meth. Appl. Sciences_ 11.6 Wiley Online Library, 1989, pp. 805–819
  * [QX17] Guohuan Qiu and Chao Xia  “Overdetermined boundary value problems in $\mathbb{S}^{n}$”  In _J. Math. Study_ 50.2, 2017, pp. 165–173  DOI: 10.4208/jms.v50n2.17.03
  * [RS14] Xavier Ros-Oton and Joaquim Serra  “The Dirichlet problem for the fractional Laplacian: regularity up to the boundary”  In _J. Math. Pures Appl. (9)_ 101.3, 2014, pp. 275–302  DOI: 10.1016/j.matpur.2013.06.003
  * [RS19] Xavier Ros-Oton and Joaquim Serra  “The boundary Harnack principle for nonlocal elliptic operators in non-divergence form”  In _Potential Anal._ 51.3, 2019, pp. 315–331  DOI: 10.1007/s11118-018-9713-7
  * [Ser13] Joaquim Serra  “Radial symmetry of solutions to diffusion equations with discontinuous nonlinearities”  In _J. Differential Equations_ 254.4, 2013, pp. 1893–1902  DOI: 10.1016/j.jde.2012.11.015
  * [Ser71] James Serrin  “A symmetry problem in potential theory”  In _Arch. Rational Mech. Anal._ 43, 1971, pp. 304–318  DOI: 10.1007/BF00250468
  * [Sha12] Henrik Shahgholian  “Diversifications of Serrin’s and related symmetry problems”  In _Complex Var. Elliptic Equ._ 57.6, 2012, pp. 653–665  DOI: 10.1080/17476933.2010.504848
  * [SS15] Luis Silvestre and Boyan Sirakov  “Overdetermined problems for fully nonlinear elliptic equations”  In _Calc. Var. Partial Differential Equations_ 54.1, 2015, pp. 989–1007  DOI: 10.1007/s00526-014-0814-x
  * [SV19] Nicola Soave and Enrico Valdinoci  “Overdetermined problems for the fractional Laplacian in exterior and annular sets”  In _J. Anal. Math._ 137.1, 2019, pp. 101–134  DOI: 10.1007/s11854-018-0067-2
  * [Tam76] Italo Tamanini  “Il problema della capillarità su domini non regolari”  In _Rend. Sem. Mat. Univ. Padova_ 56, 1976, pp. 169–191 (1978)  URL: http://www.numdam.org/item?id=RSMUP_1976__56__169_0
  * [Wei71] H.. Weinberger  “Remark on the preceding paper of Serrin”  In _Arch. Rational Mech. Anal._ 43, 1971, pp. 319–320  DOI: 10.1007/BF00250469


of $1024$ points. Unless otherwise specified, UMBC models use the softmax
activation function. We list the architectures in Sections C.2, C.2 and C.2.
All models have an additional linear output which outputs $K\times 5$
parameters for the Gaussian mixture outlined in Equation 14.

Table 6: MVN generic model (Used by all encoders)

Output Size | Layers | Amount
---|---|---
$N_{i}\times 2$ | Input Set | $\times 1$
$N_{i}\times 128$ | Linear(2, 128), ReLU | $\times 1$
$N_{i}\times 128$ | Set Encoder | $\times 1$
$K\times 5$ | Decoder | $\times 3$

Table 7: Set Encoder Specific settings for baseline models.

Name | Set Encoder | Output Size | Set Decoder | Output Size
---|---|---|---|---
DeepSets (Zaheer et al., 2017) | Mean Pooling | $128$ | Linear, ReLU | $128$
SSE (Bruno et al., 2021) | Slot Set Encoder | $K\times 128$ | Linear, ReLU | $128$
FSPool (Zhang et al., 2020) | Featurewise Sort Pooling | $128$ | Linear, ReLU | $128$
Diff. EM. (Kim, 2022) | Expectation Maximization Layer | $1286$ | Linear, ReLU | 128
Set Transformer (Lee et al., 2019) | Pooling by Multihead Attention | $K\times 128$ | Set Attention Block | $K\times 128$

Table 8: Set Encoder Specific settings for UMBC models. UMBC models account
for the extra encoder by using fewer layers in the decoder (2 layers instead
of 3).

Name | MBC Set Encoder | Output Size | non-MBC Set Encoder | Output Size | Set Decoder | Output Size
---|---|---|---|---|---|---
(Ours) UMBC+FSPool | UMBC Layer | $K\times 128$ | Featurewise Sort Pooling | $128$ | Linear, ReLU | $128$
(Ours) UMBC+Diff EM | UMBC Layer | $K\times 128$ | Expectation Maximization Layer | $1286$ | Linear, ReLU | $128$
(Ours) UMBC+Set Transformer | UMBC Layer | $K\times 128$ | Set Attention Block | $K\times 128$ | Linear, ReLU | $K\times 128$

##  Appendix D Measuring the Variance of Pooled Features

In Figure 3, we show the quantitative effect on the pooled representation
between the plain Set Transformer, UMBC+Set Transformer, FSPool and DiffEM.
The UMBC model always shows 0 variance, while the non-MBC models produce
variance between aggregated encodings of random partitions. For a single
chunk, however, non-MBC models show no variance, as random partitions of a
single chunk would be equivalent to permuting the elements within the chunk
(_i.e_. non-MBC models still produce an encoding which is permutation
invariant). The variance increases drastically when the set is partitioned
into two chunks and then the behavior differs between the non-MBC models. Set
Transformer happens to show decreasing variance as the number of chunks
increases. Note that as the number of chunks increases, the cardinality of
each chunk decreases. Therefore, the variance decreases as the chunk
cardinality also decreases, but this _does not_ indicate that the models is
performing better. For example, in Figure 2, when a singleton set is input to
the Set Transformer, the predictions become almost meaningless even though
they may have lower variance. The procedure for aggregating the encodings of
set partitions for the non-MBC models is outlined in Appendix E.

Figure 9: Distributions used in sampling random inputs for the encoding variance experiment in Figure 2 Distribution | Dimension | Number of Points
---|---|---
Normal(0, 1) | $128$ | $256$
Uniform(-3, 3) | $128$ | $256$
Exponential(1) | $128$ | $256$
Cauchy(0, 1) | $128$ | $256$

Figure 10: The number of chunks and elements per chunk.

Number of Chunks | $1$ | $2$ | $4$ | $8$ | $16$ | $32$
---|---|---|---|---|---|---
Elements per Chunk | $1024$ | $512$ | $256$ | $128$ | $64$ | $32$

To perform this experiment, we used a randomly initialized model with $128$
hidden units, and sampled $256$ set elements from four different distributions
in order to make a total set size of $1024$. We then created 100 random
partitions for various chunk sizes. Chunk sizes and distributions are shown in
Figure 10. We then encode the whole set in chunks and and report the observed
variance over the 100 different random partitions at each of the various chunk
sizes (Figure 3). Note that the encoded set representation is a vector and
Figure 3 shows a scalar value. To achieve this, we take the feature-wise
variance over the 100 encodings and report the mean and standard deviation
over the feature dimension. Specifically, given
$\mathbf{Z}=[\mathbf{z}_{1}\cdots\mathbf{z}_{100}]^{\top}\in\mathbb{R}^{100\times
128}$ representing all 100 encodings with
$\mathbf{z}_{i}=(z_{i,1},\ldots,z_{i,128})$, we compute feature-wise variance
as

$\displaystyle\hat{z}_{j}=\frac{1}{(100-1)}\sum_{i=1}^{100}(z_{i,j}-\mu_{j})^{2},\quad\mu_{j}=\frac{1}{100}\sum_{i=1}^{100}z_{i,j}$

for $j=1,\ldots,128$. We then achieve the values of y-axis and error bars in
Figure 3 by a mean and standard deviation over the feature dimension,

$\displaystyle y=\frac{1}{128}\sum_{i}\hat{z}_{i},\quad\quad
y_{\sigma}=\sqrt{\frac{1}{(128-1)}\sum_{i}(\hat{z}_{i}-y)^{2}}.$ (51)

##  Appendix E A Note on MBC Testing of non-MBC models

In the qualitative experiments Figures 2 and 3, we apply MBC testing to non-
MBC models in order to study the effects of using non-MBC models in MBC
settings. Non-MBC models do not prescribe a way to accomplish this in the
original works, so we took the approach of processing each chunk up until the
pooled representation. We then performed mean pooling over the encoded chunks
in the following way. Let $X$ be an input set and let
$\zeta(X)=\\{X_{1},\ldots,X_{n}\\}$ be a partition of the set $X$, _i.e_.
$X=\bigcup_{j=1}^{N}X_{j}$ with $X_{i}\cap X_{j}=\emptyset$ for $i\neq j$.
Denote $\tilde{f}_{\theta}$ a non-MBC set encoding function, then our pseudo-
MBC testing procedure is as follows,

$\displaystyle Z=\frac{1}{N}\sum_{j=1}^{N}\tilde{f}_{\theta}(X_{j})$ (52)

(a) Deepsets

(b) Slot Set Encoder

(c) UMBC + ST

Figure 11: Performance of (a) DeepSets (b) SSE, and (c) UMBC with varying set
sizes for image completion on CelebA dataset.

##  Appendix F Details on the Image Completion Experiments

###  F.1 Additional Experimental Results

In figure Figure 11, we evaluate our proposed unbiased full set gradient
approximation algorithm (red) with Deepsets, Slot Set Encoder (SSE) and UMBC +
Set Transformer (ST) and compare our algorithm against the one training with a
randomly sampled subset of 100 elements , which is a biased estimator, (green)
and the one computing full set gradient (blue). Across all models, our
unbiased estimator significantly outperforms the models trained with a
randomly sampled subset. Notably, the model trained with our proposed
algorithm is indistinguishable from the model trained with full set gradient
while our method only incurs constant memory overhead for any set size. These
empirical results again verify efficiency of our unbiased full set gradient
approximation.

###  F.2 Experimental Setup

We train all models on CelebA dataset for 200,000 steps with Adam optimizer
(Kingma & Ba, 2015) and 256 batch size but no weight decay. We set the
learning rate to $5\cdot 10^{-4}$ and use a cosine annealing learning rate
schedule. In Sections F.2 and F.2, we specify the architecture of Conditional
Neural Process with UMBC + Set Transformer. We use $k=128$ slots and set
dimension of each slot to $d_{s}=128$. For the attention layer, we use the
softmax for the activation function $\sigma$ and set the dimension of
attention output to $d=128$. As an input to the set encoder, we concatenate
the coordinates of each ${\mathbf{x}}_{i,c_{j}}\in\mathbb{R}^{2}$ and the
corresponding pixel value $y_{i,c_{j}}\in\mathbb{R}^{3}$ from the context
$X_{i}$ for $j=1,\ldots,N_{i}$, resulting in a $N_{i}\times 5$ matrix.

Table 9: UMBC Set Encoder of Conditional Neural Process.

Output Size | Layers
---|---
$N_{i}\times 5$ | Input Context Set
$N_{i}\times 128$ | Linear(5, 128), ReLU
$N_{i}\times 128$ | Linear(128, 128), ReLU
$N_{i}\times 128$ | Linear(128, 128), ReLU
$N_{i}\times 128$ | Linear(128, 128), ReLU
$128\times 128$ | UMBC Layer
$128\times 128$ | Layer Normalization
$128\times 128$ |  Set Attention Block (Lee et al., 2019)
$128\times 128$ | Set Attention Block
$128$ |  Pooling by Multihead Attention (Lee et al., 2019)

Table 10: Decoder of Conditional Neural Process.

Output Size | Layers
---|---
$128,1024\times 2$ | Input Set Representation and Coordinates
$1024\times 130$ | Tile & Concatenate
$1024\times 128$ | Linear(130, 128), ReLU
$1024\times 128$ | Linear(128, 128), ReLU
$1024\times 128$ | Linear(128, 128), ReLU
$1024\times 128$ | Linear(128, 128), ReLU
$1024\times 6$ | Linear(128, 6)

##  Appendix G Details on the Long Document Classification Experiments

We train all models for 30 epochs with AdamW optimizer (Loshchilov & Hutter,
2019) and batch size 8. We use constant learning rate $5\cdot 10^{-5}$. For
our UMBC model, we pretrain the model while freezing BERT for 30 epochs and
finetune the whole model for another 30 epochs. In Table 11, we specify
architecture of UMBC + BERT (Devlin et al., 2019) without positional encoding.
We use $k=256$ slots and set dimension of each slot to $d_{s}=128$. We use
slot-sigmoid for the activation function $\sigma$ and set the dimension of the
attention output to $d=768$.

Table 11: UMBC Set Encoder and BERT decoder for long document classification.

Output Size | Layers
---|---
$N_{i}$ | Input Document
$N_{i}\times 768$ | Word Embedding
$N_{i}\times 768$ | Layer Normalization
$N_{i}\times 768$ | Linear(768,768), ReLU
$N_{i}\times 768$ | Linear(768,768), ReLU
$256\times 768$ | UMBC Layer
$256\times 768$ | Layer Normalization
$256\times 768$ | BERT w/o Positional Encoding (Devlin et al., 2019)
$768$ | [CLS] token Pooler
$4271$ | Dropout(0.1), Linear(768, 4271)

##  Appendix H Details on the Camelyon16 Experiments

Figure 12: Examples of $4$ single patches from the Camelyon16 dataset. On
average, each set in the dataset contains over $9,300$ patches like the ones
pictured above.

The Camelyon16 Whole Slide Image dataset consists of 270 training instances
and 129 validation instances. The dataset was created for a competition, and
therefore the test set is hidden. We therefore follow the example set by
previous works (Li et al., 2021) and report performance achieved on the
validation set. For preprocessing, we consider the $20\times$ slide
magnification setting, and use OTSU’s thresholding method to detect regions
containing tissue within the WSI. We then split the activated regions into non
overlapping patches of size $256\times 256$. An example of single input
patches can be seen in Figure 12. The largest input set contains $37,345$
image patches which are each $\in\mathbb{R}^{256\times 256\times 3}$. All
patch extraction code can be found in the supplementary file. Table 14
contains statistics related to the numbers of patches per input for the
training and the test set as well as the distribution of positive and negative
labels.

###  H.1 Experimental Setup

We use a ResNet18 (He et al., 2016) which was pretrained with self-supervised
contrastive learning (Chen et al., 2020) by Li et al. (2021). The pretrained
ResNet18 weights can be downloaded from this repository. Following the
classification experiments done by Lee et al. (2019), we place dropout layers
before and after the PMA layer of the Set Transformer in our UMBC model. We
will describe our pretraining and finetuning steps below in detail.

Table 12: Camelyon16 generic model (Used by all encoders)

Output Size | Layers | Name | Amount
---|---|---|---
$N_{i}\times 256\times 256\times 3$ | Input Set | Bag of Instances | $\times 1$
$N_{i}\times 512$ | ResNet18(InstanceNorm) | Feature Extractor | $\times 1$
$N_{i}\times 256$ | Linear, ReLU, Linear | Projection | $\times 1$
$1$ | Linear, Max Pooling | Instance Classifier | $\times 1$
$N_{i}\times 128$ | Set Encoding Function | Bag Classifier | $\times 1$
$1$ | Set Decoder | Bag Classifier | $\times 1$

Table 13: Camelyon16 Bag Classifier Models.

Name | Set Encoder | Output Size
---|---|---
AB-MIL |  Gated Attention (Ilse et al., 2018) | $1$
DSMIL |  DS-MIL Aggregator (Li et al., 2021) | $1$
Deepsets |  Max Pooling $\rightarrow$ Linear $\rightarrow$ ReLU $\rightarrow$ Linear | $1$
Slot Set Encoder |  SSE (Bruno et al., 2021)$\rightarrow$ Linear $\rightarrow$ ReLU $\rightarrow$ Linear | $1$
UMBC+SetTransformer |  UMBC(K=$64$) $\rightarrow$ SAB $\rightarrow$ PMA(K=$1$, p=0.5) $\rightarrow$ Linear | $1$

#### Pretraining.

For pretraining, we extract the features from the pretrained ResNet18 and only
train the respective MIL models (Section H.1) on the extracted features. We
pretrain for 200 epochs with the Adam optimizer which uses a learning rate of
$5\cdot 10^{-4}$ and a cosine annealing learning rate decay which reaches the
minimum at $5\cdot 10^{-6}$. We use $\beta_{1}=0.5$, and $\beta_{2}=0.9$ for
Adam. We train with a batch size of 1 on a single GPU, and save the model
which showed the best performance on the validation set, where the performance
metric is $\left(\text{Accuracy}+\text{AUC}\right)/2$. Other details can be
found in Section 4.4. These results can be seen in the left column of Table 4.

#### Finetuning.

For finetuning, we use our unbiased gradient approximation algorithm with a
chunk size of $256$. We freeze the pretrained MIL head and only finetune the
backbone resnet model. Therefore, we sequentially process each $256$ chunk for
each input set until the entire set has been processed. We train for 10 total
epochs, and use the AdamW optimizer with a learning rate of $5\cdot 10^{-5}$,
and a weight decay of $1\cdot 10^{-2}$ which is not applied to bias or
layernorm parameters. We use a one epoch linear warmup, and then a cosine
annealing learning rate decay at every iteration which reaches a minimum at
$5\cdot 10^{-6}$. We train on 1 GPU, with a batch size of 1 and with a single
instance on each GPU.

Table 14: Statistics for the Camelyon16 training and test sets we used. Left:
Number of patches (set size) per instance. Right: The distribution of positive
and negative samples.

Metric | Train | Test
---|---|---
Mean | 9,329 | 9,376
Min | 154 | 1558
Max | 32,382 | 37,345

Metric | Train | Test
---|---|---
Positive ($+$) | 110 | 49
Negative ($-$) | 160 | 80

##  Appendix I Generalizing Attention Activations

As shown in Equation 7, any attention activation function which can be
expressed as a strictly positive elementwise function combined with sum
decomposable normalization constants $\nu_{p}$ and ${\bar{f}}_{\theta}$
represents a valid attention activation function. Table 15 shows 5 such
functions with their respective normalization constants, although there are an
infinite number of possible functions which can be used.

The softmax operation we propose $h_{1}:\mathbb{R}^{d}\mapsto(0,1)^{d}$ which
is outlined immediately before Theorem 3.4 is mathematically equivalent to the
standard softmax $h_{2}:\mathbb{R}^{d}\mapsto(0,1)^{d}$ which is commonly
implemented in deep learning libraries because $f$ and $g$ have the same
domain, the same codomain, and $h_{1}(\mathbf{x})=h_{2}(\mathbf{x})$ for all
$\mathbf{x}\in\mathbb{R}^{d}$. Therefore the functions are mathematically
equivalent, even though the implementations are not. Our proposed function
$h_{1}$ requires separately applying the exponential, and storing and updating
the normalization constant while $h_{2}$ is generally implemented in such a
way that everything is done in a single operation.

Table 15: Valid UMBC attention activation functions with slot normalization
$\nu_{1}$ and normalization over the set elements ${\bar{f}}_{\theta}$. Slot-
exp uses $\hat{A}_{i,j}-\max_{1\leq i\leq k}\hat{A}_{i,j}$ instead of
$\nu_{p}$.

function $\sigma$ | $\nu_{p}$ | ${\bar{f}}_{\theta}$ | name | reference
---|---|---|---|---
$\mathtt{sigmoid}$ | $p=1$ | - | slot-sigmoid | (Bruno et al., 2021)
$\exp()$ | $p=1$ | ✓ | slot-softmax | (Locatello et al., 2020)
$\exp()$ | $p=2$ | ✓ | softmax | (Lee et al., 2019)
$\exp()$ | -22footnotemark: 2 | ✓ | slot-exp | -
$\mathtt{sigmoid}()$ | $p=2$ | ✓ | sigmoid | -

##  Appendix J Algorithm

We outline our unbiased full set gradient approximation here.

Algorithm 1 Unbiased Full Set Gradient Estimation

1: Input: Dataset $((X_{i},y_{i}))_{i=1}^{n}$, batch size $m$, the number of
subsets $m^{\prime}$, learning rate $(\eta_{t})_{t=1}^{T}$, total steps $T$,
and functions $f_{\theta}$, and $g_{\lambda}$.

2: Randomly initialize $\theta$ and $\lambda$

3: for all $t=1,\ldots,T$ do

4: Sample $((\bar{X}_{i},{\bar{y}}_{i}))_{i=1}^{m}\sim
D[((X_{i},y_{i}))_{i=1}^{n}]$

5: $L_{t,1}(\theta,\lambda)\leftarrow 0,L_{t,2}(\theta,\lambda)\leftarrow 0$

6: for all $i=1\ldots,m$ do

7: Partition a set $\bar{X}_{i}$ to get $\zeta_{t}(\bar{X}_{i})$

8: Sample ${\bar{\zeta}}_{t}(\bar{X}_{i})\sim D[\zeta_{t}(\bar{X}_{i})]$ with
$|{\bar{\zeta}}_{t}(\bar{X}_{i})|=m^{\prime}$

9:
$f_{\theta}(\bar{X}_{i})=\sum_{{\bar{S}}\in{\bar{\zeta}}_{t}(\bar{X}_{i})}f_{\theta}({\bar{S}})+\sum_{S\in\zeta_{t}(\bar{X}_{i})\backslash{\bar{\zeta}}_{t}(\bar{X}_{i})}\texttt{StopGrad}(f_{\theta}(S))$

10: $L_{t,1}(\theta,\lambda)\leftarrow
L_{t,1}(\theta,\lambda)+\frac{1}{m}\frac{|\zeta_{t}(\bar{X}_{i})|}{|{\bar{\zeta}}_{t}(\bar{X}_{i})|}\ell(g_{\lambda}(f_{\theta}(\bar{X}_{i})),{\bar{y}}_{i})$

11: $L_{t,2}(\theta,\lambda)\leftarrow
L_{t,2}(\theta,\lambda)+\frac{1}{m}\frac{1}{\lvert{\bar{\zeta}}_{t}(\bar{X}_{i})\rvert}\ell(g_{\lambda}(f_{\theta}(\bar{X}_{i})),{\bar{y}}_{i})$

12: end for

13: $\theta\leftarrow\theta-\eta_{t}\frac{\partial
L_{t,1}(\theta,\lambda)}{\partial\theta}$

14: $\lambda\leftarrow\lambda-\eta_{t}\frac{\partial
L_{t,2}(\theta_{t},\lambda)}{\partial\lambda}$

15: end for