Constrained optimization of joint Bernoulli density parametrized by a neural network

Question

I want to learn a joint distribution on $n$ Bernoulli random variables conditioned on a random variable $b\sim D$ and parametrized by a neural network, $f$:

$$p(A = (a_1, \ldots, a_n)|b) = f_{\Lambda}(b).$$

I require that $\langle A \rangle = \frac{1}{n}\sum_{i=1}^n E[a_i] = p_0$ where $p_0$ is some fixed (typically small) number. Each $A$ is evaluated by some loss function $L(A,b)$. I have tried solving the regularized problem

$$\text{minimize} \ E_{b\sim D}[L(A,b; \Lambda) + \eta (\langle A\rangle - p_0)^2]$$

with backprop, but I find it is difficult to tune $\eta$. Is there some way to impose a hard constraint?

Answer 1

You could have the network parameterize just the first $n-1$ bernoulli RVs and set $p_{n} = np_0 - \sum_i^{n-1} p_i$ defining ($p_i = E[a_i]$). If $p_n$ isn't in $[0,1]$ then "back up" one more step, and set $p_{n-1} = p_n = np_0 - \sum_i^{n-2} p_i$. If this still doesn't keep $p_n$ in $[0,1]$, keep "backing up" until it does.