Year: 2021

Abstract:
Deep Boltzmann machines refer to deep multi-layered probabilistic models, governed by a pairwise energy function that describes the likelihood of all variables in the network. Due to the difficulty of inference in such systems, they have given way largely to restricted deep Boltzmann machines (which do not permit intralayer or skip connections). In this paper, we propose a class of model that allows for exact, efficient mean-field inference and learning in general deep Boltzmann machines. To do so, we use the tools of the recently proposed monotone Deep Equilibrium (DEQ) Model, an implicit-depth deep network that always guarantees the existence and uniqueness of its fixed points. We show that, for a class of general deep Boltzmann machine, the mean-field fixed point can be considered as the equivalent fixed point of a monotone DEQ, which gives us a recipe for deriving an efficient mean-field inference procedure with global convergence guarantees. We apply this approach to simple deep convolutional Boltzmann architectures and demonstrate that it allows for tasks such as the joint completion and classification of images, all within a single deep probabilistic setting.

heading: 1 INTRODUCTION
text: This paper considers (deep) Boltzmann machines, which are pairwise energy-based probabilistic models. Theses models specify a joint distribution over variables x given by the density
p(x) ∝ exp  ∑ (i,j)∈E x⊤i Φijxj + n∑ i=1 b⊤i xi  , (1) where each x1:n denotes a discrete random variable over ki possible values, represented as a onehot encoding xi ∈ {0, 1}ki ; E denotes the set of edges in the model; Φi,j ∈ Rki×kj represents pairwise potential; and bi ∈ Rki represents unary potential. Depending on context, these models are typically referred to as pairwise Markov random fields (MRFs) (Koller & Friedman, 2009), or (potentially deep) Boltzmann machines (Goodfellow et al., 2016; Salakhutdinov & Hinton, 2009; Hinton, 2002). In the above setting each xi may represent an observed or unobserved value, and there can be substantial structure within the variables; for instance, the collection of variables x may (and indeed will, in the main settings we consider in this paper) consist of several different “layers” in a joint convolutional structure, leading to the deep convolutional Boltzmann machine (Norouzi et al., 2009).
In this paper, we propose a new parameterization and algorithmic approach to approximate inference of these probabilistic models. There are two main contributions: First, we define a generic parameterization of the pairwise kernel function Φ, that can represent a general Boltzmann machine. Our parametrization is flexible enough to incorporate almost all operators and network topology, including fully-connected layers, convolution operators, and skip-connections, etc. Through the lens of the recently developed monotone DEQ (Winston & Kolter, 2020), we constraint Φ to satisfy certain monotonicity conditions through training and inference; this thereby assures that the mean-field approximation will always have a unique, globally-optimal fixed point under this parameterization. Second, although previous works (Krähenbühl & Koltun, 2013; Baqué et al., 2016) have made approaches to parallel mean-field updates, they either require strong conditions on ϕ, or fail to converge to the true mean-field distribution. We provide a properly-damped mean-field update method, based upon a generic proximal operator, which is guaranteed to converge to the mean-field fixed point, even if applied in parallel to all random variables simultaneously. Although there is no exact
closed form solution of this proximal operator, we derive a very efficient Newton-based implementation. To estimate the parameters of our model, we follow the marginal-based loss minimization approach in Krähenbühl & Koltun (2013); Domke (2013), where the objective is to directly maximize the likelihood induced by the approximated mean-field distribution. Such an approach grants us the ability to update parameters via taking gradient steps of the proposed mean-field iterations.
With the proposed approaches, we perform both learning and inference for a deep convolutional, multi-resolution Boltzmann machine, and apply the network to model MNIST and CIFAR-10 pixels and their classes conditioned on partially observed images. Such joint probabilistic modelling allows us to simultaneously impute missing pixels and predict the class. While these are naturally a smallscale problem, we emphasize that performing joint probabilistic inference over a complete model of this type is a relatively high-dimensional task as far as traditional mean-field inference is concerned. We also compare our inference method to previous ones and demonstrate different convergence properties of each. We conclude by highlighting limitations and directions for future work with the method.

heading: 2 BACKGROUND AND RELATED WORK
text: This paper builds upon three main avenues of work: 1) deep equilibrium models, especially their convergent version, the monotone DEQ; 2) the broad topic of energy-based deep model and Boltzmann machines in particular; and 3) work on concave potentials and parallel methods for mean-field inference. We discuss each of these below.
Equilibrium models and their provable convergence The DEQ model was first proposed by Bai et al. (2019). Based on the observation that a neural network zt+1 = σ(Wzt + Ux+ b) with input injection x usually converges to a fixed point, they modeled an effectively infinite-depth network with input injection directly via its fixed point: z∗ = σ(Wz∗ + Ux + b). Its backpropagation is done through the implicit function theorem and only requires constant memory. Bai et al. (2020) also showed that the multiscale DEQ models achieve near state-of-the-art performances on many large-scale tasks. Winston & Kolter (2020) later presented a parametrization of the DEQ (denoted as monDEQ) that guarantees provable convergence to a unique fixed point, using monotone operator theory. Specifically, they parameterize W in a way that I − W ⪰ mI (called m-strongly monotone) is always satisfied during training for some m > 0; they convert nonlinearities into proximal operators (which include ReLU, tanh, etc.), and show that using existing splitting methods like forward-backward and Peaceman-Rachford can provably find the unique fixed point.
Markov random field (MRF) and its variants MRF is a form of energy-based model, which model joint probabilities of the form pθ(x) = exp (−Eθ(x)) /Zθ for an energy function Eθ. A common type of MRF is the Boltzmann machine, the most successful variant of which is the restricted Boltzmann machines (RBM) (Hinton, 2002) and its deep (multi-layer) variant (Salakhutdinov & Hinton, 2009). Particularly, RBMs define Eθ(v, h) = −a⊤v− b⊤h− v⊤Wh, where θ = {W,a, b}, v is the set of visible variables, and h is the set of latent variables. It is usually trained using the
contrastive-divergence algorithm, and its inference can be done efficiently by a block mean-field approximation. However, a particular restriction of RBMs is that there can be no intra-layer connections, that is, each variable in v (resp. h) is independent conditioned on h (resp. v). A deep RBM allows different layers of hidden nodes, but there cannot be intra-layer connections. By contrast, our formulation allows intra-layer connections and is therefore is more expressive in this respect. See fig. 1 for the network topology of RBM, deep RBM, and general BM (we also use the term general deep BM interchangeably to emphasize the existence of deep structure). Wu et al. (2016) proposed a deep parameterization of MRF, but their setting only considers a grid of hidden variables h and the connections among hidden units are restricted to the neighboring nodes. Therefore, it is a special case of our parameterization (although their learning algorithm is orthogonal to ours). Numerous works also try to combine deep neural networks with conditional random fields (CRF) (Krähenbühl & Koltun, 2013; Zheng et al., 2015; Schwartz et al., 2017) These models either train a pre-determined kernel as an RNN or use neural networks for producing either inputs or parameters of their CRFs.
Parallel and convergent mean-field It is well-known that mean-field updates converge locally using a coordinate ascent algorithm (Blei et al., 2017). However, local convergence is only guaranteed if the update is applied sequentially. Nonetheless, several works have proposed techniques to parallelize updates. Krähenbühl & Koltun (2013) proposed a concave-convex procedure (CCCP) to minimize the KL divergence between the true distribution and the mean-field variational family. To achieve efficient inference, they use a concave approximation to the pairwise kernel, and their fast update rule only converges if the kernel function is concave. Later, Baqué et al. (2016) derived a similar parallel damped forward iteration to ours that provably converges without the concave potential constraint. However, unlike our approach, they do not use a parameterization which ensures a global mean-field optimum, and their algorithm therefore may not converge to the actual fixed point of the mean-field updates. This is because Baqué et al. (2016) used the prox1f proximal operator (described below), whereas we derive the proxαf operator to guarantee global convergence when doing mean-field updates in parallel. What’s more, Baqué et al. (2016) focused only on inference over prescribed potentials, and not on training the (fully parameterized) potentials as we do here. Lê-Huu & Alahari (2021) brought up a generalized Frank-Wolfe based framework for mean-field updates which include the methods proposed by Baqué et al. (2016); Krähenbühl & Koltun (2013). Their results only guarantee global convergence to a local optimal.

heading: 3 MONOTONE DEEP BOLTZMANN MACHINES AND APPROXIMATE INFERENCE
text: In this section, we present the main technical contributions of this work. We begin by presenting a parameterization of the pairwise potential in a Boltzmann machine that guarantees the monotonicity condition. We then illustrate the connection between a (joint) mean-field inference fixed point and the fixed point of our monotone Boltzmann machine and discuss how deep structured networks can be implemented in this form practically; this establishes that, under the monotonicity conditions on Φ, there exists a unique globally-optimal mean-field fixed point. Finally, we present an efficient parallel method for computing this mean-field fixed point, again motivated by the machinery of monotone DEQs and operator splitting methods.

heading: 3.1 A MONOTONE PARAMETERIZATION OF GENERAL BOLTZMANN MACHINES
text: In this section, we show how to parameterize our probabilistic model in a way that the pairwise potentials satisfy I − Φ ⪰ mI , which will be used later to show the existence of a unique meanfield fixed point. Additionally, since Φ defines a graphical model that has no self-loop, we further require Φ to be a block hollow matrix (that is, the ki × ki diagonal blocks corresponding to each variable must be zero). While both these conditions on Φ are convex constraints, in practice it would be extremely difficult to project a generic set of weights onto this constraint set under an ordinary parameterization of the network.
Thus, we instead advocate for a non-convex parameterization of the network weights, but one which guarantees that the monotonicity condition is always satisfied, without any constraint on the weights in the parameterization. Specifically, define the block matrix
A = [ A1 A2 · · · An ] (2)
with Ai ∈ Rd×ki matrices for each variables, and where d can be some arbitrarily chosen dimension. Then let Âi be a spectrally-normalized version of Ai
Âi = Ai ·min{ √ 1−m/∥Ai∥2, 1} (3)
i.e., a version of Ai normalized such that its largest singular value is at most √ 1−m (note that we can compute the spectral norm of Ai as ∥Ai∥2 = ∥ATi Ai∥ 1/2 2 , which involves computing the singular values of only a ki × ki matrix, and thus is very fast in practice). We define the Â matrix analogously as the block version of these normalized matrices.
Then we propose to parameterize Φ as
Φ = blkdiag(ÂT Â)− ÂT Â (4) where blkdiag denotes the block-diagonal portion of the matrix along the ki×ki block. Put another way, this parameterizes Φ as
Φij = { −ÂTi Âj if i ̸= j, 0 if i = j.
(5)
As the following simple theorem shows, this parameterization guarantees both hollowness of the Φ matrix and monotonicity of I −Φ, for any value of the A matrix. Theorem 3.1. For any choice of parameters A, under the parametrization equation 4 above, we have that 1) Φii = 0 for all i = 1, . . . , n, and 2) I −Φ ⪰ mI .
Proof. Block hollowness of the matrix follows immediately from construction. To establish monotonicity, note that
I −Φ ⪰ mI ⇐⇒ I + ÂT Â− blkdiag(ÂT Â) ⪰ mI ⇐= I − blkdiag(ÂT Â) ⪰ mI ⇐⇒ I − ÂTi Âi ⪰ mI, ∀i ⇐⇒ ∥Âi∥2 ≤ √ 1−m, ∀i.
(6)
This last property always holds by the construction of Âi.

heading: 3.2 MEAN-FIELD INFERENCE AS A MONOTONE DEQ
text: In this section, we formally present how to formulate the mean-field inference as a DEQ update. Recall from before that we are modelling a distribution of the form eq. (1). We are interested in approximating the conditional distribution p(xh|xo), where o and h denote the observed and hidden variables respectively, with a factored distribution q(xh) = ∏ i∈h qi(xi). Here, the standard mean-field updates (which minimize the KL divergence between q(xh) and p(xh|xo) over the single distribution qi(xi)) are given by the following equation,
qi(xi) := softmax  ∑ j:(i,j)∈E Φijqj(xj) + bi  (7) where overloading notation slightly, we let qj(xj) denote a one-hot encoding of the observed value for any j ∈ o (see e.g., Koller & Friedman (2009) for a full derivation). The essence of the above updates is a characterization of the joint fixed point to mean-field inference. For simplicity of notation, defining q = [q1(x1) q2(x2) . . .] T .
We see that qh is a joint fixed point of all the mean-field updates if and only if
qh = softmax (Φhhqh +Φhoxo + bh) (8)
where xo analogously denotes the stacked one-hot encoding of the observed variables.
We briefly recall the monotone DEQ framework of Winston & Kolter (2020). Given input vector x, a monotone DEQ computes the fixed point z⋆(x) that satisfies the equilibrium equation
z⋆(x) = σ(Wz⋆(x) + Ux+ b). (9)
Winston & Kolter (2020) showed that if: 1) σ is given by a proximal operator1 σ(x) = prox1f (x) for some convex closed proper (CCP) f , and 2) if we have the monotonicity condition I−W ⪰ mI (in the positive semidefinite sense) for some m > 0, then for any x there exists a unique fixed point z⋆(x), which can be computed through standard operator splitting methods, such as forwardbackward splitting.
We now state our main claim of this subsection, that under certain conditions the mean-field fixed point can be viewed as the fixed point of an analogous DEQ. This is formalized in the following proposition.
Proposition 3.1. Suppose that the pairwise kernel Φ satisfies I −Φ ⪰ mI 2 for m > 0. Then the mean-field fixed point
qh = softmax (Φhhqh +Φhoxo + bh) (10)
corresponds to the fixed point of a monotone DEQ model. Specifically, this implies that for any xo, there exists a unique, globally-optimal fixed point of the mean-field distribution qh.
As the monotonicity condition of the monotone DEQ is assumed in the proposition, the proof of the proposition rests entirely in showing that the softmax operator is given by prox1f for some CCP f . Specifically, as shown in (Krähenbühl & Koltun, 2013), this is the case for
f(z) = ∑ i zi log zi − 1 2 ∥z∥22 + I {∑ i zi = 1, zi ≥ 0 } (11)
i.e., the restriction of the entropy minus squared norm to the simplex (note that even though we are subtracting a squared norm term it is straightforward to show that this function is convex, since the second derivatives are given by 1/zi − 1, which is always non-negative over its domain).

heading: 3.3 PRACTICAL CONSIDERATIONS WHEN MODELLING MONOTONE BOLTZMANN MACHINES
text: The construction in section 3.1 guarantees monotonicity of the resulting pairwise probabilistic model. However, instantiating the model in practice, where the variables represent hidden units of a deep architecture (i.e., representing multi-channel image tensors with pairwise potentials defined by convolutional operators), requires substantial subtlety and care in implementation. In this setting, we do not want to actually represent A explicitly, but rather determine a method for multiplying Av and ATv for some vector v (as we see in section 3.2, this is all that is required for the parallel mean-field inference method we propose). This means that certain blocks of A are typically parameterized as convolutional layers, with convolution and transposed convolution operators as the main units of computation.
More specifically, we typically want to partition the full set of hidden units into some K distinct sets
q = [q1 q2 . . . qK ] T (12)
1A proximal operator is defined by proxαf (x) = argminz 1 2 ∥x− z∥2 + αf(z).
2Technically speaking, we only need I − Φhh ⪰ mI , but since we want this to hold for any choice of h, we need the condition to apply to the entire Φ matrix.
where e.g., qi would be best represented as a height × width × groups × cardinality tensor (i.e., a collection of a multiple hidden units corresponding to different locations in a typical deep network hidden layer). Note that here qi is not the same as qi(xi), but rather the collection of many different individual variables. These qi terms can be related to each other via different operators, and a natural manner of parameterizing A in this case is as an interconnected set of a convolutional or dense operators. To represent the pairwise interactions, we can create a similarly-factored matrix A, e.g., one of the form
A =  A11 0 · · · 0 A21 A22 · · · 0
... ...
. . . ...
AK1 AK2 · · · AKK  (13) where e.g., Aij is a (possibly strided) convolution mapping between the tensors representing qj and qi. In this case, we emphasize that Aij is not the kernel matrix that one “slides” along the variables. Instead, Aij is the linear mapping as if we write the convolution as a matrix-matrix multiplication. For example, a 2D convolution with stride 1 can be expressed as a doubly block circulant matrix (the case is more complicated when different striding is allowed). This parametrization is effectively a general Boltzmann machine, since each random variable in eq. (12) can interactive with any other variables except for itself. Varying Aij , the formulation in eq. (13) is rich enough for any types of architectures including convolutions, fully-connected layers, and skip-connections, etc.
An illustration of one possible network structure is shown in fig. 2. The precise details of how one computes the block diagonal elements of ATA, and how one normalizes the proper diagonal blocks (which, we emphasize, still just requires computing the singular values of matrices whose size is the cardinality of a single qi(xi)) are somewhat involved, so we defer a complete description to the Appendix (and accompanying code). The larger takeaway message, though, is that it is possible to parameterize complex convolutional multi-scale Boltzmann machines, all while ensuring monotonicity.

heading: 3.4 EFFICIENT PARALLEL SOLVING FOR THE MEAN-FIELD FIXED POINT
text: Although the monotonicity of Φ guarantees the existence of a unique solution, it does not necessarily guarantee that the simple iteration
q (t) h = softmax(Φhhq (t−1) h +Φhoxo + bh) (14)
will converge to this solution. Instead, to guarantee convergence, one needs to apply the damped iteration (see, e.g. (Winston & Kolter, 2020))
q (t) h = prox α f ( (1− α)q(t−1)h + α(Φhhq (t−1) h +Φhoxo + bh) ) . (15)
The damped forward-backward iteration converges linearly to the unique fixed point if α ≤ 2m/L2, assuming I − Φ is m-strongly monotone and L-Lipschitz (Ryu & Boyd, 2016). Crucially, this update can be formed in parallel over all the variables in the network: we do not require a coordinate descent approach as is typically needed by mean-field inference.
The key issue, though is that while prox1f (x) = softmax(x) for f defined as in eq. (11), in general this does not hold for α ̸= 1. Indeed, for α ̸= 1, there is no closed form solution to the proximal operation, and computing the solution is substantially more involved. Specifically, computing this proximal operator involves solving the optimization problem
proxαf (x) = argmin z
1 2 ∥x− z∥22 + α ∑ i zi log zi − α 2 ∥z∥22
subject to ∑ i zi = 1, z ≥ 0. (16)
The following theorem, proved in the Appendix, characterizes the solution to this problem for α ∈ (0, 1) (although it is also possible to compute solutions for α > 1, this is not needed in practice, as it corresponds to a “negatively damped” update, and it is typically better to simply use the softmax update in such cases).
Theorem 3.2. Given f as defined in eq. (11), α ∈ (0, 1), and x ∈ Rk, the proximal operator proxαf (x) is given by
proxαf (x)i = α
1− α W ( 1− α α exp ( xi − α+ λ α )) ,
where λ ∈ R is the unique solution chosen to ensure that the resulting ∑
i prox α f (xi) = 1, and
where W (·) is the principal branch of the Lambert W function.
In practice, however, this is not the most numerically stable method for computing the proximal operator, especially for small α, owing to the large term inside the exponential. Computing the proximal operation efficiently is somewhat involved, though briefly, we define the alternative function
g(y) = log α
1− α W ( 1− α α exp ( y α − 1 )) (17)
and show how to directly compute g(y) using Halley’s method (note that Halley’s method is also the preferred manner to computing the Lambert W function itself numerically (Corless et al., 1996)). Finding the prox operator then requires that we find λ such that ∑k i=1 exp(g(xi+λ)) = 1. This can be done via (one-dimensional) root finding with Newton’s method, which is guaranteed to always find a solution here, owing to the fact that this function is convex monotonic for λ ∈ (−∞, 1). We can further compute the gradients of the g function and of the proximal operator itself via implicit differentiation (i.e., we can do it analytically without requiring unrolling the Newton or Halley iteration). We describe the details in the Appendix, and include an efficient PyTorch function implementation in the supplementary material.
Comparison to Winston & Kolter (2020) Although this work uses the same monotonicity constraint as in Winston & Kolter (2020), our result further requires the linear module Φ to be hollow, and extend their work to the softmax nonlinear operator as well. These extensions introduce significant complications, but also enable us to interpret our network as a probabilistic model, while the network in Winston & Kolter (2020) cannot.

heading: 3.5 TRAINING CONSIDERATIONS
text: Finally, we discuss approaches for training these monotone Boltzmann machines, exploiting their efficient approach to mean-field inference. Probabilistic models are typically trained via approximate likelihood maximization, and since the mean-field approximation is based upon a particular likelihood approximation, it may seem most natural to use this same approximation to train parameters. In practice, however, this is often a suboptimal approach. Specifically, because our forward inference procedure ultimately uses mean-field inference, it is better to train the model directly to output the correct marginals, when running this mean-field procedure. This is known as a marginalbased loss (Domke, 2013). In the context of monotone Boltzmann machines, this procedure has a particularly convenient form, as it corresponds roughly to the “typical” training of DEQ.
In more detail, suppose we are given a sample x ∈ X (i.e., at training time the entire sample is given), along with a specification of the “observed” and “hidden” sets, o and h respectively. Note that the choice of observed and hidden sets is potentially up to the algorithm designer, and can effectively allows one to train our model in a “self-supervised” fashion, where the goal is to predict some unobserved components from others. In practice, however, one typically wants to design hidden and observed portions congruent with the eventual use of the model: e.g., if one is using the model for classification, then at training time it makes sense for the label to be “hidden” and the input to be “observed.”
Given this sample, we first solve the mean-field inference problem to find q⋆h(xh) such that q⋆h = softmax (Φhhq ⋆ h +Φhoxo + bh) . (18)
For this sample, we know that the true value of the hidden states is given by xh. Thus, we can apply some loss function ℓ(q⋆h,xh) between the prediction and true value, and update parameters of the model θ = {A, b} using their gradients
∂ℓ(q⋆h,xh)
∂θ =
∂ℓ(q⋆h,xh)
∂q⋆h
∂q⋆h ∂θ = ∂ℓ(q⋆h,xh)
∂q⋆h
( I − ∂g(q ⋆ h)
∂q⋆h
)−1 ∂g(q⋆h)
∂θ (19)
with g(q⋆h) ≡ proxαf ((1− α)q∗h + α(Φhhq∗h +Φhoxo + bh)) . and where the last equality comes from the standard application of the implicit function theorem as typical in DEQs or monotone DEQs. This backward pass can also be computed via an iterative approach, and here the details exactly mirror that of Winston & Kolter (2020).
As a final note, we also mention that owning to the restricted range of weights allowed by the monotonicty constraint, the actual output marginals qi(xi) are often more uniform in distribution than desired. Thus, we typically apply the loss to a scaled marginal
q̃i(xi) ∝ qi(xi)τi (20) where τi ∈ R+ is a variable-dependent learnable temperature parameter. Importantly, we emphasize that this is only done after convergence to the mean-field solution, and thus only applies to the marginals to which we apply a loss: the actual internal iterations of mean-field cannot have such a scaling, as it would violate the monotonicity condition.

heading: 4 EXPERIMENTAL EVALUATION
text: We evaluate the performance of our model primarily on the MNIST dataset. Our model is able to approximate any conditional distribution. Here we demonstrate how to model missing pixels and the class digits conditioned on the observed pixels, as well as model missing pixels conditioned on class digits and observed pixels. Although MNIST is of course a small-scale problem, the goal here is to demonstrate joint inference and learning over what is still a reasonably-sized joint model, considering the number of hidden units. Nonetheless, the current experiment is admittedly largely a demonstration of the proposed method rather than a full accounting of its performance, a point we highlight in the subsequent section as well. After describing the MNIST experiments, we detail a similar experiment on CIFAR-10 to demonstrate the potential for scaling the approach.
We also show how our inference method differs from the previous ones. On the joint imputation and classification task, we train models using our updates and the updates in Krähenbühl & Koltun (2013); Baqué et al. (2016), and inference each using all three update methods, with and without the monotonicity constraint. We show that our inference method is superior in both performance and convergence speed. Numerically, we also demonstrate that the other two methods could either diverge or not converging to the actual mean-field fixed point. All deferred experiments and details can be found in the appendix.
Experiment setup The original MNIST dataset has one channel representing the gray-scale intensity, ranging between 0 and 1. Here we adopt the strategy of Van Oord et al. (2016) to convert this continuous distribution to a discrete one. We bin the intensity evenly to 4 categories {0, . . . , 3}, and for each channel uses a one-hot encoding of the category so that the input data has shape 4×28×28. We remark that the number of categories is chosen arbitrarily and can be any integer. For image with missing data, if the pixels are randomly masked, we always mask each pixel off independently with probability 60%, such that in expectation only 40% pixels are observed. If a whole patch of pixels is masked, we randomly pick a 14× 14 patch. The patches/masks are chosenly differently for every image, similar to the query training in Lázaro-Gredilla et al. (2020). To make the model class richer, in the patch case we lift the monotonicity constraint, and the model converges regardless. We also conduct the same set of MNIST experiments using a 3-layer DBM and include the results in the appendix.
Results on imputation and classification Training the above described model for 40 epochs with 40% missing pixels leads to a 92.95% test accuracy, and the image reconstruction result is shown in fig. 4d. Based on the reconstructed images and our classification results, our model successfully recovers the unimodal distribution conditioned on the observed pixels. Despite working simultaneously on two different tasks, our model can simply be trained jointly, taking full advantages of the existing auto-differentiation frameworks without additional burden. It is worth noting that while the autoregressive models of Van Oord et al. (2016) are comparable to ours, they require sampling pixels one-by-one in a sequential order at inference time, whereas our model does not assume an underlying sequence and can generate all pixels at once. Comparing to the differently-parameterized monDEQ in Winston & Kolter (2020), whose linear module suffers from drastically increasing condition number (hence in later epochs taking around 20 steps to converge, even with tuned α), our
parameterization produces a much nicer convergence pattern: the average number of forward iterations over the 40 training epochs is less than 6 steps, see fig. 3. When the missing pixels form consecutive patches, our model reconstructs readable digits despite potentially large chunk of missing pixels (fig. 4e). Meanwhile, if the model is given the image labels as input injections, our model performs conditionaly generation fairly well (fig. 4f). These results demonstrate the flexibility of our parameterization for modelling different conditional distributions.
CIFAR-10 Experiment We additionally conduct an experiment on the simultaneous tasks of image pixel imputation and label prediction given partially observed features. Model architecture and training details are given in the appendix. With 50% of the pixels observed, the model obtains 58% test accuracy, and can impute the missing pixels effectively (see fig. 5).

heading: 5 CONCLUSION
text: In this work, we give a monotone parameterization for general Boltzmann machines, and connect its mean-field fixed point to a monotone DEQ model. We provide a mean-field update method that is proven to be globally convergent. Our parameterization allows for full parallelization of mean-field updates without restricting the potential function to be concave, thus addressing issues with prior approaches. Moreover, we allow complicated and hierarchical structures among the variables and show how to efficiently implement them. For parameter learning, we directly optimize the marginalbased loss over the mean-field variational family, circumventing the intractability of computing the partition function. Our model is evaluated on the MNIST and CIFAR-10 dataset for simultaneously predicting with missing data and imputing the missing data itself. There are several potential future directions. First, the monotone parameterization in theorem 3.1 is not an if-and-only-if statement, hence potentially making our model more restrictive than necessary. The second direction is that although we have a fairly efficient implementation of proxαf , it is still slower than normal nonlinearities like ReLU or softmax. It is an interesting direction to more efficiently scale these models.

heading: A APPENDIX
text: A.1 DEFERRED PROOFS
Proof of theorem 3.2. By definition, the proximal operator induced by f (the same f in eq. (11)) and α solves the following optimization problem:
min z
1 2 ∥x− z∥2 + α ∑ i zi log zi − α 2 ∥z∥2
s.t. zi ≥ 0, i = 1, . . . , d,∑ i zi = 1
of which the KKT condition is
− xi + zi + α+ α log zi − αzi + λ− µi = 0, for i ∈ [d] µi ≥ 0 zi ≥ 0∑ i∈[d] µizi = 0 d∑ i=1 zi = 1
We have that µi = 0 is feasible and the first equation of the above KKT condition can be massaged as
− xi + zi + α+ α log zi − αzi + λ− µi = 0 ⇐⇒ (1− α)zi + α log zi = xi − α− λ
⇐⇒ (1− α)zi + α log zi α = xi − α− λ α
⇐⇒ exp ( (1− α)zi + α log zi
α
) = exp ( xi − α− λ
α ) ⇐⇒ zi exp ( 1− α α zi ) = exp ( xi − α− λ α
) ⇐⇒ 1− α
α zi exp ( 1− α α zi ) = 1− α α exp ( xi − α− λ α ) ⇐⇒ 1− α
α zi = W ( 1− α α exp ( xi − α− λ α )) where W is the lambert W function. Notice here zi > 0. Our primal problem is convex and Slater’s condition holds. Hence, we conclude that
zi = α
1− α W ( 1− α α exp ( xi − α− λ α )) .
A.2 CONVOLUTION NETWORK
It is clear that the monotone parameterization in section 3 directly applies to fully-connected networks, and all the related quantities can be calculated easily. Nonetheless, the real power of the DEQ model comes in when we use more sophisticated linear operators like convolutions. In the context of Boltzmann machines, the convolution operator gives edge potentials beneficial structures. For example, when modeling the joint probability of pixels in an image, it is intuitive that only the nearby pixels depend closely on each other.
Let A ∈ Rk×k×r×r denote a convolutional tensor with kernel size r and channel size k, let x denote some input. For a convolution with stride 1, the block diagonal elements of ATA simply form a 1× 1 convolution. In particular, we apply the convolutions
−AT (A(x)) + Ã(x) (21)
where Ã is a 1× 1 convolution given by Ã[:, :] = ∑ i,j A[:, :, i, j]TA[:, :, i, j]. (22)
We can normalize by the spectral norm of Ã term to ensure strong monotonicity. Since Ã can be rewritten as a k × k matrix and k is usually small, its spectral norm can be easily calculated. It takes more effort to work out convolutions with stride other than 1. Specifically, the block diagonal terms do not form a 1 × 1 convolution anymore, instead, the computation varies depending on the location. It is easier to see the computation directly in the accompanying code.
Grouped channels It is crucial to introduce the concept of grouped channels, which allows us to represent multiple categorical variables in a single location, such as the three categorical variables representing the three (binned) color channels of an RGB pixel. In this case, each of the three RGB channels will be represented by a different group of k channels representing the k bins. The grouping is achieved by having the nonlinearity function (softmax) applied to each group separately. We remark that the convolutions themselves are not grouped, otherwise none of the red pixels would interact with green or blue pixels, etc. Instead, we want all RGB channels to interact with each other (except that channel i at position (j, k) does not interact with itself). That means in eq. (4), the blkdiag(ÂT Â) is grouped in the following way. Recall that this block diagonal term has element of size ki × ki for i ∈ [n]. This parameterization has only 1 group. With g groups, the element of the block diagonal matrix then has size ki1 × ki1 , . . . , kig × kig for i ∈ [n], where ∑ j∈[g] kij = ki. We also observe empirically that grouping the latent variables improves the performance.
A.3 EFFICIENT COMPUTATION OF proxαf
The solution to the proximal operator in damped forward iteration given in theorem 3.2 involves the Lambert W function, which does not attain an analytical solution. In this section, we show how to efficiently calculate the nonlinearity σ(xi), as well as its Jacobian matrix for backward iteration.
Let f(y) = α1−αW ( 1−α α exp ( y α − 1 )) , and we have
x = log f(y) = log α
1− α + log ey/α−1 + log 1− α α −W ( 1− α α exp ( y α − 1 )) ,
where the last equality uses the identity log(W (x)) = log x − W (x). Rewrite W ( 1−α α exp ( y α − 1 )) = f(y) 1−αα and massage the terms, we have that solving log f(y) is equivalent to finding the root of h(x) = y − α− ex(1− α)− αx. Direct calculation shows that h′(x) = −α − (1 − α)ex and h′′(x) = −(1 − α)ex. Note here y is the input and it is known to us, and x is a scalar. Hence we can efficiently solve the root finding problem using Halley’s method. For backpropagation, we need dxdy , which can be computed by implicit differentiation:
h(x) = y − α− ex(1− α)− αx = 0
=⇒dx dy = 1 α+ (1− α)ex =
1
y − αx .
Now we can find λ s.t ∑ i zi = 1 using Newton’s method on g(λ) = ∑ i e log(f(xi+λ)) − 1 = 0. Note this is still a one-dimensional optimization problem. A direct calculation shows that dgdλ =∑ i e log(f(xi+λ)) d log(f(xi+λ) dλ , and above we have already calculated that
d log(f(xi + λ)
dλ =
dx∗
dy =
1
y + λ− αx .
For backward computation, by the chain rule, we have:
delog f(xi+λ)
dxi = elog f(xi+λ)
d log(f(xi + λ))
dxi
= elog f(xi+λ) 1 + dλ/dxi
xi + λ− α log(f(xi + λ)) ,
where the last step is derived by implicit differentiation. Now to get dλ/dxi, notice that by applying the implicit function theorem on p(x, λ(x)) = ∑ i e log(f(xi+λ)) − 1 = 0, we get
dλ dxi = −
( dp
dλ
)−1 dp
dxi .
Thus we have all the terms computed, which finishes the derivation.

heading: B ADDITIONAL EXPERIMENTS AND DETAILS
text: Here we provide the model architectures and experiment details omitted in the main text.
Model architecture For MNIST experiments, using the notation in eq. (13), we design a 20-layer deep monotone DEQ with the following structure:A11 0 0 0A21 A22 0 0A31 A32 A33 0
0 0 A43 A44  , where A11 is a 20×20×3×3 convolution, A22 is a 40×40×3×3 convolution, A21 is a 40×20×3×3 convolution with stride 2, A33 is a 80×80×3×3 convolution, A31 is a 80×20×3×3 convolution with stride 4, A32 is a 80 × 40 × 3 × 3 convolution with stride 2, A43 is a (80 · 7 · 7) × 10 dense linear layer, and A44 is a 10 × 10 dense linear layer. The corresponding variable q as in eq. (12) then has 4 elements of shape (20×28×28), (40×14×14), (80×7×7), (10×1). When applying the proximal operator to q, we use 1, 10, 20, 1 as their number of groups, respectively.
Training details and hyperparameters Treating the image reconstruction as a dense classification task, we use cross-entropy loss and class weights 1−β1−βni with β = 0.9999 (Cui et al., 2019), where ni is the number of times pixels with intensity i appear in the hidden pixels. For classification, we use standard cross-entropy loss. To enable joint training, we put equal weight of 0.5 on both task losses and backpropagate through their sum. For both tasks, we put τiΦq∗i into the crossentropy loss as logits, as described in eq. (20). Since mean-field approximation is (conditionally) unimodal, the scaling grants us the ability to model more extreme distributions. To achieve faster damped forward-backward iteration, we implement Anderson acceleration (Walker & Ni, 2011), and stop the fixed point update as soon as the relative difference between two iterations (that is, ∥qt+1 − qt∥/∥qt∥) is less than 0.01, unless we hit a maximum number of 50 allowed iterations. For proxαf and the damped iteration, we set α = 0.125 (Although one can tune down α whenever the iterations do not converge, empirically this never happens on our task). We use the Adam optimizer with learning rate 0.001. Our models are trained on one GeForce GTX 2080 Ti GPU.
Comparison to past inference methods We conduct numerical experiments to compare our inference updating method to the ones proposed by Krähenbühl & Koltun (2013); Baqué et al. (2016), denoted as Krähenbühl’s and Baqué’s respectively. Krähenbühl’s fast concave-convex procedure (CCCP) essentially decomposes to eq. (14), the un-damped mean-field update with softmax. This update only converges provably when Φ is concave. Baqué’s inference method can be written as
q (t) h = softmax ( (1− α) log q(t−1)h + α(Φhhq (t−1) h +Φhoxo + bh) ) . (23)
This algorithm provably converges despite the property of the pairwise kernel function. However, this procedure converges in the sense that the variational free energy keeps decreasing. Therefore their fixed point may not be the true mean-field distribution eq. (10). In this experiment, we train the
models using three different updating methods, and perform inference using three methods as well, with and without the monotonicity condition.
Krähenbühl’s and Baqué’s methods often do not converge in the backward pass (there’s no theoretical guarantees neither). To rule out the impact of the backward iteration, during training we directly update use the gradient of the forward pass, instead of using a backward gradient hook to compute eq. (19). Figure 7 and fig. 8 demonstrate how the three update methods impute missing pixels when trained with different update rules, with and without the monotonicity condition, respectively. Krähenbühl’s usually does not converge when the model is trained with our method or Baqué’s, whereas the other two methods impute the missing pixels well. The classification results are presented in table 3 and table 4. Notice that when trained with our method or Baqué’s, the convergence issue of Krähenbühl’s leads to horrible classification accuracy. Our method is superior to other inference methods when the model is trained in a different update fashion. For example, if the model is trained by using Krähenbühl’s, it makes sense that the model performs the best if the inference is also Krähenbühl’s since the parameters are biased toward that particular inference method. However, our method in this case outperforms Baqué’s.
After these methods halt and return qTh , we run one more iteration of
qT+1h = softmax ( Φhhq T h +Φhoxo + bh ) , (24)
and record the relative update residual ∥qT+1h − qTh ∥/∥qTh ∥ for randomly selected 4000 MNIST images. The results are listed in table 1 and table 2. To alleviate the effect of numerical issues, we strength the convergence condition to either the relative residual is less than 10−3 or the number of iterations exceeds 100 steps.
It appears in table 1 and table 2 that although our method has a much lower residual compare to Baqué’s, both of them seem small and convergent. This is because the “optimal” fixed point in this setting on MNIST might be unique and both methods happen to converge to the same point. However, this is in general not true. We compare our method vs Baqué’s on 400 randomly selected MNIST test images with 40% pixels observed, and perform mean-field update until the relative residual of [0.1, 0.05, 0.01, 0.005, 0.001, 0.0005, 0.0001] is reached (without step constraint), respectively. Then we measure the TV distance between the distributions computed by these two methods on the remaining 60% pixels, as well as the convergence speed. The results are demonstrated in fig. 6. One can see that when the model is trained (using Krähenbühl’s, fig. 9a), the TV distance converges to 0 as the tolerance decreases. However, when the model is just initialized (but still constrained to be monotone), the TV distance remains large (fig. 6c). Even though in this case the optimal fixed point may be unique, our method is still superior to Baqué’s: it takes us less iterations till convergence, despite whether the model is trained or not..
CIFAR-10 model architecture Architecture is the same as for the MNIST experiments with the following exceptions: A11 is a 20× 20× 3× 3 convolution, A22 is a 24× 24× 3× 3 convolution, A21 is a 24× 20× 3× 3 convolution with stride 2, A33 is a 48× 48× 3× 3 convolution, A31 is a 48×20×3×3 convolution with stride 4, A32 is a 48×24×3×3 convolution with stride 2, A43 is a (48·8·8)×10 dense linear layer, and A44 is a 10×10 dense linear layer. The corresponding variable q as in eq. (12) then has 4 elements of shape (60× 32× 32), (24× 16× 16), (48× 8× 8), (10× 1). When applying the proximal operator to q, we use 1, 6, 12, 1 as their number of groups, respectively.
CIFAR-10 training details Training details are the same as for the MNIST experiments with the following exceptions: The model is trained for 100 epochs using standard data augmentation. During the first 10 epochs, the weight on the reconstruction loss is ramped up from 0.0 to 0.5 and the weight on the classification loss ramped down from 1.0 to 0.5. Also during the first 20 epochs, the percentage of observation pixels is ramped down from 100% to 50%.

heading: C COMPARISON TO RBM/DBM
text: We conduct the same set of experiments on MNIST using DBM for comparison. We use a 3- layer DBM where the first hidden layer has 300 neurons, and the last hidden layer (representing the digits) has 10 neurons, amounting to in total 238,200 parameters. Our proposed model in the MNIST experiments uses 165,300 parameters.
For image imputation, we randomly mask off 60% pixels, or a randomly selected 14 × 14 patch; the results are shown in fig. 9. In the experiment with 60% of pixels randomly masked, we also test
the model on predicting the actual digit simultaneously. The test accuracy is 93.58%. Our model achieves comparable test accuracy (92.95%) and imputation compared to this DBM, given fewer parameters and despite the the monotonicity constraint of our model.
The DBM is trained using CD-1 algorithm for 100 epochs with a batch size of 128 and learning rate of 0.01. For imputation and classification, the DBM uses Gibbs sampling of 100 steps, although the quality of the imputed image and test accuracy are insensitive to the number of steps.

Review summaries:
This paper proposes a class of model called monotone deep Boltzmann machines, where the underlying potentials are parameterized (e.g., by CNNs) such that they obey some monotonicity constraint. This constraint ensures that the inference problem has a global optimum, which can be found using some generalized variant of parallel mean field. The method is inspired from monotone DEQ, previously proposed by Winston & Kolter (2020). Experiments on a joint task of image denoising and classification show that the proposed method can effectively model complex data distributions such as images.

In this paper the authors propose a restricted parameterization of the Boltzmann machine that guarantees that for any set of observations, the mean field objective has a single global optimum. Furthermore, that global optimum can be provably achieved using damped parallel mean-field updates, which make inference efficient. To turn inference into learning, the model is treated as a supervised learning model: some of its variables are considered to be observed inputs and some of its variables are considered to be target outputs (known at test time). The usual, marginal cross-entropy loss is the optimization target for learning.

This paper proposes a new family of monotone deep Boltzmann machines where the pairwise potentials satisfy a monotonicity condition, giving rise to efficient mean-field iteration with provable convergence guarantees. The convergence is obtained by drawing connections with monotone deep equilibrium models. Small-scale experiments are done as proof of concept.