A “softmax”-like function for deciding on a partition

Softmax can be derived as follows. Say that we are given $$k$$ "log priors" $$b_i$$ that our data belongs to the $$i$$th category in some categorical distribution. Then we can solve for the category probabilities $$p_i$$ from the relations: $\sum_i p_i = 1 \qquad\text{and}\qquad p_i\propto e^{b_i}$ by solving for $$Z$$ below: $p_i = \frac{1}{Z} e^{b_i} \qquad\implies\qquad Z = \sum_i e^{b_i}.$ We can take this as a definition of softmax: $p = \operatorname{softmax}(b).$ is the function that builds a categorical distribution from log likelihoods.

OK, now let's change contexts a little bit. Say that we have now a square $$(k\times k$$) matrix of log priors $$B_{ij}$$, which we want to use to build a joint probability distribution $$P_{ij}$$ over choices of two categories $$(i,j)$$. So, we must turn $$B$$ into a doubly stochastic matrix $$P$$. Additionally, say that we don't care about the order of the categories, so naturally we should have symmetric matrices $$B_{ij}=B_{ji}$$ and $$P_{ij}=P_{ji}$$ for all $$i,j$$. In this case, there exists a generalization to softmax which at least some authors refer to as "softassign". First, let $E_{ij}=\exp(B_{ij}).$ Then run Sinkhorn's iterations on the matrix $$E$$:

• $$P'\gets E$$
• until "convergence" do:
• $$P'_{ij} \gets P'_{ij} / \sum_{j} P'_{ij}$$
• $$P'_{ij} \gets P'_{ij} / \sum_{i} P'_{ij}$$
• output P = P'.

This turns $$E$$ into a doubly stochastic matrix $$P$$ by some coordinate descent, alternately normalizing rows and columns. We use this to define the softassign function: $P = \operatorname{softassign}(B).$

OK, now here is my question. Let's interpret $$P_{ij}$$ as the probability that entities $$i$$ and $$j$$ belong to the same partition of $$\{1,\dots,k\}$$. Then $$P=I$$ would indicate that the partition is $$\{\{1\},\{2\},\dots,\{k\}\}$$, but another matrix would decrease the number of sets in the partition. Since $$P$$ now determines an equivalence relation, the symmetry constraint is natural.

But it would be nice to also enforce the transitivity constraint of equivalence relations. If we pretend for a second that $$P$$ is binary, then transitivity would mean that $$P_{ij}=P_{jk}\implies P_{ik}=P_{ij}=P_{jk}$$. A similar statement if we relax to $$P_{ij}\in[0,1]$$ would maybe be that if $$P_{ij}$$ and $$P_{jk}$$ are very close, then $$P_{ik}$$ should also be close to that value. Alternately, if $$P_{ij}$$ and $$P_{jk}$$ are both small (large), then $$P_{ik}$$ should be too.

So, my question is if there is anything like a Sinkhorn result for matrices with this sort of transitivity property that be used to augment softassign in this context? I would also be interested in improvements to the way I'm phrasing transitivity or alternate interpretations of this kind of matrix, since it feels a little off, or in any other notes.

A PS: It occurs to me now that Sinkhorn does not enforce symmetry, but it does preserve it, which is lucky for me since of course my log priors were symmetric (like many log priors, they are something like similarities or negative distances). But maybe with noisy priors, there is some alternative to Sinkhorn that imposes symmetry, and could be extended to impose transitivity? Going to look into a Lagrange multipliers approach, but I definitely need to work harder on formalizing transitivity for it to make sense in that language.