Stochastic Gradient Descent for iterated expectation?

Question

Normally, SGD comes up in the context like $$\min_\theta ~ \mathbb{E}(f(X, \theta))$$, where $\theta$ is some parameter, $f$ is a function like $f(X, \theta) = (X-\theta)^2$ (to find the mean), and one is given samples $x_i\sim X$.

However, consider the related problem $$\min_\theta ~\mathbb{E}\left( ~ \mathbb{E}(f(X, \theta)|Y) ~\right)$$, where the inner expectation is taken over $X$ conditioned on $Y$, and the outer expectation is taken over $Y$. Furthermore, the samples we have are of the form $y_j\sim Y$ and $x_{ij} \sim (X | Y=y_j)$. For example, we could have the samples

$y_1$, $x_{11}$, $x_{21}$, $x_{31}$

$y_2$, $x_{12}$

$y_3$, $x_{13}$, $x_{23}$

To motivate this, $Y$ could be the ID of a person, and $X$ could be whether they successfully make a free-throw. You want to know the average free-throw success rate $\mathbb{E}(X)$, but you're given lop-sided samples (some people give you 100 free-throw samples, some people only give you 6, etc.). What I don't want to happen is for 1 person who gave a lot of samples to disproportionately dominate things. So, in some sense, I want to find the average free-throw rate of each person $\mathbb{E}(X|Y)$, and THEN take the average of those, so that each person gets weighed equally, and doesn't dominate just because they gave more free-throws.

How should SGD be adjusted to account for this? What about batch SGD? I imagine the solution will have to reweight things based on sample sizes.

Furthermore, if I wanted to use something like Adam or Adadelta (which build on SGD), is there a way I can adjust gradients to trick them into working properly in this context?

Just as an example, assume that there are ~1 million $y_j$'s, and each one has between 100 and 10000 $x_{ij}$'s attached to it.

Answer 1

I think I have the answer, which works simply and can easily be adapted to batch Gradient decent, or Adam, etc.

All you do is, when passing through the data, you step through the $j$ index, and choose the $i$ index randomly over all the available $x_{ij}$'s for $y_j$.

So for the original question's sample set, you might go through the samples as follows:

$x_{11}, x_{12}, x_{13}, x_{21}, x_{12}, x_{23}, x_{31}, x_{12}, x_{13}$

Note how, for $j=1$, each of $x_{i1}$ gets used once (because there are lots of them), whereas for $j=2$, we have $x_{12}$ getting used 3 times (because there's only 1 of it).

Doing things like this has the effect of reweighting things properly so each $j$ is given equal weight.