\section{Introduction and Problem Formulation}
	\input{./sections/introduction}

	\section{Results - An Overview}
	\subsection{Binary $\rv{Y}$}
	\input{./sections/problem_definition}

	\input{sections/gaussian}

	\subsection{Modulo Additive Channels with \acrshort{pf} Constraint}
	\label{section:comib_modulo_additive_pf}
	\input{./sections/modulo_channels}

	\section{Modulo Additive Channels with \acrshort{tv} Constraint}
	\label{section:tv_constraint}
	\input{./sections/total_variation}

	\section{An Alternating Optimization Algorithm} \label{section:alternating_algorithm}
	\input{./sections/alternating_algorithm}

	\section{Alternating Optimization Algorithm for Modulo Additive Channels}
	\label{section:alternating_algorithm_modulo}
	\input{./sections/algorithm/modulo}

 	\section{Numerical Simulations}
 	\input{./sections/numerical_simulations}

\section{Selected Proofs}

	\subsection{Proof of \autoref{proposition:comib_binaryXY}}
	\label{appendix:proof_of_proposition_comib_binaryXY}
	\input{./sections/binary/proof_binaryXY}

	\subsection{Proof of \autoref{theorem:oblib_gaussianXY}} \label{section:oblib_proof_of_theorem_gaussianXY}
	\input{./sections/proof_of_theorem_gaussianXY}

 	\subsection{Privacy Funnel for Jointly Gaussian Vectors}
 	\label{appendix:pf_vector_gaussian}
 	\input{./sections/gaussian/pf_VectorGaussian}

	\subsection{Proof of \autoref{theorem:oblib_VectorGaussianXY}} \label{section:oblib_proof_of_theorem_VectorgaussianXY}
	\input{./sections/proof_of_theorem_VectorGaussianXY}

	\subsection{Proof of \autoref{theorem:oblib_main_result}}
	\label{section:oblib_proof_of_theorem_low_snr}
	\input{./sections/optimization_problem}

	\subsection{Proof of \autoref{theorem:oblib_main_result_bounds}} \label{section:oblib_proof_of_main_result_bounds}
	\input{./sections/negative_hamming}

	\subsection{Proof of \autoref{theorem:comib_modulo_high_snr}}
	\label{section:oblib_proof_of_high_modulo}
	\input{./sections/high_snr_modulo}

	\subsection{Proof of \autoref{theorem:comib_total_variation_bounds}}
	\label{appendix:oblib_proof_tv_bounds}
	\input{./sections/tv/proof_modulo_tv_bounds}

\section{Concluding Remarks}
	\input{./sections/concluding_remarks}

	\section*{Acknowledgment}
	The work has been supported by the
	European Union's Horizon 2020 Research And Innovation Programme,
	grant agreement no. 694630, by the ISF under Grant 1791/17, and by the WIN consortium via the
	Israel minister of economy and science.

\subsection{\acrlong{tito} Circulant Matrix}
\input{../chapters/oblib/tito}

\subsection{Stationary Points Characterization}

\begin{proposition} \label{proposition:oblib_stationarity_inner_minimization}
	Any stationary point $ \pmf{\rv{X}\rv{Y}}^* $ of \eqref{eq:definition_oblivious_ib_minimization_lagrangian} satisfies
	\begin{equation} \label{eq:oblib_stationarity_pXY}
		\mathsf{P}^*_{\rv{X}\rv{Y}} (x,y) =  \frac{\pmf{\rv{X}}(x) \pmf{\rv{Y}}(y) e^{- \beta_1   D\left( \pmf{\rv{Z}|\rv{Y}} (\cdot|y) || \pmf{\rv{Z}|\rv{X}} (\cdot|x)   \right)}}{Z_1 (x,y,\beta_1)},
	\end{equation}
	where $ \beta_1 \triangleq 1/ \lambda_1 $ and $ Z_1(x,y,\beta_1) $ is the proper marginalization function, which verifies that $\pmf{\rv{X}\rv{Y}}$ has the desired marginals $\pmf{\rv{X}}$ and $\pmf{\rv{Y}}$.  Furthermore,  the optimal $ \pmf{\rv{Z}|\rv{X}} (z|x) $ is given by
	\begin{equation} \label{eq:oblib_stationarity_pZgX}
		\pmf{\rv{Z}|\rv{X}} (z|x) = \frac{1}{\pmf{\rv{X}}(x)} \sum_{y \in \set{Y}} \pmf{\rv{Z}|\rv{Y}} (z|y) \mathsf{P}^*_{\rv{X} \rv{Y}} (x,y).
	\end{equation}
\end{proposition}

\begin{remark}
Note that the problem of computing $\rv{Z}_1(x,y,\beta_1)$ is of independent interest. We propose an alternating algorithm which is summarized in \autoref{algorithm:oblib_minimization_marginalization}.
\end{remark}

\begin{algorithm}[t]
	\SetAlgoLined
	\SetKwInput{Init}{Initialize}
	\KwIn{$ \pmf{\rv{X}} $, $ \pmf{\rv{Y}} $, $ \mathsf{Q}_{\rv{X}\rv{Y}} $ and $NumIter$}
	\Init{$ \mathsf{Q}_{\rv{X}\rv{Y}}^{(0)} = \mathsf{Q}_{\rv{X}\rv{Y}}  $.}
	\For{ $ t= 1$ \KwTo $IterNum$}{
		$\mathsf{Q}_{\rv{X}} (x) = \sum_{y \in \set{Y}} \mathsf{Q}_{\rv{X}\rv{Y}}^{(t-1)} (x,y)$ \;
		$\mathsf{Q}_{\rv{Y}} (y) = \sum_{x \in \set{X}} \mathsf{Q}_{\rv{X}\rv{Y}}^{(t-1)} (x,y)$ \;
		$\mathsf{Q}_{\rv{X}\rv{Y}}^{un} (x,y) = \frac{ \pmf{\rv{X}}(x) \cdot \pmf{\rv{Y}}(y) \mathsf{Q}_{\rv{X}\rv{Y}}^{(t-1)} (x,y)}{\mathsf{Q}_{\rv{X}} (x)  \cdot \mathsf{Q}_{\rv{Y}} (y) }$ \;
		$\mathsf{Q}_{\rv{X}\rv{Y}}^{(t)} (x,y) = \frac{\mathsf{Q}_{\rv{X}\rv{Y}}^{un} (x,y)}{\sum_{x,y} \mathsf{Q}_{\rv{X}\rv{Y}}^{un} (x,y)}$
	}
	\KwOut{$\mathsf{Q}_{\rv{X}\rv{Y}}^{(IterNum)} (x,y) $}
	\caption{marginalization(args)}
	\label{algorithm:oblib_minimization_marginalization}
\end{algorithm}

The system of equations characterizing the stationary points in \eqref{eq:oblib_stationarity_pXY} and \eqref{eq:oblib_stationarity_pZgX} must hold simultaneously for consistency. An alternating iteration algorithm is a common approach to solve these equations.

\begin{proposition}
	Equations \eqref{eq:oblib_stationarity_pXY} and \eqref{eq:oblib_stationarity_pZgX} are satisfied simultaneously at the minimum of the Lagrangian \eqref{eq:definition_oblivious_ib_minimization_lagrangian}
	where the minimization is performed independently over the convex sets of $ \{\pmf{\rv{X}\rv{Y}} (x,y)\} $ and $ \{\pmf{\rv{Z}|\rv{X}} (z|x) \} $,
	\begin{equation}
		\min_{\pmf{\rv{Z}|\rv{X}} (z|x) } \min_{\pmf{\rv{X}\rv{Y}} (x,y)} \mathcal{L}_{\min} (\pmf{\rv{X}\rv{Y}} , \lambda_1, \mu,\nu) .
	\end{equation}
	These independent conditions correspond precisely to alternating interactions of \eqref{eq:oblib_stationarity_pXY} and \eqref{eq:oblib_stationarity_pZgX}. Denoting by $ t $ the iteration step, we obtain \autoref{algorithm:oblib_minimization}.
\end{proposition}

\begin{algorithm}[b]
	\SetAlgoLined
	\SetKwInput{Init}{Initialize}
	\KwIn{$ \pmf{\rv{X}} $, $ \pmf{\rv{Y}} $, $ \pmf{\rv{Z}|\rv{Y}} $ and $ \beta_1 $}
	\Init{Arbitrary $ \pmf{\rv{X}\rv{Y}}^{(0)}  $ with valid marginals, $t=1$.}
	\While{Variation in $ I(\rv{X};\rv{Z}) $ is greater then $ \epsilon $}{
		Compute $ \pmf{\rv{Z}|\rv{X}}^{(t)} (z|x) = \frac{\sum_{y \in \set{Y}} \pmf{\rv{Z}|\rv{Y}} (z|y) \pmf{\rv{X} \rv{Y}} ^{(t-1)} (x,y)}{\pmf{\rv{X}}(x)}  $ \;
		Set $ \pmf{\rv{X}\rv{Y}}^{(t)} (x,y) =  \frac{ \pmf{\rv{X}} (x) \pmf{\rv{Y}} (y) e^{- \beta_1    D\left( \pmf{\rv{Z}|\rv{Y}} (\cdot|y) || \pmf{\rv{Z}|\rv{X}}^{(t)} (\cdot|x) \right) }}{Z_1(x,y,\beta_1)} $\;
		Find $ Z_1(x,y,\beta_1) $ s.t. $  \pmf{\rv{X}\rv{Y}}^{(t)}  $ has valid marginals (see \autoref{algorithm:oblib_minimization_marginalization}) \;

		$t = t+1$;
	}
	\KwOut{$\mathsf{P}^*_{\rv{X}\rv{Y}} $}
	\caption{pf\_iterator(args)}
	\label{algorithm:oblib_minimization}
\end{algorithm}

\subsection{Maximization Algorithm}
For some fixed $\pmf{\rv{V}}$ that satisfies $H(\pmf{\rv{V}}) \geq \eta_2 $, consider the following maximization problem:
\begin{equation}
    \varphi_T(\eta) \triangleq \max_{
	\scaleto{
	\mathsf{P}_{\rv{W}}\colon  H(\mathsf{P}_{\rv{W}})\leq\eta_{1}}{7pt}
	}
		H(\mathsf{P}_{\rv{W}}*\mathsf{P}_{\rv{V}}).
\end{equation}
The respective Lagrangian is given by
\begin{equation} \label{eq:comib_modulo_max_lagrangian}
    L^{\varphi}(\bv{p},\lambda_1) = -h(T_v \bv{p}) + \lambda_1(h(\bv{p})-\eta_1),
\end{equation}
where $T_v$ is a transition matrix with columns being the cyclic permutations of $\bv{p}_v$. Maximizing $L^{\varphi}(\bv{p},\lambda_1)$ can be given an exact formal solution.
\begin{proposition}
The maximizer of \eqref{eq:comib_modulo_max_lagrangian}, $\bv{p}_w^*$, satisfies
\begin{equation}
    \bv{p}_{w}^* = \frac{e^{\beta_1 T_v^T \log \bv{q}_w}}{Z_1(\beta_1)},
\end{equation}
where $Z_1(\beta_1)$ is the partition function,
and $\bv{q}_w $ is given by
    $\bv{q}_w = T_v \bv{p}_w^*$.
\end{proposition}

The self-consistent
equations can be turned into converging,
alternating iterations as given in the following proposition.

\begin{proposition}
The set of self-consistent equations is satisfied simultaneously at the maxima of \eqref{eq:comib_modulo_max_lagrangian}, where the maximization is done independently over the convex set of the normalized distributions, $\bv{p}, \bv{q} \in \simplex{n} $. Namely,
\begin{equation}
    \max_{\bv{q} \in \simplex{n}} \max_{\bv{p} \in \simplex{n}} \Phi [\bv{p},\bv{q}] = \max_{\bv{q} \in \simplex{n}} \max_{\bv{p} \in \simplex{n}} h(\bv{q}) - \lambda_1 h(\bv{p}) .
\end{equation}
This maximization is performed by the converging alternating iterations. Denoting by $t$ the iterations step, we obtain \autoref{algorithm:oblib_modulo_maximization}.
\end{proposition}

\begin{algorithm}[b]
		\SetAlgoLined
		\SetKwInput{Init}{Initialize}
		\SetKwInput{Set}{Set}
		\KwIn{$ \bv{p}_v $,  and $ \beta_1 $}
		\Set{$T_v = $ cyclic permutations of $\bv{p}_v$}
		\Init{Arbitrary valid $ \bv{p}^{(0)} \in \simplex{n} $, $\bv{q}^{(0)} = T_v \bv{p}^{(0)}$ $t=1$.}
		\While{Variation in $ h(T_v \bv{p}) $ is greater then $ \epsilon $}{
			Compute $ \bv{p}^{(t)} = \frac{e^{\beta_1 T_v^T \log \bv{q}^{(t-1)}}}{Z_1(\beta_1)} $ \;
			Set $ \bv{q}^{(t)} = T_v \bv{p}^{(t)} $\;
			$t = t+1$;
		}
		\KwOut{$\bv{p}_w^* = \bv{p}^{(t-1)} $}
		\caption{pf\_modulo\_iterator(args)}
		\label{algorithm:oblib_modulo_maximization}
\end{algorithm}

\subsection{Minimization Algorithm}
In a very similar manner, fix $\pmf{\rv{W}}$ that satisfies $H(\pmf{\rv{W}}) \leq \eta_1 $, and consider the respective minimization problem, namely, \acrshort{ib} for modulo additive channels:
\begin{equation}
    g_T(\eta) \triangleq \min_{
	\scaleto{
	\mathsf{P}_{\rv{V}}\colon  H(\mathsf{P}_{\rv{V}})\geq\eta_{2}}{7pt}
	}
		H(\mathsf{P}_{\rv{W}}*\mathsf{P}_{\rv{V}}).
\end{equation}
The respective Lagrangian is given by
\begin{equation} \label{eq:comib_modulo_min_lagrangian}
    L^{g}(\bv{p},\lambda_2) = h(T_w \bv{p}) + \lambda_2(\eta_2- h(\bv{p})),
\end{equation}
where $T_w$ is a transition matrix with columns being the cyclic permutations of $\bv{p}_w$. Minimizing $L^{g}(\bv{p},\lambda_2)$ can be given an exact formal solution.
\begin{proposition}
The minimizer of \eqref{eq:comib_modulo_min_lagrangian},  $\bv{p}_{v}^*$,  satisfies
\begin{equation}
    \bv{p}_{v}^* = \frac{e^{\beta_2 T_w^T \log \bv{q}_v}}{Z_2(\beta_2)},
\end{equation}
where $Z_2(\beta_2)$ is the partition function, and $\bv{q}_v $ is given by
    $\bv{q}_v = T_w \bv{p}_v^*$.
\end{proposition}

The self-consistent equations can be turned into converging, alternating iterations as given in the following proposition.

\begin{proposition}
The set of self-consistent equations are satisfied simultaneously at the minima of \eqref{eq:comib_modulo_min_lagrangian}, where the minimization is done independently over the convex set of the normalized distributions, $\bv{p}, \bv{q} \in \simplex{n} $. Namely,
\begin{equation}
    \min_{\bv{q} \in \simplex{n}} \min_{\bv{p} \in \simplex{n}} \mathcal{G} [\bv{p},\bv{q}] = \min_{\bv{q} \in \simplex{n}} \min_{\bv{p} \in \simplex{n}} h(\bv{q}) - \lambda_2 h(\bv{p}) .
\end{equation}
This minimization is performed by the converging alternating iterations. Denoting by $s$ the iterations step, we obtain \autoref{algorithm:oblib_modulo_minimization}.
\end{proposition}

\begin{algorithm}[t]
		\SetAlgoLined
		\SetKwInput{Init}{Initialize}
		\SetKwInput{Set}{Set}
		\KwIn{$ \bv{p}_w $,  and $ \beta_2 $}
		\Set{$T_w = $ cyclic permutations of $\bv{p}_w$}
		\Init{Arbitrary valid $ \bv{p}^{(0)} \in \simplex{n} $, $\bv{q}^{(0)} = T_w \bv{p}^{(0)}$ $t=1$.}
		\While{Variation in $ h(T_w \bv{p}) $ is greater then $ \epsilon $}{
			Compute $ \bv{p}^{(s)} = \frac{e^{\beta_2 T_w^T \log \bv{q}^{(s-1)}}}{Z_2(\beta_2)} $ \;
			Set $ \bv{q}^{(s)} = T_w \bv{p}^{(s)} $\;
			$s = s+1$;
		}
		\KwOut{$\bv{p}_v^* = \bv{p}^{(s-1)} $}
		\caption{ib\_modulo\_iterator(args)}
		\label{algorithm:oblib_modulo_minimization}
\end{algorithm}

\subsection{Compound Algorithm}
We combine the maximization and minimization methods into alternating procedure in order to solve \eqref{eq:definition_compound_ib_modulo}, as described in \autoref{algorithm:modulo_hopping_optimization}.

\begin{algorithm}
	\SetAlgoLined
	\SetKwInput{Init}{Initialize}
	\SetKwInput{Set}{Set}
	\KwIn{$ n $, $ \eta_1 $ and $ \eta_2 $}
	\Init{$ \bv{p}_w^{(0)}  $ and $ \bv{p}_w^{(0)}  $  .}
	\While{Variation in $ h(\bv{p}_w^{(0)}*\bv{p}_v^{(0)}) $ is greater then $ \epsilon $}{
		\For{$ \beta_1 \in \field{R}_+ $}{
			$\bv{p}_w^*(\beta_1) = pf\_modulo\_iterator(\bv{p}_v^{(0)},\beta_1) $;
		}
		Find $\beta_1^* $ s.t. $ h(\bv{p}^*_{w}(\beta_1^*)) = \eta_1 $ ;\\
		\Set{$ \bv{p}^*_{w} (\beta_1^*) \mapsto \bv{p}_{w}^{(0)}$;}
		\For{$ \beta_2 \in \field{R}_+ $}{
			$\bv{p}^*_{v}(\beta_2) = ib\_modulo\_iterator(\bv{p}^{(0)}_{w},\beta_2) $;
		}
		Find $\beta_2^* $ s.t. $ h(\bv{p}_v (\beta_2^*)) = \eta_2 $;\\
		\Set{$ \bv{p}^*_{v} (\beta_2^*) \mapsto \bv{p}_{v}^{(0)}$;}
	}
	\KwOut{$ \bv{p}_w^*, \bv{p}_v^* $}
	\caption{\acrshort{comib} Modulo Programming}
	\label{algorithm:modulo_hopping_optimization}
\end{algorithm}

\subsection{Optimization Kernel}
We will begin our discussion by solving a simplified version of our problem, termed here as the optimization kernel.
The proof follows and extends a basic idea from the proof of  \cite[Lemma 7]{Witsenhausen1974}.
\begin{lemma} \label{lemma:optimization_kernel}
	Let $ \bv{x} = (x_1,x_2,x_3)^T \in \field{R}^3_+ $, and consider the extreme points of the following function
	\begin{equation}
		f(\bv{x}) = -\sum_{i=1}^3 (a x_i + b) \log (a x_i + b),
	\end{equation}
	where $ 0 \neq a \in \field{R} $ and $ b > 0 $, over the set defined by the following system of equations
	\begin{align}
		x_1 \geq x_2 \geq x_3 &\geq 0, \\
		x_1+x_2+x_3 &= c_1 \label{eq:oblib_optimization_kernel_sumX},\\
		-\sum_{i=1}^3 x_i \log x_i &= c_2 \label{eq:oblib_optimization_kernel_entropyX}.
	\end{align}
	The function $ f(\bv{x}) $ decreases as $ x_1 $ increases along the arc in $ \field{R}^3 $ defined by \eqref{eq:oblib_optimization_kernel_sumX} and \eqref{eq:oblib_optimization_kernel_entropyX}. That is  if $ c_2 < \log 2 $, then the maximum is obtained for $ \bv{x} = (x_1^*,c_1-x_1^*,0) $ where $ x_1^* $ is the root of
	\begin{equation} \label{eq:kernel_root1}
		c_2+x_1\log x_1 + (c_1-x_1) \log (c_1-x_1) = 0.
	\end{equation}
	Otherwise, if $ c_2 \geq \log 2 $, then the maximum is obtained for $ \bv{x} = (x_1^*,x_1^*, c_1-2x_1^*) $ where $ x_1^* $ is the root of
	\begin{equation} \label{eq:kernel_root2}
		c_2+2x_1\log x_1 + (c_1-2x_1) \log (c_1-2x_1) = 0,
	\end{equation}
	and the minimum is obtained for $ \bv{x} = (\hat{x}_1,(c_1-\hat{x}_1)/2,(c_1-\hat{x}_1)/2) $ where $ \hat{x}_1 $ is the root of
	\begin{equation} \label{eq:kernel_root3}
		c_2+x_1\log x_1 + (c_1-x_1) \log (c_1-x_1)/2 = 0.
	\end{equation}
\end{lemma}

\begin{proof}

	The relations  \eqref{eq:oblib_optimization_kernel_sumX} and \eqref{eq:oblib_optimization_kernel_entropyX} imply
		$\mathrm{d} x_1 + \mathrm{d} x_2 + \mathrm{d} x_3 = 0$, and
		$\sum_{i=1}^3 (1+\log x_i) \mathrm{d} x_i = 0$,
	which further indicate
	\begin{equation}
		\mathrm{d} x_3 = \frac{\log x_1 - \log x_2}{\log x_2 - \log x_3} \mathrm{d}x_1.
	\end{equation}
	Thus
	\begin{align}
		&\mathrm{d} f(\bv{x})
		= -\sum_{i=1}^3 a(1+\log(ax_i+b)) \mathrm{d}x_i
		= -a\sum_{i=1}^3 \log(ax_i+b) \mathrm{d}x_i \\
		&= -a \mathrm{d} x_1 \bigg[
		\log (ax_1+b) - \log (ax_2+b)
		 - \frac{\log x_1 - \log x_2}{\log x_2 - \log x_3} \left( \log (ax_2+b) - \log (ax_3+b) \right)
		\bigg] \label{eq:oblib_optimization_kernel_df}.
	\end{align}
	Consider the function $ \phi(t) \triangleq a\log (ae^t+b) $. Note that
		$\frac{\mathrm{d} \phi}{\mathrm{d} t} = \frac{a^2 e^t}{a e^t+b}$,
	and
	\begin{equation}
		\frac{\mathrm{d}^2 \phi}{\mathrm{d} t^2} = \frac{a^2 e^t(a e^t+b) - a^3 e^{2t}}{(a e^t+b)^2} = \frac{a^2 b e^t }{(a e^t+b)^2} > 0,
	\end{equation}
	where the last inequality follows since $ b > 0 $. Thus $ \phi(t) $ is a convex function and by Chordal Slope Lemma \cite[Ch. 6.6]{Royden2010} we have with $ t_i = \log x_i $ that
	\begin{equation}
		\frac{a  \log  (  ax_1 +  b  )  -   a  \log ( ax_2  +  b  )}{\log x_1 - \log x_2} \geq  \frac{a  \log   (  ax_2  +  b  )  -   a \log  ( ax_3  +  b )}{\log x_2 - \log x_3}.
	\end{equation}
	Plugging this inequality in \eqref{eq:oblib_optimization_kernel_df} implies that the expression inside the brackets is strictly positive, and thus increasing $ x_1 $ results in decreasing of $ f(\bv{x}) $.

	If $ c_2 < \log 2 $ the endpoint that corresponds to the maximum value satisfies $ (x_1,c_1-x_1,0) $, where $ x_1 $ can be found by solving \eqref{eq:kernel_root1}.
	If $ c_2 \geq \log 2 $ the endpoint that corresponds to the maximum value satisfies $(x_1,x_1,c_1-2x_1) $ where $ x_1 $ can be found by solving  \eqref{eq:kernel_root2}.
	For any $ c_2 $, the minimum value is obtained for $ (x_1,(c_1-x_1)/2,(c_1-x_1)/2,) $ where $ x_1 $ is found by solving  \eqref{eq:kernel_root3}.
\end{proof}

We proceed to solve the first direction of the saddle point property, i.e., we will solve a maximization problem. The result is summarized in the following proposition.

\begin{proposition} \label{proposition:comib_modulo_maximization}
	Suppose that $ \mathsf{P}_{\rv{V}} $ is a regular $ (\alpha,n) $ Hamming. Then the solution of
	\begin{equation} \label{eq:oblib_maximize_problem2}
		\max_{\mathsf{P}_{\rv{W}}\colon H(\mathsf{P}_{\rv{W}}) \leq \eta_{1}}  H(\mathsf{P}_{\rv{W}}*\mathsf{P}_{\rv{V}}),
	\end{equation}
	is a negative Hamming channel $ (\beta,n,k) $.
\end{proposition}
\begin{remark}
	Note that here we do not have any constraint on $ k $, i.e., this result holds for any entropy constraint (i.e., value of $\eta_1$).
\end{remark}

\begin{proof}
	The underlying optimization problem is a maximization of a concave function over the complement of a convex set, therefore, the optimal value lies on the boundary of the set, that is, we may restrict to all $ \mathsf{P}_{\rv{W}} $ satisfying $ H(\mathsf{P}_{\rv{W}}) = \eta_1 $.  Since $ h_n\left( \alpha \bv{p} + (1-\alpha) \bv{u}\right) $ and $ h_n(\bv{p}) $ are both invariant under permutations, one may assume $ p_1 \geq p_2 \geq \cdots \geq p_n $.
	Thus, \eqref{eq:oblib_maximize_problem2} may be reformulated in the following standard form:
	\begin{equation} \label{eq:oblib_maximize_problem3}
		\begin{aligned}
			\max_{\bv{p} \in  \field{R}^n} \quad & -\sum_{i=1}^n \left( \alpha p_i + \frac{(1-\alpha)}{n} \right) \log \left( \alpha p_i + \frac{(1-\alpha)}{n} \right) \\
			\textrm{s.t.} \quad & -\sum_{i=1}^n p_i \log p_i = \eta_1,
			\quad \sum_{i=1}^n p_i = 1\\
			& 1\geq p_1 \geq p_{2} \geq \cdots \geq p_n \geq 0  .
		\end{aligned}
	\end{equation}

	\begin{itemize}
		\item For $ n=3 $, \eqref{eq:oblib_maximize_problem3} is exactly the problem defined in \autoref{lemma:optimization_kernel} with $ a= \alpha  > 0$, $ b= \bar{\alpha} /n >0$, $ c_1=1 $ and $ c_2 = \eta_1 $, thus
		\begin{itemize}
			\item if $ \eta_1 < \log _2 $, the maximizer is $ \bv{p}^* = (p_1^*,1-p_1^*,0) $, where $ p_1^* = h_b^{-1}(\eta_1) $.
			\item otherwise, if $ \eta_1 \geq \log _2 $, the maximizer is $ \bv{p}^* = (p_1^*,p_1^*,1-2p_1^*) $, where $ p_1^* $ is the root of $ h_b( 2p_1^*) + 2p_1^* \log 2 = \eta_1 $ .
		\end{itemize}
		\item Suppose $ n>3 $ and assume by contradiction that $ \bv{p}^* $ is not a negative Hamming. Thus, there exist $ k_1,k_2,k_3 \in [n] $ with $ 1 > p^*_{k_1} > p^*_{k_2} > p^*_{k_3} >0 $. We will show that the output entropy can be further increased, thus contradicting the optimality of $ \bv{p}^* $. Assume that the remaining indices are kept fixed, thus they contribution to the output entropy is not changed. We are interested in the following problem:
		\begin{equation} \label{eq:oblib_maximize_problem_converse}
			\begin{aligned}
				&\max_{
					[p_{k_1}, p_{k_2}, p_{k_3}]^T
					 \in  \field{R}^3}
				&& -  \sum_{i=1}^3    \left(  \alpha p_{k_i}   +  \frac{(1-\alpha)}{n}   \right)   \log   \left(   \alpha p_{k_i}   +  \frac{(1-\alpha)}{n}   \right) \\
				&\textrm{subject to}
				&& -\sum_{i=1}^3 p_{k_i} \log p_{k_i} = \eta_1 + \sum_{i \notin \{k_1,k_2,k_3\} }    p_{i} \log p_{i} \\
				&&& 1\geq p_{k_1} \geq p_{k_2} \geq p_{k_3} \geq 0 ,
				\quad \sum_{i=1}^n p_{k_i} = 1 - \sum_{i \notin \{k_1,k_2,k_3\} } p_{i}.
			\end{aligned}
		\end{equation}
		The problem defined in \autoref{lemma:optimization_kernel} is identical to \eqref{eq:oblib_maximize_problem_converse} with $ a= \alpha  > 0$, $ b= \bar{\alpha} /n >0$, $ c_1=1 - \sum_{i \notin \{k_1,k_2,k_3\} } p_{i}$ and $ c_2 = \eta_1 + \sum_{i \notin \{k_1,k_2,k_3\} } p_{i} $, but it has a different maximizer, thus contradicting the optimality of $ \bv{p}^* $.
	\end{itemize}

\end{proof}

Now, we go ahead to solve the reverse direction of the saddle point property, i.e., we will solve a minimization problem. The result is summarized in the following proposition.

\begin{proposition} \label{proposition:comib_modulo_minimization}
	Suppose that $ \mathsf{P}_{\rv{W}} $ is a negative $(\beta,n,n) $ Hamming satisfying $ H(\mathsf{P}_{\rv{W}}) \geq \log (n-1) $. Then, the solution of
	\begin{equation} \label{eq:oblib_minimize_problem2}
	 \min_{\mathsf{P}_{\rv{V}}\colon H(\mathsf{P}_{\rv{V}}) \geq \eta_{2}}  H(\mathsf{P}_{\rv{W}}*\mathsf{P}_{\rv{V}}),
	\end{equation}
	is a regular Hamming $  \mathsf{P}_{\rv{V}} $ with parameter $\alpha$.
\end{proposition}
The proof of \autoref{proposition:comib_modulo_minimization} is omitted due to space limitation and its resemblance to the proof of \autoref{proposition:comib_modulo_maximization}.

\subsection{Minimization Problem}
For a fixed $ \mathsf{P}_{\rv{Z}|\rv{Y}} $, we denote by $ f(C_1;\mathsf{P}_{\rv{Z}|\rv{Y}}) $  the output of the inner minimization problem given in \eqref{eq:definition_oblivious_ib}, i.e.,
\begin{equation}
	f(C_1;\mathsf{P}_{\rv{Z}|\rv{Y}}) \triangleq  \min_{\mathsf{P}_{\rv{X} \rv{Y}}: I(\rv{X};\rv{Y}) \geq C_1} I(\rv{X};\rv{Z}).
\end{equation}
This problem is equivalent (it has the same optimal joint \acrshort{pmf} $ \mathsf{P}_{\rv{X}\rv{Y}} $), to
\begin{equation} \label{eq:definition_oblivious_ib_minimization}
	\psi(\eta)  = \psi(\eta;\mathsf{P}_{\rv{Z}|\rv{Y}}) \triangleq  \max_{\mathsf{P}_{\rv{X} \rv{Y}}: H(\rv{Y}|\rv{X}) = \eta} H(\rv{Z}|\rv{X}).
\end{equation}
We denote the upper concave envelope of $ \psi(\eta) $ as $ \varphi(\eta) $. Namely, $ \varphi(\eta) $ is the minimal concave function such that $ \varphi(\eta) \geq \psi(\eta) $.

\begin{remark}
	Note that in general $ \psi(\eta) $ is not concave, but, as was shown in, if the condition $ \sum_{y}
	 \mathsf{P}_{\rv{X}\rv{Y}}(x,y) = \mathsf{P}_{\rv{X}} (x) $ is omitted, then it is a concave function.
\end{remark}

\begin{theorem}
	The supremum of the conditional output entropy rate of a sequence of $ K $ outputs $ \rv{Z}^K $ with respective input $ \rv{X}^K $ over all joint distributions $ \mathsf{P}_{\rv{X}^K\rv{Y}^k} $ and $ K >0 $ is
	\begin{equation}
		\varphi(\eta) = \sup_{K >0 . \mathsf{P}_{\rv{X} \rv{Y}}: \frac{1}{K} H(\rv{Y}^K|\rv{X}^K) = \eta} \frac{1}{K} H(\rv{Z}^K|\rv{X}^K),
	\end{equation}
	namely, the RHS is a concave function of $ \eta $.
\end{theorem}

\begin{proof}
	Since the underlying channels are memoryless, and utilizing the definitions of $ \psi(\cdot) $ and $ \varphi(\cdot) $, we obtain the following list of inequalities:
	\begin{align}
		\frac{1}{K} H(\rv{Z}^K|\rv{X}^K)
		&= \frac{1}{K} \sum_{k=1}^{K} H(\rv{Z}_k|\rv{X}_k) \\
		&\leq \frac{1}{K} \sum_{k=1}^{K} \psi (H(\rv{Y}_k|\rv{X}_k)) \\
		&\leq \frac{1}{K} \sum_{k=1}^{K} \varphi (H(\rv{Y}_k|\rv{X}_k)) \\
		&\leq \varphi \left( \frac{1}{K} \sum_{k=1}^{K}  H(\rv{Y}_k|\rv{X}_k) \right) \\
		&= \varphi \left(\frac{1}{K} H(\rv{Y}^K|\rv{X}^K)\right) \\
		&= \varphi(\eta).
	\end{align}
	If $ \psi(\eta) = \varphi(\eta ) $, then equality achieved for any $  K \geq 1 $. If $ \psi(\eta) < \varphi(\eta ) $, then the point $ (\eta,\varphi(\eta)) $ is a convex combination of two points $ (\eta_1,\psi(\eta_1)) $ and $ (\eta_2,\psi(\eta_2)) $, of the graph of $ \psi(\cdot) $, corresponding to joint probability vectors $ \mathsf{P}_{\rv{X}\rv{Y}}^{(1)} $ and $ \mathsf{P}_{\rv{X}\rv{Y}}^{(2)} $. Namely. there exists $ \lambda \in [0,1] $, s.t.,
	\begin{equation}
		\eta = \lambda \eta_1 + (1-\lambda)\eta_2  \qquad \text{ and } \qquad
		\varphi(\eta) = \lambda \psi(\eta_1) + (1-\lambda) \psi(\eta_2).
	\end{equation}
	If $ \lambda $ is rational, i.e, there exist $ a,b  \in \field{N} $ and $ \lambda = a/b $, take independent $ \{(\rv{X}_i,\rv{Y}_i)\}_{i=1}^b $ of which $ a $ have distribution $  \mathsf{P}_{\rv{X}\rv{Y}}^{(1)} $ and $ b-a $ have distribution $ \mathsf{P}_{\rv{X}\rv{Y}}^{(2)} $, i.e.,
	\begin{equation}
		\mathsf{P}_{\rv{X}^b \rv{Y}^b} (x^b,y^b) = \prod_{i=1}^{a} \mathsf{P}_{\rv{X}\rv{Y}}^{(1)}(x_i,y_i)
		\cdot
		\prod_{j=a+1}^{b} \mathsf{P}_{\rv{X}\rv{Y}}^{(2)}(x_j,y_j)
	\end{equation}
	we obtain the point $ (\eta,\varphi(\eta)) $, and achieve equality. If $ \lambda $ is irrational, then for any $ 0 < \epsilon < \eta_2 $, we have
	\begin{equation}
		\eta = \lambda_\epsilon \eta_1 + (1-\lambda_\epsilon) (\eta_2 - \epsilon)
	\end{equation}
	with
	\begin{equation}
		\lambda_\epsilon \triangleq \frac{\lambda (\eta_2-\eta_1)-\epsilon}{(\eta_2-\eta_1) - \epsilon}.
	\end{equation}
	Thus,
	\begin{equation}
		\zeta_{\epsilon} = \lambda_\epsilon \psi(\eta_1) + (1-\lambda_\epsilon) \psi(\eta_2 - \epsilon)
	\end{equation}
	is achievable when $ \lambda_\epsilon $ is rational. Taking $ \epsilon \rightarrow 0 $, $ \zeta_{\epsilon} \rightarrow \varphi(\eta) $ and since $ \lambda_\epsilon \rightarrow \lambda $ through an interval, in which rationals are dense, the supremum cannot be less than $ \varphi(\eta) $, completing the proof.
\end{proof}

\begin{corollary}
	Let
	\begin{equation}
		f^{(K)}(C_1;\mathsf{P}_{\rv{Z}^K|\rv{Y}^K}) \triangleq  \min_{\mathsf{P}_{\rv{X}^K \rv{Y}^K}: \frac{1}{K} I(\rv{X}^K;\rv{Y}^K) \geq C_1} \frac{1}{K} I(\rv{X}^K;\rv{Z}^K).
	\end{equation}
	then $ \inf_{K >0} f^{(K)}(C_1;\mathsf{P}_{\rv{Z}^K|\rv{Y}^K}) $  is a convex function of $ C_1 $.
\end{corollary}

\subsection{The Composite Problem}
Consider next the compound setting,
\begin{equation}
	R^{(K)}(C_1,C_2) = \max_{\mathsf{P}_{\rv{Z}^K| \rv{Y}^K}:  \frac{1}{K}I(\rv{Y}^K;\rv{Z}^K)\leq C_2} \inf_{K >0} f^{(K)}(C_1;\mathsf{P}_{\rv{Z}^K|\rv{Y}^K})
\end{equation}

\begin{proposition}
	$ R^{(K)} (C_1,C_2) $ is convex in $ C_2 $ for fixed $ C_1 $ as a pointwise supremum over an infinite set of convex functions (see \cite[p. 81]{Boyd2014}).
\end{proposition}

\begin{remark}
	For a fixed $ \mathsf{P}_{\rv{X}\rv{Y}} $ the function
	\begin{equation}
		R^{\text{ib}}(C) = \max_{\mathsf{P}_{\rv{Z}| \rv{Y}}: I(\rv{Y};\rv{Z})\leq C}  I(\rv{X};\rv{Z}).
	\end{equation}
	is concave in $ C $. This is a known result for the standard \acrshort{ib} problem.
\end{remark}

\subsection{A special case of $ \alpha = 0.5 $}

\subsubsection{The Supper-Symmetric \acrshort{tito}}
First, consider a specific, super-symmetric case where $ \alpha = \beta $ in \eqref{eq:tito_transition_matrix}. This is also known as  the third order Hamming channel with parameter $ \alpha $ \cite[Sec. IV.E]{Witsenhausen1975}. For this, super-symmetric formulation we have the following characterization of the channel $ \mathsf{P}_{\rv{Y}|\rv{Z}} $.
\begin{proposition} \label{proposition:tito3_hamming}
	The optimal channel from $ \rv{Y} $ to $ \rv{Z} $ is also a super-symmetric channel with parameter $ \delta $, where
	$ \delta = g_3^{-1}(C)  $, with $ g_3(x) \triangleq \log 3 +2x\log x + (1-2x) \log(1-2x) $.
\end{proposition}
\begin{proof}
	For $ \alpha = \beta $,
	\begin{equation}
		\mat{T} \bv{p} =
		\begin{pmatrix}
			(1-2\alpha)  p + \alpha r + \alpha(1-p-r) \\
			\alpha p+(1-2\alpha)  r  + \alpha(1-p-r) \\
			\alpha p+\alpha  r  + (1-2\alpha)(1-p-r)
		\end{pmatrix}
		=
		\begin{pmatrix}
			(1-3\alpha)  p  + \alpha \\
			(1-3\alpha)  r  + \alpha \\
			(1-3\alpha)(1-p-r) + \alpha
		\end{pmatrix}
	\end{equation}
	Therefore,
	\begin{align}
		\nabla_{\bv{p}} \phi(\bv{p},\lambda)
		&=
		\begin{pmatrix}
			\frac{\partial \phi}{\partial p} \\
			\frac{\partial \phi}{\partial r}
		\end{pmatrix} \\
		&=
		\begin{pmatrix}
			(1-3\alpha) \log \frac{(1-3\alpha)(1-p-r) + \alpha}{(1-3\alpha)  p + \alpha} + \lambda \log \frac{p}{1-p-r} \\
			(1-3\alpha) \log \frac{(1-3\alpha)(1-p-r) + \alpha}{(1-3\alpha)  r + \alpha} + \lambda \log \frac{r}{1-p-r}
		\end{pmatrix}
	\end{align}
	The local minimum condition implies that $ \nabla_{\bv{p}} \phi(\bv{p},\lambda) = \bv{0} $, therefore,
	\begin{equation}
		\lambda \log \frac{p}{r} + (1-3\alpha) \log \frac{(1-3\alpha)r+\alpha}{(1-3\alpha)p+\alpha} = 0.
	\end{equation}
	Thus, the minimum point satisfies $ \bv{p}_1= \bv{p}_\lambda = (p_\lambda,p_\lambda,1-2p_{\lambda}) $. Since $ \bv{p}_2 = \mat{\Pi}_2 \bv{p}_{\lambda} $ and $ \bv{p}_3 = \mat{\Pi}_3 \bv{p}_{\lambda} $ are also minimal points of $ \phi(\bv{p},\lambda) $, the transition matrix from $ \rv{Z} $ to $ \rv{Y} $ which columns are the vectors $ \bv{p}_1 $, $ \bv{p}_2 $ and $ \bv{p}_3 $, is also a super-symmetric transition matrix. Moreover, for the Hamming setting, $ p_\lambda $ can be easily computed via the following procedure.
	\begin{align}
		C &= H(\rv{Y}) -   H(\rv{Y}|\rv{Z})\\
		& =\log(3) - h_3(\bv{p}_\lambda) \\
		&= \log(3) +p_\lambda \log p_\lambda + p_\lambda \log p_\lambda+(1-2 p_\lambda) \log (1-2p_\lambda) \\
		&= \log 3  +2p_\lambda \log p_\lambda + (1-2 p_\lambda) \log (1-2p_\lambda) \\
		&\triangleq g_3(p_\lambda)
	\end{align}
	Thus, $ p_\lambda = g_3^{-1}(C) $.
\end{proof}

We further verify \autoref{proposition:tito3_hamming} via numerical simulations. \autoref{figure:tito_hamming_vs_alpha_C_1} shows the probability vector $ \rv{W} $ and $ R_{T_\alpha}(C) $ for various values of $ \alpha $. We observe that the numerical optimization agrees with theoretic arguments in \autoref{proposition:tito3_hamming}.

Since $ \rv{W} $ is independent of the choice of $ \alpha $, we further freeze $ \alpha = 0.5 $ and compare it with respect to the value of $ C $. The results are shown in \autoref{figure:tito_hamming_vs_C_alpha_0p5}.

\begin{figure}
	\centering
	\begin{subfigure}[t]{0.49\textwidth}
		\input{../figures/titoIB/tito_hamming_Pw_vs_C_alpha_0p5}
		\caption{}
	\end{subfigure}
	\begin{subfigure}[t]{0.49\textwidth}
		\input{../figures/titoIB/tito_hamming_R_vs_C_alpha_0p5}
		\caption{}
	\end{subfigure}
	\caption{(a) Optimal $ \bv{w} $ for $ \alpha = 0.5 $ vs $ C $. (b) $ R_{T_{\alpha}} (C) $ vs $ C $ for $ \alpha=0.5 $.}
	\label{figure:tito_hamming_vs_C_alpha_0p5}
\end{figure}

Are there other  channels that their Information Bottleneck function can be fully characterized? In the following section we consider  a yet another channel defined by a single parameter.

\subsubsection{The Noisy Typewriter \acrshort{tito}}
Here we consider the case where $ \beta = 0 $ in \eqref{eq:tito_transition_matrix}. This is also known as  the noisy typewriter channel with parameter $ \alpha $. Unfortunately, for this formulation we don't have a simple characterization of the test-channel from $ \rv{Z} $ to $ \rv{Y} $. For $ \beta = 0 $, the transition matrix \eqref{eq:tito_transition_matrix} is given by:
	\begin{equation}
		\mat{T} \bv{p} =
		\begin{pmatrix}
			(1-\alpha)  p + \alpha r  \\
			(1-\alpha)  r  + \alpha(1-p-r) \\
			\alpha p+ (1-\alpha)(1-p-r)
		\end{pmatrix}
		=
		\begin{pmatrix}
			(1-\alpha)  p  + \alpha r \\
			(1-2\alpha)  r  + \alpha(1-p) \\
			(1-\alpha)(1-r) -(1-2\alpha) p
		\end{pmatrix}.
	\end{equation}
	Therefore,
	\begin{align}
		\nabla_{\bv{p}} \phi(\bv{p},\lambda)
		&=
		\begin{pmatrix}
			\frac{\partial \phi}{\partial p} \\
			\frac{\partial \phi}{\partial r}
		\end{pmatrix} \\
		&=
		\begin{pmatrix}
			(1-2\alpha) \log \left( \bar{\alpha} \left( 1-p-r \right) + \alpha p \right) + \alpha \log (\alpha \bar{p}+ (1-2\alpha)r) - \bar{\alpha} \log (\bar{\alpha} p +\alpha r)
			+ \lambda \log \frac{p}{1-p-r} \\
			\bar{\alpha} \log \left( \bar{\alpha} \left( 1-p-r \right) + \alpha p \right) - (1-2\alpha) \log (\alpha \bar{p}+ (1-2\alpha)r) - \alpha \log (\bar{\alpha} p +\alpha r) + \lambda \log \frac{r}{1-p-r}
		\end{pmatrix}
	\end{align}
	The local minimum condition implies that $ \nabla_{\bv{p}} \phi(\bv{p},\lambda) = \bv{0} $. Our initial conjecture was that the optimal test-channel here from $ \rv{Z} $ to $ \rv{Y} $ would be also noisy typewriter channel, but as you may observe,  choosing $ p+r=1 $  will not nullify the gradient.

	To further convince the reader, we find the optimal test-channel using numerical optimization methods. The result for various values of $ \alpha $ is shown in \autoref{figure:tito_typewriter_vs_alpha_C_0p8}. As can be observed, the optimal test-channel from $ \rv{Z} $ to $ \rv{Y} $ does depend on $ \alpha $.

	We further freeze $ \alpha = 0.2 $ and plot the optimal test-channels for various values of $ C $ as seen in \autoref{figure:tito_typewriter_vs_C_alpha_0p2}. Our only observation is that the optimal test-channel is not noisy-typewriter,  however, one of the elements of $ P_{\rv{W}} $ rapidly tends to zero (red curve in \autoref{figure:tito_typewriter_vs_C_alpha_0p2} (b)).

\begin{figure}
	\centering
	\begin{subfigure}[t]{0.49\textwidth}
		\input{../figures/titoIB/tito_typewriter_Pw_vs_alpha_C_0p8}
		\caption{}
	\end{subfigure}
	\begin{subfigure}[t]{0.49\textwidth}
		\input{../figures/titoIB/tito_typewriter_R_vs_alpha_C_0p8}
		\caption{}
	\end{subfigure}
	\caption{(a) Optimal $ \bv{w} $ for $ C= 0.8 $ vs $ C $. (b) $ R_{T_{\alpha}} (C) $ vs $ \alpha $ for $ C=0.8 $.}
	\label{figure:tito_typewriter_vs_alpha_C_0p8}
\end{figure}

\begin{figure}
	\centering
	\begin{subfigure}[t]{0.49\textwidth}
		\input{../figures/titoIB/tito_typewriter_Pw_vs_C_alpha_0p2}
		\caption{}
	\end{subfigure}
	\begin{subfigure}[t]{0.49\textwidth}
		\input{../figures/titoIB/tito_typewriter_R_vs_C_alpha_0p2}
		\caption{}
	\end{subfigure}
	\caption{(a) Optimal $ \bv{w} $ for $ \alpha= 0.2 $ vs $ C $. (b) $ R_{T_{\alpha}} (C) $ vs $ C $ for $ \alpha=0.2 $.}
	\label{figure:tito_typewriter_vs_C_alpha_0p2}
\end{figure}

\subsection{The oblivious symmetric bottleneck with $ m=3 $.}
In this section we solve numerically the optimization problem defined in \eqref{eq:definition_sym_oblivious_ib} for $ m=3 $. \autoref{figure:tito_maxmin_iXZ_T_vs_Pw_C_1} shows $ \min_{\mathsf{P}_{\rv{W}}: H(\rv{W})= \log_2(3) - C_1} I(\rv{X};\rv{Z}) $ for various values of $ P_{\rv{V}} $ that satisfy $ H(\rv{V}) = \log_2(3)-C_2 $. We observe that for any choice of $ P_{\rv{V}} $ the minimizer of $ I(\rv{X};\rv{Z}) $ is a noisy-typewriter channel. The optimal $ \mathsf{P}_{\rv{W}} $ and $ \mathsf{P}_{\rv{V}} $ versus  $ C_1 = C_2 = C $ are plotted in \autoref{figure:tito_maxmin_iXZ_T_vs_C}. Note the bifurcation phenomena for $ \mathsf{P}_{\rv{W}} $ when the entropy of $ \rv{W} $ crosses $ H(\rv{W}) = 1 $. We observe that the optimal $ \mathsf{P}_{\rv{W}} $ is the noisy typewriter for $ H(\rv{W}) <1 $, while for $ H(\rv{W}) >1 $ it is a Hamming channel. On the other hand, the optimal $ \mathsf{P}_{\rv{V}} $ is a Hamming channel for any $ C $.

\begin{figure}
	\centering
	\input{../figures/titoIB/tito_maxmin_iXZ_T_vs_Pw_C_1}
	\caption{Evaluation of $ \min I(\rv{X};\rv{Z}) $ and the respective $ P_{\rv{W}} $ that achieves the minimum for various values of $ P_{\rv{V}} $ for $ C_1=C_2=1 $ .}
	\label{figure:tito_maxmin_iXZ_T_vs_Pw_C_1}
\end{figure}

\begin{figure}
	\centering
	\input{../figures/titoIB/tito_maxmin_iXZ_T_vs_C}
	\caption{Evaluation of $ R(C,C) $ and the respective optimal $ \mathsf{P}_{\rv{W}} $ and $ \mathsf{P}_{\rv{V}} $ for various values of  $ C_1=C_2=C $ .}
	\label{figure:tito_maxmin_iXZ_T_vs_C}
\end{figure}

\subsection{Multilayer \acrlong{ib}}

In this work, we address the compound setting for the \acrshort{ib} problem, with the goal of providing similar results. First, we address the \acrshort{dsbs} and Gaussian (scalar and vector) settings. Second, we analyze the \acrshort{kld} class for $\set{P}_{\set{X}\set{Y}}$ for the particular choice of jointly Gaussian random variables.
Then, we consider general modulo additive channels, with modulo additive representations, and provide various bounds on the \acrshort{comib} function with \acrshort{pf}-based compound set, and then with \acrshort{tv}-based compound set. Finally, we return to the general discrete alphabet case with \acrshort{pf} based compound set, and propose an alternating algorithm, which essentially iterates between the maximization over $\mathsf{P}_{\rv{Z}| \rv{Y}}$ (an \acrshort{ib} problem) and  minimization over $\mathsf{P}_{\rv{X} \rv{Y}}$ (a \acrshort{pf} problem).
We further specialize this algorithm to the modulo-additive setting, obtaining an elegant and efficient computational method.

\subsection{Related work} \label{section:related_work}
In many problems in learning, there is an interest to represent data $\rv{Y}$, with a compressed version  $\rv{Z}$ that captures as much relevant information as possible with a fixed number of bits. One possible approach to handle such problem is via rate distortion theory for lossy source coding. However, the utilization of rate distortion theory, requires specifying a distortion function first, and it is usually intractable finding such function for real data scenarios. A pioneering work by Tishby et al. \cite{Tishby1999}, suggested the \acrshort{ib} framework, where additional variable $\rv{X}$ determines relevance (for example, it can be the labeling of the data). The quality of distortion is measured by the mutual information between $\rv{X}$ and $\rv{Z}$, thus revealing a more natural distortion measure.
This framework is closely related to a variety of problems in information theory, such as \textit{remote source coding} \cite{Wolf1970},  \acrfull{ceb} \cite{Witsenhausen1975}, \acrlong{cr} \cite{Steinberg2009}, and \acrfull{infcomb} \cite{Sutskover2005,Land2006}. See an overview in a recent comprehensive tutorial on the
IB method and related problems \cite{Zaidi2020}.
Applications of the \acrshort{ib} problem in machine-learning are detailed in \cite{Chechik2005,Tishby2015,Farajiparvar2018,Goldfeld2020}.

In the coding-theoretic context, it has been recently shown that the \acrshort{ib} method can be used to reduce the data transfer rate and computational complexity in 5G low-density parity check (LDPC) decoders \cite{Lewandowsky2016,Stark2020}. Furthermore, it is also related to construction of good polar codes \cite{Bhatt2021}.  In this problem, the value of the capacity of the polarized channels is required in order to identify the location of "frozen bits" in the codeword. However, the  output-alphabet size of the polarized channels increases exponentially, and so quantization is employed in order to reduce the computational complexity.
The quality of the quantization scheme is then assessed via  mutual information preservation. It can be shown that the corresponding \acrshort{ib} problem upper bounds the mutual-information after quantization technique.

General quantization algorithms based upon the \acrshort{ib} method were considered in \cite{Stark2018,Shah2019, Shah2019a}.
Furthermore, a relationship between the KL means algorithm, and the \acrshort{ib} method has been discovered in \cite{Kurkoski2017}.
In \cite{Pensia2020} a robust \acrshort{ib} program was proposed, with the goal of extracting features that are simultaneously relevant and robust. Unlike in this paper, therein the channel from $ \rv{Y} $ to $ \rv{Z} $ is made robust, as measured in terms of Fisher Information.

With more generality, the \acrshort{ib} problem connects to many other timely aspects. These include game theory and Nash equilibrium \cite{Kazikli2020}, \emph{capital investment} \cite{Erkip1998}, \emph{distributed learning} \cite{Farajiparvar2018}, \emph{deep learning} \cite{Tishby2015,Alemi2017,ShwartzZiv2017,Gabrie2019,Goldfeld2018} and \emph{convolutional neural networks} \cite{Cheng2018,Yu2021}.

\subsection{Notations and Preliminaries}
\input{./sections/notations}

\input{./sections/preliminaries}

\subsection{The Inner Lagrangian}
\input{./sections/lagrangian}
\subsection{The Outer Lagrangian}
\input{./sections/outer_lagrangian}

\subsection{The Compound Algorithm}
\input{./sections/compound_algorithm}

\subsection{Witsenhausen Approach}
Here we apply a directional derivative technique similar to \cite{Witsenhausen1974}. We will need to consider three different configurations for the the triple $ k_1,k_2,k_3 $: the case where they are consecutive; the case where two of them are adjacent; and the case where none of them are adjacent.
\begin{itemize}
	\item In the case $ k_1 < k_1+1 < k_2 < k_2+1< k_3 $, the minimization problem takes the following form:
	\begin{equation}
		\begin{aligned}
			\min_{(p_{k_1},p_{k_2},p_{k_3}) \in  \field{R}^3} \quad & \sum_{i=1}^{3} h\left(  \beta^+ \gamma_{k_i} + \beta^-  p_{k_i} \right) + h\left( \beta^+ p_{k_i} + \beta^-  \bar{\gamma}_{k_i} \right) \\
			\textrm{s.t.} \quad & -\sum_{i=1}^3 p_{k_i} \log p_{k_i} = c_2\\
			& \sum_{i=1}^3 p_{k_i} = c_1.
		\end{aligned}
	\end{equation}

	The directional derivative of the objective function is given by
	\begin{align}
		&-\left[\beta^- \log(\beta^+ \gamma_{k_1} + \beta^- p_{k_1}) + \beta^+ \log(\beta^+ p_{k_1} + \beta^- \bar{\gamma}_{k_1}) \right] \mathrm{d} p_{k_1} \\
		&-\left[\beta^- \log(\beta^+ \gamma_{k_2} + \beta^- p_{k_2}) + \beta^+ \log(\beta^+ p_{k_1} + \beta^- \bar{\gamma}_{k_2}) \right] \mathrm{d} p_{k_2} \\
		&-\left[\beta^- \log(\beta^+ \gamma_{k_3} + \beta^- p_{k_3}) + \beta^+ \log(\beta^+ p_{k_3} + \beta^- \bar{\gamma}_{k_3}) \right] \mathrm{d} p_{k_3}  \\
		&=- \mathrm{d} p_{k_1} \bigg[\beta^- \log \left(\frac{\beta^+ \gamma_{k_1} + \beta^- p_{k_1}}{\beta^+ \gamma_{k_2} + \beta^- p_{k_2}} \right)  + \beta^+ \log \left( \frac{\beta^+ p_{k_1} + \beta^- \bar{\gamma}_{k_1}}{\beta^+ p_{k_2} + \beta^- \bar{\gamma}_{k_2}} \right)  \\
		&-\frac{\log p_{k_1} - \log p_{k_2}}{\log p_{k_2} - \log p_{k_1}} \left[
		\beta^- \log \left(\frac{\beta^+ \gamma_{k_2} + \beta^- p_{k_2}}{\beta^+ \gamma_{k_3} + \beta^- p_{k_3}} \right)  + \beta^+ \log \left( \frac{\beta^+ p_{k_2} + \beta^- \bar{\gamma}_{k_2}}{\beta^+ p_{k_3} + \beta^- \bar{\gamma}_{k_3}} \right)
		\right]
		\bigg] .
	\end{align}
	\item If $ k_1 + 1 = k_2 $ and $ k_3  \neq \{k_1-1, k_2+1\}  $, the respective minimization problem becomes:
	\begin{equation}
		\begin{aligned}
			\min_{(p_{k_1},p_{k_2},p_{k_3}) \in  \field{R}^3} \quad & h\left(  \beta^+ \gamma_{k_1} + \beta^-  p_{k_1} \right) + h\left( \beta^+ p_{k_1} + \beta^-  p_{k_2} \right) + h\left( \beta^+ p_{k_2} + \beta^-  \bar{\gamma}_{k_2} \right)  \\
			&+ h\left( \beta^+ \gamma_{k_3} + \beta^-  p_{k_3} \right) +  h\left( \beta^+ p_{k_3} + \beta^-  \bar{\gamma}_{k_3} \right)\\
			\textrm{s.t.} \quad & -\sum_{i=1}^3 p_{k_i} \log p_{k_i}  = c_2\\
			& \sum_{i=1}^3 p_{k_i}  = c_1.
		\end{aligned}
	\end{equation}
	The directional derivative of the objective function is given by
	\begin{align}
		&-\left[\beta^- \log(\beta^+ \gamma_{k_1} + \beta^- p_{k_1}) + \beta^+ \log(\beta^+ p_{k_1} + \beta^- p_{k_2}) \right] \mathrm{d} p_{k_1} \\
		&-\left[\beta^- \log(\beta^+ p_{k_1} + \beta^- p_{k_2}) + \beta^+ \log(\beta^+ p_{k_2} + \beta^- \bar{\gamma}_{k_2}) \right] \mathrm{d} p_{k_2} \\
		&-\left[\beta^- \log(\beta^+ \gamma_{k_3} + \beta^- p_{k_3}) + \beta^+ \log(\beta^+ p_{k_3} + \beta^- \bar{\gamma}_{k_3}) \right] \mathrm{d} p_{k_3}  \\
		&=- \mathrm{d} p_{k_1} \bigg[\beta^- \log \left(\frac{\beta^+ \gamma_{k_1} + \beta^- p_{k_1}}{\beta^+ p_{k_1} + \beta^- p_{k_2}} \right)  + \beta^+ \log \left( \frac{\beta^+ p_{k_1} + \beta^- p_{k_2}}{\beta^+ p_{k_2} + \beta^- \bar{\gamma}_{k_2}} \right)  \\
		&-\frac{\log p_{k_1} - \log p_{k_2}}{\log p_{k_2} - \log p_{k_1}} \left[
		\beta^- \log \left(\frac{\beta^+ p_{k_1} + \beta^- p_{k_2}}{\beta^+ \gamma_{k_3} + \beta^- p_{k_3}} \right)  + \beta^+ \log \left( \frac{\beta^+ p_{k_2} + \beta^- \bar{\gamma}_{k_2}}{\beta^+ p_{k_3} + \beta^- \bar{\gamma}_{k_3}} \right)
		\right]
		\bigg]
	\end{align}
	\item As for the case $ k_1 + 1 = k_2 = k_3-1 $,  the minimization problem has the following form:
	\begin{equation}
		\begin{aligned}
			\min_{(p_{k_1},p_{k_2},p_{k_3}) \in  \field{R}^3} \quad & h\left(  \beta^+ \gamma_{k_1} + \beta^-  p_{k_1} \right) + h\left( \beta^+ p_{k_1} + \beta^-  p_{k_2} \right) + h\left( \beta^+ p_{k_2} + \beta^-  p_{k_3} \right)  + h\left( \beta^+ p_{k_3} + \beta^-  \bar{\gamma}_{k_3} \right)\\
			\textrm{s.t.} \quad & -\sum_{i=1}^3 p_{k_i} \log p_{k_i}  = c_2\\
			& \sum_{i=1}^3 p_{k_i}  = c_1.
		\end{aligned}
	\end{equation}
	The directional derivative of the objective function is given by
	\begin{align}
		&-\left[\beta^- \log(\beta^+ \gamma_{k_1} + \beta^- p_{k_1}) + \beta^+ \log(\beta^+ p_{k_1} + \beta^- p_{k_2}) \right] \mathrm{d} p_{k_1} \\
		&-\left[\beta^- \log(\beta^+ p_{k_1} + \beta^- p_{k_2}) + \beta^+ \log(\beta^+ p_{k_2} + \beta^- p_{k_3}) \right] \mathrm{d} p_{k_2} \\
		&-\left[\beta^- \log(\beta^+ p_{k_2} + \beta^- p_{k_3}) + \beta^+ \log(\beta^+ p_{k_3} + \beta^- \bar{\gamma}_{k_3}) \right] \mathrm{d} p_{k_3}  \\
		&=- \mathrm{d} p_{k_1} \bigg[\beta^- \log \left(\frac{\beta^+ \gamma_{k_1} + \beta^- p_{k_1}}{\beta^+ p_{k_1} + \beta^- p_{k_2}} \right)  + \beta^+ \log \left( \frac{\beta^+ p_{k_1} + \beta^- p_{k_2}}{\beta^+ p_{k_2} + \beta^- p_{k_3}} \right)  \\
		&-\frac{\log p_{k_1} - \log p_{k_2}}{\log p_{k_2} - \log p_{k_1}} \left[
		\beta^- \log \left(\frac{\beta^+ p_{k_1} + \beta^- p_{k_2}}{\beta^+ p_{k_2} + \beta^- p_{k_3}} \right)  + \beta^+ \log \left( \frac{\beta^+ p_{k_2} + \beta^- p_{k_3}}{\beta^+ p_{k_3} + \beta^- \bar{\gamma}_{k_3}} \right)
		\right]
		\bigg] .
	\end{align}
\end{itemize}

\subsection{Algorithm}
Now we propose a minimization algorithm.
\begin{enumerate}
	\item Draw a point $ \bv{p} \in \simplex{n} $  that satisfies $ h_n(\bv{p}) = \eta_2 $.  Choose an arbitrary triple $ (p_{k_1},p_{k_2},p_{k_3}) $ from $ \bv{p} $ and consider an arc in $ \simplex{3} $ satisfying $ p_{k_1} + p_{k_2} + p_{k_3} = c_1 $,  and $ \sum_{i=1}^3 p_{k_i} \log p_{k_i} = c_2 $. Parameterize the path along the arc with parameter $ t $, i.e., $ (p_{k_1}(t),p_{k_2}(t),p_{k_3}(t)) $, $ t \in [0,1] $, with $ (p_{k_1}(0),p_{k_2}(0),p_{k_3}(0)) $ being the negative Hamming "distribution" and $ (p_{k_1}(1),p_{k_2}(1),p_{k_3}(1)) $ the positive Hamming "distribution" (since $ c_1\leq 1 $ the triple $ (p_{k_1},p_{k_2},p_{k_3}) $  might not be a proper probability vector summing to one).
	\begin{itemize}
		\item If $ k_1+1=k_2=k_3-1 $, increase $ t $ as long as
		\begin{multline}
			  \beta^- \log \left(\frac{\beta^+ \gamma_{k_1} + \beta^- p_{k_1}}{\beta^+ p_{k_1} + \beta^- p_{k_2}} \right)  + \beta^+ \log \left( \frac{\beta^+ p_{k_1} + \beta^- p_{k_2}}{\beta^+ p_{k_2} + \beta^- p_{k_3}} \right)  \\
			  -\frac{\log p_{k_1} - \log p_{k_2}}{\log p_{k_2} - \log p_{k_1}} \left[
			  \beta^- \log \left(\frac{\beta^+ p_{k_1} + \beta^- p_{k_2}}{\beta^+ p_{k_2} + \beta^- p_{k_3}} \right)  + \beta^+ \log \left( \frac{\beta^+ p_{k_2} + \beta^- p_{k_3}}{\beta^+ p_{k_3} + \beta^- \bar{\gamma}_{k_3}} \right)
			  \right]
		\end{multline}
		is positive. If it is positive for every $ t $ then choose $ t=1 $ and $ (p^*_{k_1},p^*_{k_2},p^*_{k_3}) = (p_{k_1}(1),p_{k_2}(1),p_{k_3}(1)) $.
		\item If $ k_1+1=k_2 $, $ k_3 \notin \{k_1-1,k_2+1\} $ then increase $ t $ as long as
		\begin{multline}
			\beta^- \log \left(\frac{\beta^+ \gamma_{k_1} + \beta^- p_{k_1}}{\beta^+ p_{k_1} + \beta^- p_{k_2}} \right)  + \beta^+ \log \left( \frac{\beta^+ p_{k_1} + \beta^- p_{k_2}}{\beta^+ p_{k_2} + \beta^- \bar{\gamma}_{k_2}} \right)  \\
			 -\frac{\log p_{k_1} - \log p_{k_2}}{\log p_{k_2} - \log p_{k_1}} \left[
			\beta^- \log \left(\frac{\beta^+ p_{k_1} + \beta^- p_{k_2}}{\beta^+ \gamma_{k_3} + \beta^- p_{k_3}} \right)  + \beta^+ \log \left( \frac{\beta^+ p_{k_2} + \beta^- \bar{\gamma}_{k_2}}{\beta^+ p_{k_3} + \beta^- \bar{\gamma}_{k_3}} \right)
			\right]
		\end{multline}
		is positive. If it is positive for every $ t $ then choose $ t=1 $ and $ (p^*_{k_1},p^*_{k_2},p^*_{k_3}) = (p_{k_1}(1),p_{k_2}(1),p_{k_3}(1)) $.
		\item If $ k_1 < k_1+1 < k_2 < k_2+1< k_3 $ then increase $ t $ as long as
		\begin{multline}
			\beta^- \log \left(\frac{\beta^+ \gamma_{k_1} + \beta^- p_{k_1}}{\beta^+ \gamma_{k_2} + \beta^- p_{k_2}} \right)  + \beta^+ \log \left( \frac{\beta^+ p_{k_1} + \beta^- \bar{\gamma}_{k_1}}{\beta^+ p_{k_2} + \beta^- \bar{\gamma}_{k_2}} \right)  \\
			-\frac{\log p_{k_1} - \log p_{k_2}}{\log p_{k_2} - \log p_{k_1}} \left[
			\beta^- \log \left(\frac{\beta^+ \gamma_{k_2} + \beta^- p_{k_2}}{\beta^+ \gamma_{k_3} + \beta^- p_{k_3}} \right)  + \beta^+ \log \left( \frac{\beta^+ p_{k_2} + \beta^- \bar{\gamma}_{k_2}}{\beta^+ p_{k_3} + \beta^- \bar{\gamma}_{k_3}} \right)
			\right]
		\end{multline}
		is positive. If it is positive for every $ t $ then choose $ t=1 $ and $ (p^*_{k_1},p^*_{k_2},p^*_{k_3}) = (p_{k_1}(1),p_{k_2}(1),p_{k_3}(1)) $.
	\end{itemize}
	\item Let $ t^* $ be the optimal parameter on the arc. Replace the triple $ (p_{k_1},p_{k_2},p_{k_3}) $ with $ (p_{k_1}(t^*),p_{k_2}(t^*),p_{k_3}(t^*)) $.
	\item Go to step 1 until a stopping criterion is satisfied: either maximum number of iteration or no further decrease of the output entropy function is possible.
\end{enumerate}

\subsection{Lagrangian Approach}
Next we derive a method to solve \eqref{eq:oblib_low_snr_local_optimization} using Lagrangian multipliers and KKT method.
Similarly here, we are required to consider three different configurations for the the triple $ k_1,k_2,k_3 $: the case where they are consecutive; the case where two of them are adjacent; and the case where none of the are adjacent.
\begin{itemize}
	\item If $ k_1 < k_1+1 < k_2 < k_2+1< k_3 $, the minimization problem obtains the following form:
	\begin{equation}
		\begin{aligned}
			\min_{(p_{k_1},p_{k_2},p_{k_3}) \in  \field{R}^3} \quad & \sum_{i=1}^{3} h\left(  \beta^+ p_{k_i-1} + \beta^-  p_{k_i} \right) + h\left( \beta^+ p_{k_i} + \beta^-  p_{k_i+1} \right) \\
			\textrm{s.t.} \quad & -\sum_{i=1}^3 p_{k_i} \log p_{k_i} =  c_2\\
			& \sum_{i=1}^3 p_{k_i} = c_1.
		\end{aligned}
	\end{equation}

	The respective Lagrangian takes the form
	\begin{multline}
		L(p_{k_1},p_{k_2},p_{k_3},\lambda ,\mu) \\
		= \sum_{i=1}^{3} h\left(  \beta^+ p_{k_i-1} + \beta^-  p_{k_i} \right) + h\left( \beta^+ p_{k_i} + \beta^-  p_{k_i+1} \right)  + \lambda \left(\sum_{i=1}^3 p_{k_i} \log p_{k_i} + c_2\right) + \mu \left(\sum_{i=1}^3 p_{k_i}  - c_1\right).
	\end{multline}
	Assume that a local optimum of the problem is regular. Then the KKT conditions
	are necessary for a solution, and except for primal and dual feasibility
	state that:
	\begin{itemize}
		\item Stationarity (S):
		\begin{equation}
			\frac{\partial L}{\partial p_{k_i}}=
			- \beta^- \log (\beta^+  p_{k_i-1} + \beta^{-} p_{k_i})   - \beta^+ \log (\beta^+  p_{k_i} + \beta^{-} p_{k_i+1}) + \lambda + \lambda \log p_{k_i} + \mu = 0
		\end{equation}
		which is equivalent to
		\begin{equation}
			\frac{\left(\beta^+ \gamma + \beta^- p_{k_j} \right) ^{\beta^-} \left(\beta^+ p_{k_j} + \beta^- \bar{\gamma} \right) ^{\beta^+}}{p_{k_j}^{\lambda}} = \exp(\lambda + \mu)
			, \qquad \forall j \in [3].
		\end{equation}
		Since $(\lambda,\mu)\in\field{R}^2 $ we may replace $ \exp(\lambda + \mu) $ with  $\mu_{1}\triangleq\exp(\lambda + \mu)$. So we obtain
		\begin{equation} \label{eq:oblib_highsnr_noneighbors_stationarity}
			\frac{\left(\beta^+ \gamma + \beta^- p_{k_j} \right) ^{\beta^-} \left(\beta^+ p_{k_j} + \beta^- \bar{\gamma} \right) ^{\beta^+}}{p_{k_j}^{\lambda}} =\mu_1
			, \qquad \forall j \in [3].
		\end{equation}
		We further denote the LHS of \eqref{eq:oblib_highsnr_noneighbors_stationarity} as $ \varphi(p_{k_j},\lambda, \gamma_{k_j},\bar{\gamma}_{k_j}) $, i.e.,
		\begin{equation} \label{eq:psi_definiition}
			\varphi(p,\lambda ,\gamma , \bar{\gamma}) \triangleq \frac{\left(\beta^+ \gamma + \beta^- p \right) ^{\beta^-} \left(\beta^+ p + \beta^- \bar{\gamma} \right) ^{\beta^+}}{p^{\lambda}} .
		\end{equation}
		\item Complementary slackness (CS): $\sum_{i=1}^{3}p_{k_i}=c_1$, and $\sum_{i=1}^{3}p_{k_i} \log p_{k_i}=c_2$.
	\end{itemize}
	\item The case $ k_1 + 1 = k_2 $ and $ k_3  \neq \{k_1-1, k_2+1\}  $. In such case the minimization problem has the following form.
	\begin{equation}
		\begin{aligned}
			\min_{(p_{k_1},p_{k_2},p_{k_3}) \in  \field{R}^3} \quad & h\left(  \beta^+ \gamma_{k_1} + \beta^-  p_{k_1} \right)
			+ h\left( \beta^+ p_{k_1} + \beta^-  p_{k_2} \right)
			+ h\left( \beta^+ p_{k_2} + \beta^-  \bar{\gamma}_{k_2} \right)  \\
			&+ h\left( \beta^+ \gamma{k_3} + \beta^-  p_{k_3} \right)
			+  h\left( \beta^+ p_{k_3} + \beta^-  \bar{\gamma}_{k_3} \right)\\
			\textrm{s.t.} \quad & -\sum_{i=1}^3 p_{k_i} \log p_{k_i}  = c_2\\
			& \sum_{i=1}^3 p_{k_i}  = c_1.
		\end{aligned}
	\end{equation}
	\begin{itemize}
		\item Stationarity (S):
		\begin{align}
			\frac{\partial L}{\partial p_{k_1}} &=
			- \beta^- \log (\beta^+  \gamma_{k_1} + \beta^{-} p_{k_1})   - \beta^+ \log (\beta^+  p_{k_1} + \beta^{-} p_{k_2}) + \lambda + \lambda \log p_{k_1} + \mu = 0 \\
			\frac{\partial L}{\partial p_{k_2}} &=
			- \beta^- \log (\beta^+  p_{k_1} + \beta^{-} p_{k_2})   - \beta^+ \log (\beta^+  p_{k_2} + \beta^{-} \gamma_{k_2}) + \lambda + \lambda \log p_{k_2} + \mu = 0 \\
			\frac{\partial L}{\partial p_{k_3}} &=
			- \beta^- \log (\beta^+  \gamma_{k_3} + \beta^{-} p_{k_3})   - \beta^+ \log (\beta^+  p_{k_3} + \beta^{-} \gamma_{k_3}) + \lambda + \lambda \log p_{k_3} + \mu = 0
		\end{align}
		which is equivalent to
		\begin{equation}
			\frac{\left(\beta^+ \gamma_{k_1} + \beta^- p_{k_1} \right) ^{\beta^-} \left(\beta^+ p_{k_1} + \beta^- p_{k_2} \right) ^{\beta^+}}{p_{k_1}^{\lambda}} =
			\frac{\left(\beta^+ p_{k_1} + \beta^- p_{k_2} \right) ^{\beta^-} \left(\beta^+ p_{k_2} + \beta^- \bar{\gamma}_{k_2} \right) ^{\beta^+}}{p_{k_2}^{\lambda}}
			=\mu_1
		\end{equation}
		and
		\begin{equation}
			\frac{\left(\beta^+ \gamma_{k_3} + \beta^- p_{k_3} \right) ^{\beta^-} \left(\beta^+ p_{k_3} + \beta^- \bar{\gamma}_{k_3} \right) ^{\beta^+}}{p_{k_3}^{\lambda}} = \mu_1.
		\end{equation}
		\item Complementary slackness (CS): $\sum_{i=1}^{3}p_{k_i}=c_1$, and $\sum_{i=1}^{3}p_{k_i} \log p_{k_i}=c_2$.
	\end{itemize}
	\item In the case $ k_1 + 1 = k_2 = k_3-1 $  the minimization problem takes the following form.
	\begin{equation}
		\begin{aligned}
			\min_{(p_{k_1},p_{k_2},p_{k_3}) \in  \field{R}^3} \quad & h\left(  \beta^+ \gamma_{k_1} + \beta^-  p_{k_1} \right)
			+ h\left( \beta^+ p_{k_1} + \beta^-  p_{k_2} \right)
			+ h\left( \beta^+ p_{k_2} + \beta^-  p_{k_3} \right)
			+ h\left( \beta^+ p_{k_3} + \beta^-  \bar{\gamma}_{k_3} \right)\\
			\textrm{s.t.} \quad & -\sum_{i=1}^3 p_{k_i} \log p_{k_i}  = c_2\\
			& \sum_{i=1}^3 p_{k_i}  = c_1.
		\end{aligned}
	\end{equation}
	\begin{itemize}
		\item Stationarity (S):
		\begin{align}
			\frac{\partial L}{\partial p_{k_1}} &=
			- \beta^- \log (\beta^+  \gamma_{k_1} + \beta^{-} p_{k_1})   - \beta^+ \log (\beta^+  p_{k_1} + \beta^{-} p_{k_2}) + \lambda + \lambda \log p_{k_1} + \mu = 0 \\
			\frac{\partial L}{\partial p_{k_2}} &=
			- \beta^- \log (\beta^+  p_{k_1} + \beta^{-} p_{k_2})   - \beta^+ \log (\beta^+  p_{k_2} + \beta^{-} p_{k_3}) + \lambda + \lambda \log p_{k_2} + \mu = 0 \\
			\frac{\partial L}{\partial p_{k_3}} &=
			- \beta^- \log (\beta^+  p_{k_2} + \beta^{-} p_{k_3})   - \beta^+ \log (\beta^+  p_{k_3} + \beta^{-} \gamma_{k_3}) + \lambda + \lambda \log p_{k_3} + \mu = 0
		\end{align}
		which is equivalent to
		\begin{multline}
			\frac{\left(\beta^+ \gamma_{k_1} + \beta^- p_{k_1} \right) ^{\beta^-} \left(\beta^+ p_{k_1} + \beta^- p_{k_2} \right) ^{\beta^+}}{p_{k_1}^{\lambda}} \\
			=
			\frac{\left(\beta^+ p_{k_1} + \beta^- p_{k_2} \right) ^{\beta^-} \left(\beta^+ p_{k_2} + \beta^- p_{k_3} \right) ^{\beta^+}}{p_{k_2}^{\lambda}}
			=\frac{\left(\beta^+ p_{k_2} + \beta^- p_{k_3} \right) ^{\beta^-} \left(\beta^+ p_{k_3} + \beta^- \bar{\gamma}_{k_3} \right) ^{\beta^+}}{p_{k_3}^{\lambda}} = \mu_1.
		\end{multline}
		\item Complementary slackness (CS): $\sum_{i=1}^{3}p_{k_i}=c_1$, and $\sum_{i=1}^{3}p_{k_i} \log p_{k_i}=c_2$.
	\end{itemize}
\end{itemize}

\subsection{Algorithm}
Now we propose a minimization algorithm.
\begin{enumerate}
	\item Draw a point $ \bv{p} \in \simplex{n} $  that satisfies $ h_n(\bv{p}) = \eta_2 $.  Choose an arbitrary triple $ (p_{k_1},p_{k_2},p_{k_3}) $ from $ \bv{p} $ and consider an arc in $ \simplex{3} $ satisfying $ p_{k_1} + p_{k_2} + p_{k_3} = c_1 $,  and $ \sum_{i=1}^3 p_{k_i} \log p_{k_i} = c_2 $. Parameterize the path along the arc with parameter $ t $, i.e., $ (p_{k_1}(t),p_{k_2}(t),p_{k_3}(t)) $, $ t \in [0,1] $, , with $ (p_{k_1}(0),p_{k_2}(0),p_{k_3}(0)) $ being the negative Hamming "distribution" and $ (p_{k_1}(1),p_{k_2}(1),p_{k_3}(1)) $ the positive Hamming "distribution" (since $ c_1\leq 1 $ the triple $ (p_{k_1},p_{k_2},p_{k_3}) $  might not be a proper probability vector summing to one).
	\begin{itemize}
		\item If $ k_1+1=k_2=k_3-1 $, then find $ t \in [0,1] $ and $ \lambda \in \field{R} $ that satisfy
		\begin{multline}
			\frac{\left(\beta^+ \gamma_{k_1} + \beta^- p_{k_1}(t) \right) ^{\beta^-} \left(\beta^+ p_{k_1}(t)+ \beta^- p_{k_2}(t) \right) ^{\beta^+}}{p_{k_1}^{\lambda}(t)} \\
			=
			\frac{\left(\beta^+ p_{k_1} (t)+ \beta^- p_{k_2}(t) \right) ^{\beta^-} \left(\beta^+ p_{k_2} (t)+ \beta^- p_{k_3}(t) \right) ^{\beta^+}}{p_{k_2}^{\lambda}(t)}
			=\frac{\left(\beta^+ p_{k_2}(t) + \beta^- p_{k_3} (t) \right) ^{\beta^-} \left(\beta^+ p_{k_3} (t)+ \beta^- \bar{\gamma}_{k_3} \right) ^{\beta^+}}{p_{k_3}^{\lambda}(t)} .
		\end{multline}
		\item If $ k_1+1=k_2 $, $ k_3 \notin \{k_1-1,k_2+1\} $ then find $ t \in [0,1] $ and $ \lambda \in \field{R} $ that satisfy
		\begin{multline}
			\frac{\left(\beta^+ \gamma_{k_1} + \beta^- p_{k_1} \right) ^{\beta^-} \left(\beta^+ p_{k_1} + \beta^- p_{k_2} \right) ^{\beta^+}}{p_{k_1}^{\lambda}} \\
			=
			\frac{\left(\beta^+ p_{k_1} + \beta^- p_{k_2} \right) ^{\beta^-} \left(\beta^+ p_{k_2} + \beta^- \bar{\gamma}_{k_2} \right) ^{\beta^+}}{p_{k_2}^{\lambda}}
			= \frac{\left(\beta^+ \gamma_{k_3} + \beta^- p_{k_3} \right) ^{\beta^-} \left(\beta^+ p_{k_3} + \beta^- \bar{\gamma}_{k_3} \right) ^{\beta^+}}{p_{k_3}^{\lambda}}
		\end{multline}
		\item If $ k_1 < k_1+1 < k_2 < k_2+1< k_3 $ then find $ t \in [0,1] $ and $ \lambda \in \field{R} $ that satisfy
		\begin{equation}
			\varphi(p_{k_1}(t),\lambda ,\gamma _{k_1}, \bar{\gamma}_{k_1}) = \varphi(p_{k_2}(t),\lambda ,\gamma _{k_1}, \bar{\gamma}_{k_1}) = \varphi(p_{k_3}(t),\lambda ,\gamma _{k_1}, \bar{\gamma}_{k_1}) .
		\end{equation}
	\end{itemize}
\end{enumerate}

\subsection{Scalar Gaussian $ (\rv{X},\rv{Y}) $}
    \input{./sections/gaussianXY}

    \subsection{Vector Gaussian $ (\bv{X},\bv{Y}) $}
    \input{./sections/VectorGaussianXY}

\subsection{Additive Channels with KL-divergence Constraint}
\input{./sections/kl_constraints}

\subsection{Stationary Points Characterization}
	\begin{algorithm} [t]
		\SetAlgoLined
		\SetKwInput{Init}{Initialize}
		\SetKwInput{Cmp}{Compute}
		\SetKwInput{Set}{Set}
		\KwIn{$ \mathsf{P}_{\rv{X} \rv{Y}} $,  and $ \beta_2 $}
		\Init{Arbitrary $ \mathsf{P}_{\rv{Z}|\rv{Y}}^{(0)}  $, $ s=1 $.}
		\While{Variation in $ I(\rv{X};\rv{Z}) $ is greater then $ \epsilon $}{
			$ \mathsf{P}_{\rv{Z}|\rv{Y}}^{(s)} (z|y)
			= \frac{\mathsf{P}_{\rv{Z}}^{(s - 1)} (z)}{Z(y,\beta_2)}  \cdot
			e^{ \scalebox{0.75}[1.0]{\( - \)} \beta_2   D\left( \mathsf{P}_{\rv{X}|\rv{Y}} (\cdot|y) || \mathsf{P}_{\rv{X}|\rv{Z}}^{(s-1)} (\cdot|z) \right) }$\;
			 $ \mathsf{P}_{\rv{Z}}^{(s)} (z) =  \sum_{y \in \set{Y}} \mathsf{P}_{\rv{Y}} (y) \mathsf{P}_{\rv{Z}|\rv{Y}}^{(s-1)} (z|y) $\;
			 $ \mathsf{P}_{\rv{X}|\rv{Z}}^{(s)} (x|z) =  \sum_{y \in \set{Y}} \mathsf{P}_{\rv{X}|\rv{Y}} (x|y) \mathsf{P}_{\rv{Y}|\rv{Z}}^{(s)} (y|z) $\;
			$ s = s+1 $ \;
		}
		\KwOut{$\mathsf{P}^*_{\rv{Z}|\rv{Y}} $}
		\caption{ib\_iterator(args)}
		\label{algorithm:oblib_maximization}
	\end{algorithm}


\section[Introduction]{Introduction}

A surge of interest in building deep learning-based data-driven models for complex systems such as chaotic dynamical systems \cite{pathak2018model,chattopadhyay2020data}, fully turbulent flow \cite{chattopadhyay2020deep}, and weather and climate models \cite{weyn2019can,weyn2020improving,weyn2021sub,rasp_2020_resnet,chattopadhyay2021towards} has been seen in the recent past. This interest in data-driven modeling stems from the hope that if these data-driven models are trained on observational data, (i) they would not suffer from some of the biases of physics-based numerical climate and weather models \citep{balaji2021climbing,schultz2021can}, e.g., due to ad-hoc parameterizations, (ii) they can be used to generate large ensemble forecasts for data assimilation \cite{chattopadhyay2021towards}, and (iii) they might seamlessly be integrated to perform climate simulations \cite{watson2021machine} that would allow for generating large synthetic datasets to study extreme events \citep{chattopadhyay2019analog}.

Despite the promise of data-driven weather prediction  (DDWP) models, most DDWP models cannot compete with state-of-the-art numerical weather prediction (NWP) models \citep{scher2019weather,rasp_2020_resnet,weyn2019can,weyn2020improving,chattopadhyay2021towards} although more recently FourCastNet \cite{pathak2022fourcastnet} has shown promise of being quite competitive with the best NWP models. There have also been some previous work, where pre-training models on climate simulations have also resulted in improved short-term forecasts on reanalysis datasets \cite{rasp_2020_resnet}. Moreover, these DDWP models, while comparable with NWP in terms of short-term forecasts, are unstable when integrated for longer time scales. These instabilities either appear in the form of blow-ups or unphysical climate drifts \citep{scher2019weather,chattopadhyay2020deep}.

In this paper, we consider the following realistic problem set-up, wherein, we have long-term climate simulations from an imperfect 2-layer quasi-geostrophic (QG) flow (we will call it "imperfect system" from here on) at our disposal to pre-train a DDWP. Following that, we have a few noisy observations from a perfect 2-layer QG flow ("perfect system" from here on) on which we are able to fine-tune our DDWP. This observationally-constrained DDWP model is then expected to generalize to the perfect system with a noisy initial condition (obtained from the available observations) and perform stochastic short-term forecasting. It is also expected to be seamlessly time-integrated to obtain long-term climatology. We show that in this regard, a stochastic forecasting approach with a generative model such as a convolutional variational autoencoder (VAE) outperforms a deterministic encoder-decoder-based convolutional architecture in terms of short-term forecasts and remains stable without unphysical climate drift for long-term integration. A deterministic encoder-decoder would fail to do so and shows unphysical climate drift followed by numerical blow-up.

\section{Methods}

\subsection{Imperfect and Perfect Systems}

In this paper, we consider 2-layer QG flow as the geophysical system on which we intend to perform both short-term forecasting and compute long-term statistics.

The dimensionless dynamical equations of the 2-layer QG flow are derived following Lutsko et al.~\cite{lutsko2015applying}, and Nabizadeh et al.~\cite{nabizadeh2019size}. The system consists of two constant density layers with a $\beta$-plane approximation in which the meridional temperature gradient is relaxed towards an equilibrium profile. The equation of the system is described as:

\begin{align}
\label{eq_QG}
\begin{split}
    \frac{\partial q_j}{\partial t}+J(\psi_j,q_j)=-\frac{1}{\tau_d}(-1)^{j}\left(\psi_1-\psi_2-\psi_R\right) \\
    -\frac{1}{\tau_f}\delta_{k2}\nabla^2\psi_j - \nu \nabla^{8}q_j.
       \end{split}
\end{align}
Here, $q$ is potential vorticity and is expressed as
\begin{align}
\label{q_def}
q_j=\nabla^2\psi_j +(-1)^{j}\left(\psi_1-\psi_2\right)+\beta y,
\end{align}
where $\psi_j$ is the stream function of the system.
In Eqs.~(\ref{eq_QG}) and~(\ref{q_def}), $j$ denotes the upper ($j=1$) and lower ($j=2$) layers.~$\tau_d$ is the Newtonian relaxation time scale while $\tau_f$ is the Rayleigh friction time scale, which only acts on the lower layers represented by the Kronecker $\delta$ function ($\delta_{k2}$).~$J$ denotes the jacobian, $\beta$ is the $y-$gradient of the coriolis parameter, and $\nu$ denotes the hyperdiffusion coefficient. We have induced a baroclinically unstable jet at the center of a horizontally (zonally) periodic channel by setting $\psi_1-\psi_2$ to be equal to a hyperbolic secant centered at $y=0$. When eddy fluxes are absent, $\psi_2$ is identically zero, making zonal velocity in the upper layer $u(y)=-\frac{\partial \psi_1}{\partial y}=-\frac{\partial \psi_R}{\partial y}$ where we set
\begin{align}
\label{init_eq}
-\frac{\partial \psi_R}{\partial y}=\sech^2\left(\frac{y}{\sigma}\right),
\end{align}
$\sigma$ being the width of the jet. Parameters of the model are set following the previous studies~\cite{lutsko2015applying,nabizadeh2019size}, in which $\beta= 0.19$, $\sigma=3.5$, $\tau_f=15$, and $\tau_d=100$.

To non-dimensionalize the equations, we have used the maximum strength of the equilibrium velocity profile as the velocity scale ($U$) and the deformation radius ($L$) for the length scale. The system’s time scale ($L/U$) is referred to as the “advective time scale” ($\tau_{adv}$).

The spatial discretization is spectral in both $x$ and $y$ where we have retained $96$ and $192$ Fourier modes, respectively. The length and width of the domain is equal to $46$ and $68$ respectively after non-dimensionalizing the numbers. The spurious waves on the northern and southern boundaries are damped by applying sponge layers. Note that the domain is wide enough that the sponges do not affect the dynamics. Here $5 \tau_{adv} \approx$ $1$ Earth day $\approx 200 \Delta t_n$, where $\Delta t_n=0.025$ is the time step of the leapfrog time integrator used in the numerical scheme.

The imperfect system in this paper has an increased value of $\beta$, given by $\beta^* = 3\beta$ and a decreased size of the jet, given by $\sigma*=\frac{4}{5}\sigma$. We assume that we have long-term climate simulations of the imperfect system while we have only a few noisy observations of the perfect system. The difference in the long-term averaged zonal-mean velocity in the upper layer of the perfect and imperfect systems is shown in Fig.~\ref{2-system}. It is to be noted that the motivation behind having an imperfect and perfect system is the fact that we assume that the \textit{imperfect system} represents a climate model with biases from which we have long simulations at our disposal, while the \textit{perfect system} is analogous to actual observations from our atmosphere, where the sample size is small and contaminated with measurement noise.

\begin{figure*}[t]
\includegraphics[width=\textwidth]{Picture1.png}
\caption{Time-averaged zonal-mean velocity (over $20000$ days), $\left<\bar{u}\right>$, of the imperfect and perfect system. The difference between $\left<\bar{u}\right>$ of the perfect and imperfect systems indicates a challenge for DDWP models to seamlessly generalize from one system to the other}
\label{2-system}
\end{figure*}

\subsection{VAE and Transfer Learning Frameworks}
Here, we propose a generative modeling approach to perform both short- and long-term forecasting for the perfect QG system after being trained on the imperfect QG system. We assume that we have a long-term climate simulation from the imperfect system, with states $x^m(k\Delta t)$, where $k\in{1,2,\cdots K}$ and $x^m \in \mathcal{R}^{2\times 96\times 192}$. A convolutional VAE is trained on the states of this imperfect system to predict an ensemble of states with a probability density function, $p\left(x^m\left(t+\Delta t\right)\right)$, from $x^m(t)$. Here, $\Delta t = 40 \Delta t_n$, which is also the sampling interval for training the model. Extensive trial and error-based search has been performed for hyperparameter optimization for the VAE. The details on the VAE architecture is given in Table~\ref{tab:VAE}.

This pre-trained VAE would then undergo transfer learning on a small number of noisy observations of the states of the perfect system $x^o(n\Delta t)$ where $n\in{1,2,\cdots N}$ and $N \ll K$. The VAE would then stochastically forecast an ensemble of states of the perfect system with a noisy initial condition from the perfect system. The noise in the initial condition is sampled from a Gaussian normal distribution with $0$ mean and standard deviation being $\eta \sigma_Z$, where $\sigma_Z$ is the standard deviation of $\psi_k$ and $\eta$ is a fraction that determines the amplitude of the noise vector. A schematic of this framework is shown in Fig.~\ref{framework}. In this work, we consider the stream function $\psi_k$ ($k=1,2$) to be the states of the system on which we would train our VAE.

The VAE is trained on $9$ independent ensembles of the imperfect system with $1400$ consecutive days of climate simulation each with a sampling interval of $\Delta t$. Hence, the training size for the VAE is about $12600$ days of data. For transfer learning, we assume that only $10\%$ of the number of training samples is available as noisy observations from the perfect system.

\begin{table}[]
    \centering
    \begin{tabular}{|l||l|p{0.1\textwidth}|}
    \hline
          \textbf{Number}& \textbf{Layer} & \textbf{Number of Filters} \\
         \hline
         1 & $5\times 5$ 2D Convolution & 32 \\
         \hline
         \hline
         2 & $2\times 2$ Max Pooling & --\\
         \hline
         3 & $5\times 5$ 2D Convolution & 32 \\
         \hline
         \hline
         4 & $2\times 2$ Max Pooling & -- \\
         \hline
         5 & $5\times 5$ 2D Convolution & 32 \\
         \hline
          6 & $2\times 2$ Max Pooling & -- \\
         \hline
         7 & flatten & --  \\
         \hline
         8 & 128 neurons for mean and standard deviation & -- \\
         \hline
         9 & Dense layer & --\\
         \hline
         10 & $5\times 5$ 2D Convolution & 32\\
          \hline
          11 & Up-sampling & --\\
          \hline
         12 & $5\times 5$ 2D Convolution & 32 \\
          \hline
         13 &  Up-sampling & --\\
          \hline
        14 & $5\times 5$ 2D Convolution & 32 \\
          \hline
         15 &  Up-sampling & --\\
         \hline
                 15 & $5\times 5$ 2D Convolution & 32 \\
          \hline
         16 &  Up-sampling & --\\
         \hline
         17 & $5\times 5$ 2D Convolution & 2 \\
         \hline
    \end{tabular}
    \caption{Number of layers and filters in the convolutional VAE architecture used in this paper. The baseline convolutional encoder-decoder model would have the exact same architecture and the same latent space size as that of the VAE. }
    \label{tab:VAE}
\end{table}

\begin{figure*}
\includegraphics[width=\textwidth]{Picture2.png}
\caption{A schematic of the transfer learning framework with a VAE pre-trained on imperfect simulations and transfer learned on noisy observations from the perfect system. Here, $x^m\left(k\Delta t\right)$ are states obtained from the imperfect system and $x^o\left(n\Delta t\right)$ are noisy observations from the perfect system. Here, $\Delta t = 40 \Delta t_n$. Note that, our proposed VAE is convolutional and the schematic is just representative. Details about the architecture are given in Table~\ref{tab:VAE}}
\label{framework}
\end{figure*}

\section{Results}

\subsection{Short-term Stochastic Forecasting on the Imperfect System}
\begin{figure*}[t]
\includegraphics[width=\textwidth]{ACC_RMSE_noisyIC_perfect.png}
\caption{Performance of the stochastic convolutional VAE as compared to baseline convolutional encoder-decoder-based model when trained on the imperfect system and predicts from a noisy initial condition sampled from the imperfect system.
(a) Anomaly correlation coefficient (ACC) between predicted $\psi_1$ and true $\psi_1$ \citep{murphy1989skill} from the imperfect system with $\eta = 5\%$ (noise level) for the initial condition. (b) Same as (a) but for RMSE. (c) Prediction horizon (number of $\Delta t_n$ until ACC $\leq 0.60$) for different noise levels added to the initial condition. (d) Averaged error over $2$ days for different noise levels of initial condition. The shading shows the standard deviation across $100$ ensembles generated by the VAE model during inference}
\label{forecast_imperfect}
\end{figure*}
In this section, we show how well the stochastic VAE performs in terms of short-term forecasting on the imperfect system when the initial condition is sampled from the imperfect system. Our baseline model is a convolutional encoder-decoder-based deterministic model that has the exact same architecture as that of the convolutional VAE with the same size of the latent space. During inference, the stochastic VAE generates $100$ ensembles at every $\Delta t$ and the mean of these ensembles is fed back into the VAE autoregressively to predict future time steps. We see from Fig.~\ref{forecast_imperfect} that the stochastic VAE model outperforms the deterministic model. Our experiments suggest that increasing the noise level of the initial condition does not affect the prediction horizon of VAE if we consider ACC $= 0.60$ as the limit of prediction (although it is an ad-hoc choice) as shown in Fig.~\ref{forecast_imperfect}(c). This maybe attributed to the reduction of initial condition error (due to noise) with ensembling. The averaged RMSE error over $2$ days grows faster with an increase in the noise level in the deterministic model as compared to the stochastic VAE model as shown in Fig.~\ref{forecast_imperfect}(d).

\subsection{Short-term Stochastic Forecasting on the Perfect System}
In this section, we show how well the pre-trained VAE (as compared to the baseline encoder-decoder and imperfect numerical model) on the imperfect climate simulations performs on the perfect system when constrained by noisy observations from the perfect system on which it is re-trained. Only two layers before and after the bottleneck layer are re-trained for both the VAE and the baseline model. Figure~\ref{forecasts_perfect} (a) and (b) show that for short-term forecasting, VAE outperforms both the baseline deterministic encoder-decoder model as well as the imperfect numerical model. Figure~\ref{forecasts_perfect}(c) shows a similar behaviour for VAE with increase in noise level in the initial condition as Fig.~\ref{forecast_imperfect}(c). Figure~\ref{forecasts_perfect}(d) shows the effect of the number of re-training samples used as noisy observations for transfer learning on the prediction horizon. While the numerical model is unaffected by the impact of noisy observations (since it does not undergo any learning), an increase in noise level to the initial condition makes it more susceptible to error as compared to the data-driven models which are inherently more robust to noise (due to re-training on noisy observations). In all cases, the VAE outperforms both the deterministic encoder-decoder model and the imperfect numerical model for short-term forecasts.

\begin{figure*}
\includegraphics[width=\textwidth]{ACC_RMSE_TL_scaling.png}
\caption{Performance of the stochastic convolutional VAE compared to that of the baseline convolutional encoder-decoder-based model and the imperfect numerical model when trained on the imperfect system, re-trained on noisy observations from the perfect system, and initialized with a noisy initial condition sampled from the perfect system.
(a) Anomaly correlation coefficient (ACC) between predicted $\psi_1$ and true $\psi_1$ from the perfect system with $\eta = 5\%$ (noise level) for the initial condition. (b) Same as (a) but for RMSE. (c) Prediction horizon (number of $\Delta t_n$ until ACC $\leq 0.60$) for different noise levels added to initial condition. (d) Prediction horizon of the models with different sample sizes of noisy observation data (as percentages of original training sample size) from perfect system. Note that the numerical model is not trained on data and hence would show no effect.   }
\label{forecasts_perfect}
\end{figure*}

\subsection{Long-term Climate Statistics}
In this section, we show the comparison between long-term climatology obtained from seamlessly integrating the VAE for 20000 days and the true climatology of the imperfect system. It must be noted that the deterministic encoder-decoder model is not stable and would show unphysical drifts in the climate similar to other deterministic data-driven models within a few days of seamless integration \citep{chattopadhyay2020deep}. Figure.~\ref{long_term_forecasts} shows stable and non-drifting physical climate obtained from the VAE whose mean of $\psi_j$, $u_j$, and EOF1 of $u_1$ closely match those of the true system.

\begin{figure*}[ht]
\includegraphics[width=\textwidth]{long_run_2000snapshot.png}
\caption{Long-term climatology obtained from integrating the models for $20000$ days. $\left<\bar{\psi}_1\right>$ is the time mean zonally-averaged $\psi_1$, $\left<\bar{u}\right>$ is the time mean zonally averaged upper-layer velocity, and EOF1, is the first EOF (empiricial orthogonal function) of $\left<\bar{u}\right>$. The VAE shows non-drifting physical climate. Baseline encoder-decoder model is not shown since it becomes unstable within $100$ days of seamless integration. }
\label{long_term_forecasts}
\end{figure*}

\section{Conclusion}
In this paper, we have developed a convolutional VAE for stochastic short-term forecasting and long-term climatology of the stream function of a 2-layer QG flow. The VAE is trained on climate simulations from an imperfect QG system before being fine-tuned on noisy observations from a perfect QG system on which it performs both short- and long-term forecasting. The VAE outperforms both the baseline convolutional encoder-decoder model as well as the imperfect numerical model.

One of the main advantages of using the stochastic VAE with multiple ensemble members is the reduction of the effect of noise in the initial condition for short-term forecasting. Moreover, the VAE remains stable through a seamless $20000$ days integration yielding a physical climatology that matches the numerical model. Deterministic data-driven models do not remain stable for long-term integration and hence a stochastic generative modeling approach may be a potential candidate when developing data-driven weather forecasting models that can seamlessly be integrated to yield long-term climate simulations.

Despite the improvement in performance in short-term skills as well as long-term stability, one of the caveats of the VAE is that it is less interpretable as compared to a deterministic model. Moreover, the VAE is more difficult to scale as compared to a deterministic model, especially for high-dimensional systems. This is because one would need to evolve a large number of ensembles for high-dimensional systems, and the exact relationship between short-term skills and long-term stability with the number of ensembles is not well known. Finally, while the application of VAE has resulted in long-term stability, the underlying causal mechanism by which instability is exhibited in a determinisitc data-driven model is still largely unknown. Future work needs to be undertaken to understand the cause of this instability such that determinisitc  models that may be more scalable than VAE can also be used for data-driven weather and climate prediction.

\paragraph{Acknowledgments}
This work was started as an internship project by A.C. in Lawrence Berkeley National Laboratory in the Summer of 2021 under J.P. and W.B. It was then continued under P.H. as a part of A.C.'s PhD at Rice University. This research used resources of NERSC, a U.S. Department of Energy Office of Science User Facility operated under Contract No. DE-AC02-05CH11231.  A.C., E.N., and P.H. were supported by ONR grant N00014-20-1-2722, NASA grant 80NSSC17K0266, and NSF CSSI grant OAC-2005123. A.C. also thanks the Rice University Ken Kennedy
Institute for a BP HPC Graduate Fellowship. We are grateful to Kamyar Azizzadenesheli, Mustafa Mustafa, and Sanjeev Raja for providing invaluable advice for this work.

\paragraph{Competing Interests}
None

\paragraph{Data Availability Statement}
The codes for this work would be available to the public upon publication in this Github link: \url{https://github.com/ashesh6810/Stochastic-VAE-for-Digital-Twins.git}.

\paragraph{Ethical Standards}
The research meets all ethical guidelines, including adherence to the legal requirements of the study country.

\bibliographystyle{unsrtnat}