Minimize $f(A,B)$ s.t. $\text{exp}(A)^T \text{exp}(B)=J_K$

Question

I have a function $f(A,B)$ that maps a pair of two (tall) matrices, $A$ and $B$, to a scalar cost that I want to minimize. $A$ and $B$ both have $K$ columns. I also want to impose a set of equality constraints, s.t. $\text{exp}(A)^T \text{exp}(B)=J_K$, where $J_K$ is the $K\times K$ matrix of ones. In other words, for every pair of columns in $\text{exp}(A)$ and $\text{exp}(B)$, the dot product of these columns must equal 1.

How would you approach this problem? I've been looking into different kinds of constrained optimization techniques but I haven't managed to get a good sense of which (if any) would be the most promising for my setting.

Edit: To clarify, the $\text{exp}$ function is applied entry-wise to these matrices. An example of a solution that satisfies the constraints is: $$ A = \log \left( \frac{1}{12} \begin{bmatrix} 11 & 1\\ 11 & 1\\ 1 & 11\\ 1 & 11 \end{bmatrix} \right), \text{ } B = \log \left( \frac{1}{12} \begin{bmatrix} 11 & 1\\ 1 & 11\\ 11 & 1\\ 1 & 11 \end{bmatrix} \right) $$ (where the $\log$ function is also applied element-wise.)

If it helps, you can assume that $f$ has the following form: $$ f(A,B) = -(\sum_i {\vec{x}^T}^{(i)} A\vec{y}^{(i)} + {\vec{x}^T}^{(i)} B\vec{z}^{(i)} ) $$ where $\vec{x}$, $\vec{y}$, and $\vec{z}$ are all strictly non-negative column vectors, which happen have the further property that $\vec{x}^T\vec{1}=\vec{y}^T\vec{1}=\vec{z}^T\vec{1}=1$ (i.e. the sum of the entries of each vector equals 1 - not sure if this is relevant but I'll mention it just in case it is). $i$ indexes separate instances of these vectors that we sum over.

Answer 1

I'm posting some work I've done on your problem, this is not a full answer but I think it almost covers it all.

Loss function

$f$, as you wrote it, is linear. This is good enough, but that equation can be developed a bit to simplify the problem. Let's define two new matrices:

$$ W_A= \sum_i \vec{x}_{(i)}\vec{y}^T_{(i)}; \qquad W_B= \sum_i \vec{x}_{(i)}\vec{z}^T_{(i)}. $$

$W_A$ and $W_B$ have the same shape of $A$ and $B$ and provide the weight for each of their elements, for computing $f$. Actually, the utility function is just the weighted sum of all elements in $A$ and $B$.

Parametrization

Let's call $A' = exp(A)$ and $B' = exp(B)$. You can of course optimize $A$ and $B$, with $f$ strictly linear, but a very complex bounded domain, or optimize $A'$ and $B'$, with the addictional constrain of positiveness, but still a simpler domain (linearly bounded). Of course from $A'$ and $B'$ you can immediately find $A$ and $B$.

I would rather do the latter: even if this way $f$ is not linear anymore, it's still trivially derivable. A quick google search for a viable solver brought me here.

However...

If you multiply any row of $A'$ to some (strictly positive) value $k$, and you divide the corresponding row of $B'$ by the same value, $A'^T B'$ is unchanged (constraints respected), but $f(A, B)$ may do change, and if it does it changes linearly with $k$, because of the properties of logarithm. This means that $k$ either is undetermined or diverges.