Easy Variational Inference for Categorical Observations via a New View of Diagonal Orthant Probit Models
Abstract
In pursuit of tractable Bayesian analysis of categorical data, auxiliary variable methods hold promise, but impose asymmetries on the truly unordered categories or spoil scalability via strong dependencies in posteriors over parameters.
The Diagonal Orthant Probit (DO-Probit) model proposed by Johndrow, Lum, and Dunson (AISTATS 2013) avoids these difﬁculties, treating all categories symmetrically while yielding tractable conditionally conjugate inference.
However, we show that the intended DO-Probit likelihood for categorical observations, when paired with a normal prior, does not yield a conjugate posterior.
Instead, we clarify that their posterior analysis is only correct for a different model that treats observations as multiple independent binary draws.
This raises two questions: Other than tractability, what justiﬁes the binary model for categorical data?
And how should a binary model make categorical predictions?
To resolve these issues, using variational methods we obtain a lower bound of a categorical model’s marginal likelihood that can be optimized by ﬁtting the conjugate binary model.
Optimizing this bound enjoys all beneﬁts advocated in the original DO-Probit work.
We further extend this fast, reliable covariate-informed modeling of categorical outcomes to groups or sequences of data related in a hierarchy.
INTRODUCTION
We consider the problem of modeling categorical data informed by covariates using the machinery of generalized linear models.
Because our intended big data applications may involve rare events or little available data for some quantities of interest, we pursue Bayesian analysis in order to estimate distributions over unknown parameters given available data, and then average over these distributions when making predictions.
While many generalized linear models for categorical observations have been tried, Bayesian analysis of these models remains a difﬁcult problem with substantial active research due to the need for methods that are simultaneously accurate, tractable, and scalable.
The most common modeling choice for categorical data is multi-class logistic regression, which uses a softmax (a.k.a.
multi-logit) function to produce category probabilities.
The model is not conjugate, and so estimating posteriors over weight parameters requires expensive sampling methods [Hoffman and Gelman, 2014] or non-conjugate variational optimization methods [Wang and Blei, 2013,
Braun and McAuliffe, 2010, Kucukelbir et al., 2017].
Recent auxiliary variable methods [Polson et al., 2013] have yielded expanded binary logistic models with conjugate conditionals, but extensions to multiple categories require stick-breaking [Linderman et al., 2015].
Stick-breaking imposes an asymmetric order over categories, yet in many cases it is unnatural to view category selection as a sequential process.
In practice, this asymmetry complicates prior speciﬁcation and inference quality [Zhang and Zhou, 2017].
An alternative model is multi-class probit regression, whose link function is the cumulative distribution function of the
Normal distribution.
The probit admits conjugate inference under a well-known auxiliary variable representation [Albert and Chib, 1993, Held and Holmes, 2006].
However, multiclass probit models encode strong posterior dependence between entries of the auxiliary parameter vectors.
This dependence requires one-entry-at-a-time sampling instead of joint sampling [Johndrow et al., 2013], yielding poor mixing performance as the number of categories grows.
Furthermore, implementations often require picking a “base category”; this choice can impact the practical results of inference [Burgette et al., 2021].
Finally, the multinomial probit’s strong dependence results in a lack of closed-form category probabilities [Johndrow et al., 2013], which alone has prevented adoption within more complicated models (e.g., see [Holsclaw et al., 2017]).
Motivated by difﬁculties that arise from these previous efforts, Johndrow et al.
[2013] introduced the diagonal orthant multinomial probit (DO-Probit) model.
Their proposed construction treats all categories symmetrically, yields tractable category probabilities, and achieves conditionally conjugate posteriors while avoiding any cross-category dependence among auxiliary variables.
These features should make the
DO-Probit a prime candidate for fast, scalable Bayesian modeling of categorical data.
However, we have uncovered subtle but critical gaps in the technical justiﬁcation for the Bayesian analysis of the DOProbit provided by Johndrow et al.
In particular, we ﬁnd that the DO-Probit categorical likelihood, when paired with a normal prior, does not admit conditionally-conjugate posteriors even with auxiliary variables.
The Gibbs sampler presented in Sec.
3 of Johndrow et al.
[2013] is not a correct sampler when the intended posterior is formed from a categorical likelihood.
Instead, it is a correct sampler only when understood as ﬁtting a different model : the independent binary model, which is conditionally conjugate after augmentation via the classic arguments for the binary probit [Albert and Chib, 1993].
This “model swap” may not be obvious to readers of the original DO-Probit text and has apparently gone unnoticed in some later uses of this method.
Magnusson et al.
[2020] use the Gibbs sampler of Johndrow et al.
[2013] directly for a categorical likelihood, which is not a valid posterior sampling technique.
The need for a model swap raises important questions: given categorical data, what justiﬁes using the independent binary model, which is not intended for one-of-K categories, other than tractability?
If we insist on using the binary model, how can we use it to make valid categorical predictions?
This paper makes several key contributions in attempting to answer these questions:
We show that the DO-Probit model for categorical data does not have a conditionally-conjugate posterior over weights even with auxiliary variables (Sec.
We suggest an effective way to make predictions about heldout categorical data using the independent binary model (Sec. 2.6), which in the intercepts-only case surprisingly achieves indistinguishable prediction quality compared to models designed for categorical data.
We introduce a new generalized linear model for categorical data, which we call the Simpliﬁed DO-Probit model (SDO-Probit, Sec. 2).
This model has similar beneﬁts as the original DO-Probit (symmetric treatment of categories and closed-form category probabilities computed using the probit function), but yields the effective prediction strategy naturally.
We derive a tractable lower bound for the marginal likelihood of the Simpliﬁed DO-Probit categorical model, and show that optimizing this bound corresponds to ﬁtting the observed data using the independent binary model (Sec. 3).
This allows us to justify applying the conjugate machinery suggested by Johndrow et al.
[2013] to otherwise intractable categorical models.
Our bound holds for the original DO-Probit model as well.
Using our bound, we provide an easily-implemented variational inference strategy for categorical SDOProbit model that is fast and reliable.
In particular, the inference has three beneﬁcial properties: (a) closedform updates, (b) decoupled auxiliary variables, and (c) symmetric handling of categories.
No prior work known to us achieves all three properties.
We provide an algorithm which allows the technique to be extended to any model which uses the SDO- (or DO-) Probit complete likelihood (see Algorithm 1 in Sec. 3.2).
We highlight the advantages of this ﬂexibility by introducing a hierarchical DO-Probit model for modeling groups or sequences of categorical data (Sec. 4).
In summary, our new variational methods provide a principled route to fast, scalable modeling of categorical data using a surrogate binary model.
MODELS
We now formally introduce the models of interest and establish their properties.
First, we introduce one model for observed categorical data, the diagonal orthant probit, denoted throughout as DO, ﬁrst described in Johndrow et al.
Next, we introduce an alternative model for multivariate binary data, which we refer to as “independent binary” or IB, ﬁrst described in Albert and Chib [1993].
Finally, we introduce a new model for categorical data, which we call the Simpliﬁed-DO Probit (SDO).
We brieﬂy discuss maximum likelihood estimation for each model.
We then expose that in a posterior estimation context, neither categorical model is conjugate even with auxiliary variables, due to an intractable normalization constant.
Thus, the Gibbs sampler presented by Johndrow et al.
[2013] for the “DO-Probit” model is in fact only a valid sampler for the IB.
This raises the question of why one should apply the IB to categorical data, which we address in Sec. 3.
2.1 DO-PROBIT MODEL FOR CATEGORY DATA
The diagonal orthant probit model [Johndrow et al., 2013] describes the generative process for N categorical outcomes, where each observation (indexed by i ∈ { 1 , . . . N } ) is a one-of-K indicator y i ∈ { 1 , 2 , . . . K } .
These outcomes are generated by transforming known covariate vector x i ∈ R M with unknown weight parameters β ∈ R M × K .
Let β k designate the k -th column of β .
To generate observation i , we ﬁrst draw a latent continuous vector z i ∈ Ω ⊂ R K from a multivariate normal whose mean is determined by the dot product of covariates and weights but whose possible values are truncated to the Diagonal Orthant region Ω .
This region con- tains all K -dimensional vectors with one entry positive and the rest negative.
Formally, Ω := (cid:83)
Kk
=1 Ω k , where
[Equation]
Given vector z i , our second generative step is simply to set the outcome y i to indicate which entry of z i is positive.
This two-step sampling is formalized for each i ∈ { 1 , . . . N } :
[Equation]
Here, T N ( η , V , R ) deﬁnes a truncated multivariate normal distribution formed when a normal with mean η and covariance V is truncated to region R .
The DO’s complete-data likelihood is:
[Equation]
where C i [Equation] ) :=
[Equation]
and where Φ is the CDF of a standard normal.
By construction, C i ( β ) is the fraction of probability mass of the unconstrained Gaussian contained in the set Ω , i.e. the orthants with only one positive entry, and thus 0 < C i ( β ) < 1 .
This fact will be useful in Sec. 3.
We can view z i as an auxiliary that can be marginalized out, yielding category probabilities conditioned only on β :
[Equation]
This tractable formula was ﬁrst given in Johndrow et al.
[2013, p. 4].
We emphasize that the standard multi-class probit model has no such closed form.
2.2 A PROBIT MODEL FOR BINARY VECTORS
We now consider another model, which we call the Independent Binary probit model (IB).
This model produces N observations (indexed by i ), where each observation a binary vector ¯ y i of size K .
Crucially, ¯ y i is not a one-hot vector: any number of entries could be 1 or 0.
The variables here play similar roles to counterparts in the earlier DO model, but have different domains.
Thus, we use symbols with bars to denote IB variables.
Following Albert and Chib [1993], to generate observation i using the IB model, we ﬁrst sample an auxiliary variable ¯ z ik for each label k from a Normal whose mean is the dot product of covariates x i and weights ¯ β k ∈ R M :
[Equation]
Second, we use the sign of the value of ¯ z ik to deterministically set the binary observation for label k :
[Equation]
The complete data likelihood for the IB model is thus:
[Equation]
2.3 A SIMPLIFIED DO-PROBIT MODEL
We introduce a third model, which we call the simpliﬁed diagonal orthant probit or just SDO, which produces categorical observations.
Given covariates x i and weight parameters β , we produce categorical outcome y i via:
[Equation]
This is a simpler way to use the probit link function to produce closed-form categorical probabilities than in the original DO model in Eq.
(2.1.4), hence the name “simpliﬁed” DO-Probit.
While a similar construction appears in
Magnusson et al.
[2020, p. 6], here for the ﬁrst time we detail important properties including this model’s complete data likelihood with auxiliaries and resulting non-conjugacy.
The SDO is completed by sampling auxiliary variable z i ∈
Ω y i
⊂ R K from a truncated Normal whose pre-truncated mean is a dot product of weights and covariates like previous models:
[Equation]
By deﬁnition, z i has only one positive entry (at index y i ) and all other entries are negative.
Thus, this model still uses a diagonal orthant construction.
The SDO complete data likelihood is given by:
[Equation]
Here, the product of the indicators is sufﬁcient to restrict z i ∈
Ω y
i as required.
C SDO is a normalization term,
[Equation]
which interestingly for this model depends on the chosen category index
y i =
k .
Like the previous normalization term, C SDO always produces values between 0 and 1.
While this SDO model enjoys several properties shared by the original DO model (closed-form category probabilities {+DO posterior at N =4 3 2 1 0 1 2 
 DO posterior at N = DO MLEs+} IB {+MLE 
+}
{+SDO posterior at N =4 3 2 1 0 1 2 
 SDO posterior at N = SDO MLEs IB MLE 
 IB posterior at N =4 3 2 1 0 1 2 
+}
IB posterior at N =
IB MLE and symmetric treatment of categories), SDO is particularly useful for connecting the IB model to categorical outcomes.
2.4 POINT ESTIMATION AND IDENTIFIABILITY
In general, there are inﬁnitely many possible β parameters that would yield the same category probabilities under the DO model in Eq.
Thus, a maximum likelihood estimate (MLE) for β under the DO likelihood is not unique. Similarly, the SDO -likelihood also has a non-unique MLE.
Below, we outline key properties about MLE estimators for all models in the intercepts-only (no covariates) setting.
Unique MLE for intercept-only IB.
Let p 1 , . . .
p K be empirical category frequencies in a K -class intercepts-only dataset: p k = 1 N (cid:80) Ni = 1 1 y i = k .
As Johndrow et al.
[2013] argue, if we encode the observations
y i as one-hot vectors, there is a unique maximum likelihood estimate for ¯ β under the IB model, available by setting the weight for each k as:
[Equation]
MLE for intercept-only SDO.
A similar argument can be used to deﬁne an ML estimate for the SDO in an interceptsonly setting.
Suppose we restrict (cid:80) Kk =1 Φ( β k ) =
r for some r ∈ (0 , min (cid:96) 1 p (cid:96) ) .
Fixing r , the unique MLE for SDO is:
[Equation]
When r = 1 , the SDO MLE for β here is exactly equal to the IB MLE from Eq.
MLE for intercept-only DO.
Similarly, in an interceptsonly version of the DO model, if we enforce the constraint (cid:80) Kk =1 Φ( β k ) Φ( − β k ) = s for a speciﬁc s > 0 , then we can com1,2    3 , DO posterior at N =  22 1,      3 ,
DO posterior at N =1,2    3 ,
SDO posterior at N =  22 1,      3 ,
SDO posterior at N =1,2    3 ,
IB posterior at N =  22 1,      3 , IB posterior at N = pute the MLE for each category k as: [Equation] [Equation]
Plugging this value into Eq.
(2.1.4) yields category probabilities that exactly match the observed frequencies p .
Stepping back, we emphasize that Johndrow et al.
[2013]’s only technical justiﬁcation for using the IB model seems to be that in an intercepts-only setting, the IB has unique MLE weights that exactly match the MLE for a categorical DO model 1 .
We ﬁnd this argument for using IB insufﬁcient, as it applies only to the maximum likelihood point estimation setting without covariates.
This leaves a critical gap: we lack justiﬁcation for the IB as a model for categorical data with covariates when we seek posteriors , not just point estimates.
2.5 BAYESIAN POSTERIOR ANALYSIS
The three models discussed here – DO, SDO, and IB – each induce a distinct posterior over weights when given the same data and same prior.
Fig.
1 shows each model’s posterior over two weight coefﬁcients given toy categorical data in the intercepts-only setting. Fig. 2 does the same given a synthetic dataset with covariates .
Posteriors are estimated via the NUTS sampler [Hoffman and Gelman, 2014], as implemented in NumPyro [Phan et al., 2019] using JAX for automatic differentiation [Bradbury et al., 2018].
For large N , we also show the unique MLE for IB and the manifold of possible (non-unique) ML estimates of weights for DO and
SDO, using formulas in Eq. (2.4.1)-(2.4.3) (these formulas only apply in Fig. 1). In the no-covariate case of Fig. 1, SDO and IB share an MLE, and thus for large N the IB posterior concentrates its mass where the SDO does.
However, in Fig. 2, all three models have distinct posteriors and even with N =examples, the IB posterior does not necessarily concentrate where either DO or SDO do.
We now assess the conjugacy properties of each model.
We assume the weight parameters are given a normal prior.
Observation 1: The IB model’s complete conditionals all have conjugate form.
Classic arguments about the binary probit [Albert and Chib, 1993, Held and Holmes, 2006] show that the intended posterior conditionals can be written as a Normal for p IB (¯ β | ¯ z , ¯ y ) and a truncated Normal for p IB (¯ z | ¯ β , ¯ y ) .
This posterior is clearly the one ﬁt by the
Gibbs sampler in Johndrow et al.
Observation 2: The DO model’s complete conditional for weights β is not normal and thus not conjugate.
Combining a normal prior over β with the DO-Probit likelihood in Eq.
(2.1.2), the resulting posterior p ( β | z , y ) includes the normalizing constant C i ( β ) in Eq.
(2.1.3), whose functional form spoils conjugacy: including this term means the log posterior is not a quadratic function of β .
Observation 3: The SDO model’s complete conditional for weights β is not normal and thus not conjugate.
We follow a similar argument as the DO model.
The complete data likelihood of Eq.
(2.3.2) contains the normalizing constant C SDO i,y i ( β ) (Eq.
that spoils conjugacy.
Thus, both DO and SDO categorical models are not conjugate, simply by inspection of the complete data likelihood.
This is not stated clearly in Johndrow et al.
and has led others to assume their Gibbs sampler is correct for a categorical likelihood, which is not true.
So while Johndrow et al.
[2013] advocate for Bayesian analysis of categorical data with DO Probit models, in fact the Gibbs sampler they derive is only correct for the independent binary generative model.
This leaves two gaps in understanding: How should the IB make predictions for categorical data?
More importantly, why is the IB appropriate for Bayesian analysis of categorical data?
We answer the ﬁrst question in Sec. 2.6, and the second in Sec. 3.
2.6 PREDICTIONS FOR CATEGORICAL DATA
Suppose we have speciﬁc weight values ¯ β for the IB that have been ﬁt to purely categorical data.
[-DO ML estimators IB ML Estimator 4 3 2 1 0 1 2 
-]
[-SDO posterior at N =4 3 2 1 0 1 
-]
[-SDO ML Estimator IB ML Estimator 4 3 2 1 0 1 2 
-]
[-IB posterior at N =4 3 2 1 0 1 
-]
For example, this may be an ML estimate of the weights or a posterior sample.
We wish to predict a new observation, which we are reliably informed will be categorical.
The IB’s support is then misspeciﬁed, because it could produce any binary vector, not just those with one positive entry.
IB+DO.
It seems natural that a new observation (indexed by ∗ ) belongs to category k with probability proportional to the IB likelihood producing the one-hot vector
e k :
[Equation]
Here, we normalize over the set of possible one-hot vectors of size K .
We call this the IB+DO estimator of heldout likelihood, because mathematically this reduces to plugging
IB weights ¯ β into DO’s probability formula in Eq.
IB+SDO.
An alternative estimator of heldout likelihood is:
[Equation]
(2.6.2) which we refer to as “IB+SDO”, because it uses our Simpliﬁed DO formula for translating the IB weights ¯ β into category probabilities.
For another motivation, see Sec.
B.
These two estimators are compared both visually and quantitatively in Fig.
3, where posterior samples of weights from the same IB model are plugged into each estimator to predict heldout categorical data.
We see signiﬁcant practical reasons to prefer the IB+SDO estimator in this interceptsonly setting, as it matches the true category frequencies far better than IB+DO.
Despite the fact that the IB model is misspeciﬁed (it is not naturally a model for categorical data), by ﬁtting the IB model and then applying the SDO estimator (which ensures valid categorical predictions) we can surprisingly deliver heldout likelihoods that for all dataset sizes
N > 16 are indistinguishable from models like DO or SDO that directly capture the categorical nature of data.
VARIATIONAL METHODS
We have clearly established that [Johndrow et al., 2013]’s tractable posterior is targeted at the independent binary model, which is not a model for categorical data.
In this section, we provide a principled justiﬁcation for applying a binary model to categorical data, based on relating the binary model and our suggested Simpliﬁed DO-Probit categorical model via variational methods.
Variational inference [Blei et al., 2017] deterministically approximates a posterior distribution by ﬁnding the member
Q ∈ Q of a tractable family of distributions which maximizes a lower bound on the logarithm of the evidence (the marginal likelihood of the data).
This lower bound is known as the evidence lower bound or “ELBO”.
Our strategy is to show that this lower bound under the misspeciﬁed (but tractable, conjugate) IB model, is in turn a lower bound for the same data under the desired (but intractable, non-conjugate) SDO model.
Thus, by using the conjugate machinery suggested in Johndrow et al.
[2013] to obtain closed-form coordinate ascent updates that maximize this surrogate bound, we are not merely swapping in a convenient model that undesirably is unaware of the categorical nature of our data, but instead exactly improving the ﬁt of
[Figure]
N
IB + DO
IB + SDO
SDO
DO
Cat.
MLE -1.
estimates category frequencies via maximize likelihood.
Setting: We ﬁt to the ﬁrst N examples of a toy dataset, generated with K = 3 categories with ground truth frequencies [0.02, 0.7, 0.28].
All models use a standard Normal prior over weights.
a suitable generative model.
Our strategy will later allow us to easily extend the SDO to more complicated graphical models, as we illustrate in Sec. 4.
3.1 RELATING MODELS VIA BOUNDS ON MARGINAL LIKELIHOODS
We begin with a variational treatment of the SDO model (Eq. with a Normal prior density π on β . The evidence of interest is p SDO ( y ) = (cid:82) (cid:82) p SDO ( y , z | β ) π ( β ) d z d β .
If we select Q as any distribution over ( z ∈ Ω , β ∈ R M × K ) , and let q ( · ) be the density of Q , then the traditional lower bound of the log of the evidence, which we denote ELBO SDO ≤ log p SDO ( y ) , follows from Jensen’s inequality:
[Equation]
Unfortunately, this lower bound is intractable due to the
SDO normalizing constant (Eq. (2.3.3))
; we must handle
N terms of the form E q [log C SDO i,y i ( β )] , which lack closedform expression and thus require high-dimensional integral approximations when the number of categories is large.
However, we observe that given the same data y i and a corresponding auxiliary variable z i ∈
Ω y i , the SDO model will always assign a higher likelihood than the IB:
[Equation]
1.1) where e y i is the one-hot vector where index y i ∈ { 1 , . . .
K } is positive.
The relation in Eq.
(3.1.1) follows immediately from the complete data likelihoods (2.3.2) and (2.2.3), along with the observation that C SDO i,y i ( β ) produces values between 0 and 1.
It is always valid to provide values from the domain of the SDO model to the IB model.
Next, we argue by monotonicity of the integral, that the expected complete likelihood of SDO is bounded by the corresponding expectation under IB for all observations i :
[Equation]
This left hand side is a critical additive piece of the
SDO ELBO, and the right hand side is a tractable surrogate bound with no troublesome normalizing constant C SDO .
We therefore deﬁne a surrogate objective L SDO ( q ) that lower bounds SDO’s log marginal likelihood (Eq.
[Equation]
We call this a surrogate lower bound because there are two bounds at work: the traditional ELBO (via Jensen’s inequality) and the bound relating SDO to IB in Eq.
Our surrogate objective L SDO ( q ) can also be seen as exactly the traditional ELBO applied directly to the IB model.
Our bound argument for applying the IB model to maximize the evidence of categorical data in Eq.
(3.1.3) relies on crucial (but achievable) assumptions.
First, both SDO and IB models need to have the same prior density π over weights (which should be Normal for conjugacy under the IB model).
Second, that our approximate posterior Q with density q can conceptually generate the unknown variables for either SDO or IB.
This is possible because given the same categorical data y , these distributions will have the same support.
Relation between DO and IB.
The arguments above that justify the IB ELBO as a lower bound for the SDO ELBO would naturally hold for the original DO-Probit model as well, not just our simpliﬁed model, because the DO normalization constant is also guaranteed to be between 0 and 1.
However, the posterior visualizations in Fig.
1 suggest the SDO is a better target model at least for large N in the intercepts-only regime, as the IB’s regions of high posterior density have more overlap with the SDO than the DO.
3.2 ALGORITHM FOR POSTERIOR ESTIMATION
Via the surrogate bound relation established in the previous section, we can provably optimize our intended categorical
SDO model by instead maximizing the traditional ELBO of the independent binary model.
This argument holds for any selected approximate posterior Q over z , β .
In this section, we develop an algorithm for coordinate ascent variational inference (CAVI) that can beneﬁt from the conjugacy of the IB model.
We assume that Q has a mean-ﬁeld factorization: q ( z , β )
=
q ( z ) q ( β ) .
Under these choices of algorithm and factorization, deriving the steps of the algorithm follows a standard general recipe [Blei et al., 2017] and has been previously applied to binary probit models [Consonni and Marin, 2007, Armagan and Zaretzki, 2011, Fasano et al., 2021].
We can further generalize this strategy to any model M whose joint distribution over all visible and unobserved variables (denoted u ) includes an SDO-Probit complete data likelihood (Eq.
(2.3.2)) for observed y and latent z as well as a conditionally-conjugate prior on weights β .
This generalization is possible since the lower bounding argument hinges solely on the expected complete data likelihood in
Eq.
Generalization allows us to handle more ﬂexible models such as the hierarchical SDO given in Section 4.
We summarize our VI strategy in Algorithm 1.
Algorithm
Take Q to be a mean-ﬁeld family with factorization: q ( u 1 , ..., u V ) =
(cid:81)
Vv =1 q v ( u v ) , where each u v is an unobserved variable ( z and β are included in this set).
Construct the surrogate model , M IB , by swapping the IBP likelihood (2.2.3) for the SDO likelihood and use it to compute surrogate complete conditionals : { log p M IB ( u
v | u − v , y Deﬁne the objective: L M ( q ) =
E q [log p M IB ( y , u ) q ( u ) ] .
Optimize L M ( q ) using optimal coordinate ascent updates [Blei et al., 2017]): q v ( u v ) ∝ exp (cid:8) E q − v (cid:2) log p M IB
( u v | u
− v ,
y ) (cid:3)(cid:9) .
If the complete conditional is an exponential family with natural parameter η v , so is its optimal update, with natural parameter given by
[Equation]
The updates in Algorithm 1 will yield a density q ∗ that is a local maximum of the ELBO of the surrogate model [Ormerod and Wand, 2010], and therefore a local maximum of a surrogate bound on our intended SDO categorical model.
All conditionals directly related to our SDO complete likelihood enjoy the closed-form updates given by Eq.
Note that the same strategy also applies to models with DO-Probit complete data likelihoods.
THE HIERARCHICAL SDO-PROBIT
We now illustrate the extensibility of our framework by introducing hierarchical SDO-Probit models, designed for datasets where observed categorical responses are nested within groups.
For instance, the model can handle collections of categorical time series (see Section 4.1).
Let our dataset consist of J groups (indexed by j ), each one containing N j units (indexed by i ).
We assume that each group is given a group-speciﬁc weight parameter β j , and that these are generated in a hierarchical fashion with a common prior:
[Equation]
Here, we introduce unknown location vectors and covariance matrices µ k , Σ k deﬁning regression weights across groups for each category k , which have standard Normal and Inverse Wishart priors.
Given group level parameters β j , we can combine these with the covariates x ij at each unit i within group j to produce categorical observations y ij ∈ { 1 , 2 , . . . K } :
[Equation]
Auxiliary variables z ij can be also introduced from an appropriate truncated Normal, following the SDO generative process (Eq. (2.3.2)).
Given this model, developing a coordinate ascent variational inference follows naturally from
Algorithm 1.
Details are worked out in Appendix C.
4.1 HIERARCHICAL AUTOREGRESSION
In order to model multiple related sequences of categorical observations over time, we develop a hierarchical multiclass autoregressive model as a special case of Eq.
In this case, the units (indexed by i ) are timesteps, and the groups (indexed by j ) are the sequences.
We simply need to adjust the covariates so that they include the previous timestep’s categorical outcome.
The dynamic covariates have the block structure
[Equation]
The ﬁrst block b ji is [1] if an intercept is used and empty otherwise.
The second block e y j,i − 1 indicates the previous category y j,i − 1 via a one-hot vector of length K .
The third block x exog ∈ R M e are exogenous dynamic covariates.
We can write M = 1 + K + M e , assuming an intercept is used.
Correspondingly, the regression weights decompose as
[Equation]
[Figure]
[Figure]
R e l a t i v e e rr o r r e du c t
i o n [Figure] Despite the complexities of this model, using our variational framework we can develop a coordinate ascent algorithm with closed-form updates .
EXPERIMENTS 5.1 PREDICTING CATEGORICAL RESPONSES WITH AN INDEPENDENT BINARY MODEL
Since the IB model via lower bound arguments approximates the posteriors of both the DO-Probit and SDO-Probit models, our goal is to assess in practice which model’s estimator performs better at predicting heldout data.
The methodology for data simulation and model ﬁtting is given in Sec.
D.1.
We consider synthetic datasets of
N = 4000 training examples, K = 3 categories, and M = 3 covariates drawn from different ground truth models.
2 Given a variational posterior obtained via Algorithm 1, we compute the posterior mean weights as a point estimate, ¯ β , and then make predictions with either the IB+DO (Eq. (2.6.1)) or IB+SDO (Eq. (2.6.2)) estimators.
To assess the relative performance of estimators, we report the relative reduction in error from using IB+SDO over
IB+DO as (¯ (cid:96) SDO − ¯ (cid:96) DO ) / (¯ (cid:96) DO − ¯ (cid:96) GT ) , where ¯ (cid:96) is the mean log probability of the observed category across observations in the heldout test set.
The ground truth likelihood, denoted GT, is obtained by plugging the known data-generating regression weights into the known model.
We report the spread of these error reduction values across 10 replicates of each data-generating process in Fig. 4.
This ﬁgure suggests that IB+SDO produces superior perfor- mance on a typical dataset, although IB+DO is superior in certain cases.
We thus recommend the IB+SDO estimator as a reasonable default.
We assess absolute performance by reporting the geometric means (cid:0) exp(¯ (cid:96) GT ) , exp(¯ (cid:96) SDO ) (cid:1) at the bottom of Figure 4, and ﬁnd that our variational scheme provides a good approximation of the ground truth, despite intentionally using a mis-speciﬁed model.
5.2 MODELING MULTIPLE CATEGORICAL TIME SERIES
Here we demonstrate how applying Algorithm 1 to the autoregressive models of Sec.
4 enables easy variational inference for collections of categorical time series.
We generated simulated data (as described in Sec.
D.2) withsequences, 5 categories, and 3 exogenous covariates.
We ﬁt a VI posterior approximating the hierarchical autoregressive SDO-Probit model of Sec.
Overall, we found that the hierarchical model provides a modest but consistent boost in heldout likelihood: the mean likelihood of test data increased from 59.9% to 62.5% with a minimum gain across groups of 1.3%. The hierarchical structure was particularly advantageous for predicting rare categories, as shown in Fig.
The hierarchical model shares statistical strength across time series, a capacity supported by our extensible framework.
DISCUSSION
In summary, our new variational methods provide a principled route to fast, scalable modeling of categorical data.
While our variational methods are most practically viewed as simply ﬁtting the conjugate binary model, we show this can be motivated via a rigorous lower bound on the marginal likelihood of a categorical model.
Future work could apply these methods to real datasets, investigate the tightness of the objective, and extend these methods to models with latent (unobserved) categorical variables. ACKNOWLEDGEMENTS This research was sponsored by the U.S. Army DEVCOM Soldier Center, and was accomplished under Cooperative
Agreement Number W911QY-19-2-0003.
The views and conclusions contained in this document are those of the authors and should not be interpreted as representing the ofﬁcial policies, either expressed or implied, of the U
S. Army DEVCOM Soldier Center, or the U S. Government. The U. S. Government is authorized to reproduce and distribute reprints for Government purposes notwithstanding any copyright notation hereon.
We also acknowledge support from the U.S. National Science Foundation under award HDR-1934553 for the Tufts
T-TRIPODS Institute.
MCH is supported in part by NSF
IIS-1908617.
A PRELIMINARIES A.1 THE INVERSE WISHART DISTRIBUTION
The Inverse Wishart is a distribution on symmetric, positive deﬁnite matrices.
Under a natural parametrization, the Inverse Wishart distribution, denoted W − 1 ( ν, Ψ ) , has density p [Equation]
[Equation]
where Σ (cid:31) 0 and ν > d − 1 to have a proper prior.
The expected value of an Inverse Wishart random variable parametrized as in (A.1.1) is given by E
[ Σ ] = Ψ ν − d − 1 .
The expected value of the precision matrix is E [ Σ − 1 ] =
ν Ψ − 1 .
A.2 UNIVARIATE NORMALS TRUNCATED TO POSITIVE OR NEGATIVE REALS
We will work with distributions truncated to the positive or negative reals, and so we deﬁne special notation: N + ( µ, σ 2 ) :
= T N ( µ, σ 2 , [0 , ∞ )) and N − ( µ, σ 2 ) :=
T N ( µ, σ 2 , ( −∞ , 0)) .
In particular, we will work with random variables of the form T + ∼ N + ( µ, 1) and T − ∼ N − ( µ, 1) .
Based on this construction, it is straightforward to derive f
[Equation]
E [ T + ] = µ + φ ( − µ ) 1 − Φ( − µ ) , E
[Equation]
Var [ T + ] = 1 − µ ( E [ T + ] − µ ) − ( E [ T + ] − µ )(A.2.2)
Var [ T − ] = 1 − µ ( E [ T − ] − µ ) − ( E [ T − ] − µ )(A.2.3) H [ T + ] = ln (cid:0) √ 2 πe [1 − Φ( − µ )] (cid:1) − µφ ( − µ ) 2(1 − Φ( − µ )) (A.2.4) H [ T − ] = ln (cid:0) √ 2 πe Φ( − µ ) (cid:1) + µφ ( − µ ) 2Φ( − µ ) (A.2.5) where we use φ and Φ to refer to the pdf and cdf, respectively, of the standard normal distribution, and where H [ X ] = − (cid:82) f ( x ) ln f ( x ) dx represents the differential entropy of a random variable X with density f .
Remark A.2.1.
( Representation in terms of perturbation of parent mean )
It is sometimes convenient to express the expectation of a truncated random variable as a perturbation of the expectation of its parent (pre-truncated)
Gaussian random variable.
To this end, for T s ∈ { T + , T − } , we write
[Equation]
[Equation] s ] [Equation]
which holds by (A.2.1).
[Equation]
B ANOTHER MOTIVATION FOR THE IB-SDO ESTIMATOR
Another way to motivate the IB+SDO estimator is to consider an intercepts-only version of the IB model.
Given a training set where the true frequency of category k is p k , a simple “moment-matching” argument suggests that a good point estimate of ¯ β should satisfy Φ(¯ β k )
≈ p k .
Thus, making predictions by computing Φ(¯ β k ) for each category index k and normalizing should given a decent estimate of the training set’s frequency of each category.
C VARIATIONAL INFERENCE FOR THE HIERARCHICAL SDO C.1 VARIATIONAL FAMILY
We take the mean-ﬁeld variational family for the Hierarchical SDO-probit model (4.0.1) to have density given by
[Equation]
(C.1.
[Equation]
where Ω ijk = (cid:40) R + , y ij = k
R − , otherwise (C.1.2)
Equality (1) is by mean-ﬁeld assumption, and, as we see below, (2) is the optimal such form as per Algorithm 1.
C.2 SURROGATE COMPLETE CONDITIONALS
The hierarchical Bayesian DO-Probit model (4.0.1) has a surrogate model which is a hierarchical Bayesian linear regression on the auxilliary variables z .
In this way, we obtain the surrogate complete conditionals
[Equation]
[Equation]
where N + and N − are truncated normal distributions deﬁned in
Section A.2, and
[Equation]
and
[Equation]
and
[Equation]
C.3 COORDINATE ASCENT UPDATES
All of the surrogate complete conditionals are in the exponential family, and hence we know from Algorithm 1 that the optimal variational factors with respect to the surrogate lower bound L H-SDO are in the same exponential family, with parameters given by (3.2.1).
Here we derive the parameters for the updates, using the notation of (C.1.1).
Updates to { q ( β jk ) } jk
Since the natural parameters of a multivariate Gaussian are the precision and precisionweighted mean, we reparametrize the surrogate complete conditional for each β jk in (C.2.2) before taking variational expectations of the parameters.
Hence, the optimal update to each q ( β jk | (cid:101) µ jk , (cid:101) Σ jk ) with respect to the objective L H-SDO is given by
[Equation]
where E q [ z jk ] is given explicitly below.
Equations (1), (2) follow from Sec.
A.
[Equation]
Thus, in standard parameterization [Equation] we update
[Equation]
where E q [ z jk ] ∈ R N j has i -th entry given by E q [ z ijk ] =
[Equation]
[Equation]
otherwise by properties of the truncated normal distribution (Section
A.2).
Recall that φ and Φ refer to the pdf and cdf, respectively, of the standard normal.
Updates to { q ( µ k ) } k
Since the natural parameters of a multivariate Gaussian are the precision and precisionweighted mean, we reparametrize the surrogate complete conditionals for µ k in (C.2.4) before taking variational expectations of the parameters.
Hence, the optimal update to q ( µ k | (cid:102) m k , (cid:101) V k ) with respect to the objective L H-SDO is given by
[Equation]
Equations (1), (2) follow from Sec.
A.1.
Thus
[Equation]
in standard parameterization [Equation] we update
[Equation]
Updates to { q ( Σ k ) } k
Since the natural parameters of an
Inverse Wishart are identical to the conventional parameters (up to a multiplicative scalar constant) [Equation] we do not need to reparametrize the surrogate complete conditional for Σ k in (C.2.3) before taking variational expectations of the parameters.
Hence [Equation] the optimal update to q ( Σ k | (cid:101) S k , (cid:101) ν k ) with respect to the objective L H-SDO is given by
[Equation]
where in (+) we used
[Equation]
and
[Equation]
with E q [ µ k β Tjk ] handled similarly as E q [ β jk µ Tk ] .
Updates to { q ( z ijk ) } ijk
In (C.2.1) , we saw that the surrogate complete conditional for each z ijk has the form T N ( η ijk , 1 , Ω ijk ) , where Ω ijk is deﬁned as in (C.1.2).
But since each such distribution is in the exponential family with natural parameter η ijk [Equation] the optimal update for each q ( z ijk ) is given by
[Equation]
C.4 LOWER BOUND
A tractable surrogate lower bound on the marginal log likeli- hood for the Hierarchical DO-Probit model (4.0.1) – useful for tracking the monotonic increase in the objective function and setting a convergence criterion – can immediately be obtained by handling the expected complete data likelihood precisely as in Section 3.
[Equation]
We expand this lower bound using:
[Equation]
j = N j (cid:88) i = K (cid:88) k
=
H [ q ( z ijk )] (cid:124) (cid:123)(cid:122) (cid:125) (B) + K (cid:88) k
= − KL (cid:0) q ( µ k ) (cid:12)(cid:12)(cid:12)(cid:12) p H-SDO ( µ k ) (cid:1) (cid:124) (cid:123)(cid:122) (cid:125) (C) +
K (cid:88) k = − KL (cid:0) q ( Σ k ) (cid:12)(cid:12)(cid:12)(cid:12) p H-SDO ( Σ k ) (cid:1) (cid:124) (cid:123)(cid:122) (cid:125)
(D) + J (cid:88) j
=
K (cid:88) k = − E q (cid:20) KL (cid:0) q ( β jk ) (cid:12)(cid:12)(cid:12)(cid:12) p H-SDO ( β jk | µ k , Σ k ) (cid:1)(cid:21) (cid:124) (cid:123)(cid:122) (cid:125) (E) .
In particular term (A) is given by
[-E q [ log p IB ( ¯ y ijk , z ijk | βββ jk )] = − 1 2 ( log2 π + 1 )+ 1 2 (cid:101) η ijk δ ¯ y ijk ( (cid:101) η ijk ) − 1 
 xxx-][Equation]
where (cid:101) η ijk = x Tij (cid:101) µ k and δ b ( x ) :=  
[Equation]
− [Equation] ) , b =
Term (B) is the entropy of a truncated normal distribution, where the ( ijk ) -th element in the sum is
• the entropy of N + ( x Tij (cid:101) µ k , 1) when y ij =
k , in which case the entropy is given by (A.2.4) • the entropy of N − ( x Tij (cid:101) µ k , 1) when y ij (cid:54) =
k , in which case the entropy is given by (A.2.5).
Term C is the KL divergence between two multivariate Gaussians; in this case we have
[Equation]
Term D is the KL divergence between two Inverse Wisharts; in this case, we have
[Equation]
where Γ M is the multivariate gamma function and ψ is the digamma function given by ψ ( x ) :=
ddx ln Γ 1 ( x ) .
Term E gives the expected KL divergence between two multivariate Gaussians, in the speciﬁc case where the expectation is taken with respect to independent Gaussian and
Inverse Wishart distributions on the parameters of the second argument of the KL divergence.
Utilizing the mean-ﬁeld assumption, we ﬁnd
[Equation]
D METHODOLOGY FOR EXPERIMENTS
All models in both experiments were run to a convergence criterion satisﬁed when consecutive iterations produced a drop of 1.0 or less in the surrogate lower bound on the log-evidence.
D.1 METHODOLOGY FOR EXPERIMENT
We generated a toy dataset of iid categorical observations as follows.
We randomly sampled regression weights β k iid ∼ N M ( 0 , I ) for each of k = 1 , ..., K categories ( K = 3 ) and M = 4 covariates.
For each of N = 5 , 000 samples we sampled exogenous covariates x exog i ∼ N 3 ( 0 , I ) and added an intercept to obtain x i = (1 , x exog i ) .
Given the covariates and regression weights, we then sampled categorical observations y i ∈ { 1 , ..., K } according to three different generative models: the DO-Probit, SDO-Probit, and Multinomial Logit model.
We performed variational inference using Algorithm 1 for the SDO-Probit model.
We trained on the ﬁrst 4,000 observations and tested on the remaining 1,000.
D.2 METHODOLOGY FOR EXPERIMENT
We generated a toy collection of categorical sequences as follows.
We assumed K = 5 categorical responses.
For covariates, we assumed B = 1 intercept term, K (ﬁrst-order) autoregessive terms, and M e = 3 exogenous covariates; thus, each x ji had M = B + K + M e = 9 elements.
We assumed no interactions.
We generated data for J = 10 sequences ( T train j ≡ 50 training timesteps and T test j ≡ 950 testing timesteps) via the generating process
[Table]
y
[Table]
= k |
[Table]
j according to (
[Table]
.4) where µ k is the k th column of the matrix µ ∈ R M × K .
Note that µ discourages the appearance of the 1st category, but that the rare binary covariate x ji 1 is highly predictive of its appearance.
We performed inference with a hierarchical DO-Probit autoregression (4.0.1) (see Section 4.1), as well as a collection of separate ﬂat DO-Probit autoregressions (2.1.1), one for each group.
After training, we computed the likelihood (i.e. model probabilities) of categorical observations in a hold-out test set.
Category probabilities were computed using (2.1.4) with E q [ β | y ] .