You should keep the gradient multiplied by the sample_weight even if you change the sampling of idx, otherwise you are not computing the correct loss gradient.

The loss with $\ell_2$ regularizer to minimize wrt $w$=weights coefficient of the linear model (skipping the intercept for now):
$loss(w) = \frac{1}{S} \sum_i s_i l(y_i, w^Tx_i) + \frac{\alpha}{2} \Vert w \Vert^2$
can be put in several ways as the finite sum minimized by sag(a):
$g(w) = \frac{1}{n} \sum_i f_i(w) \propto loss(w)$

Let $l_i(w) = l(y_i, w^Tx_i)$ be the data $i$ loss term. We can choose:

1. $f_i(w) = s_i l_i(w) + s_i \frac{\alpha}{2} \Vert w \Vert^2 $. The gradient updated and stored in memory is then update = sample_weight[idx] * (entry * gradient + alpha * weights)
2. $f_i(w) = s_i l_i(w) + \frac{S}{n} \frac{\alpha}{2} \Vert w \Vert^2$. This gives the code in main where only the data $i$ loss term is weighted.
3. $f_i(w) = s_i l_i(w)$. The gradient updated and stored in memory is then update = sample_weight[idx] * entry * gradient, ie only the data $i$ loss term. The regularizer is then taken into account when doing the gradient descent step on weights, for example weights -= alpha*step_size*weights.

I feel 2. (the current code) is a bit weird, for instance there is maybe a weird scaling of alpha as $\frac{S}{n}\alpha$. However all three options lead to the same updated weights at the end, so it shouldn't matter too much.

I think I have a slight preference for option 1. with the meaning of adding a term proportional to $s_i$, in which case it would make sense to redefine n_seen as the weighted sum of seen indices, which will converge to $S$ instead of $n$.

You should keep the gradient multiplied by the sample_weight even if you change the sampling of idx, otherwise you are not computing the correct loss gradient.

I am not entirely sure about this, in the case of sampling with frequency weighted probabilities, is this not equivalent to uniformly sampling from the repeated set of data in which case we can maintain the update step as is? Not sure if I have misunderstood something here?

In the above case it seems like we are considering weights to be more of the form of precision weights no?

I think the confusion comes from the comparison with SGD.

Let's step back. For the weighted version consider a function $f(w) = \frac{1}{S} \sum_i s_i f_i(w)$ to minimize.

The (FG) full gradient step is $w_{k+1} = w_k - \gamma f'(w_k)$
where the correct gradient is $f'(w_k) = \frac{1}{S} \sum_i s_i f'_i(w)$

SG (stochastic gradient) and SAG (stochastic average gradient) are two stochastic methods to approximate the FG but in different ways.

SG: the descent step uses the gradient at a randomly selected data point. In the unweighted case, you must sample uniformly the data points. Otherwise the gradient (averaged over some iterations) would be biased and your descent does not converge toward the right minima. In the weighted case therefore you have two choices: either you sample uniformly $i$ and multiply the gradient by the sample weight

$w_{k+1} = w_k - \gamma s_i f'_i(w_k)$

or you sample $j$ proportionally to the sample weights but not multiply the gradient:

$w_{k+1} = w_k - \gamma f'_j(w_k)$

Notice that, in expectation, and averaged over some iterations, the gradient approximates the FG $f'(w) = \frac{1}{S} \sum_i s_i f'_i(w)$. If you were to use the sample weight two times, you would get a biased estimate of the FG.

SAG(A): the descent step uses an approximation of the average gradient $f'(w) = \frac{1}{S} \sum_i s_i \nabla f_i(w)$. The trick is to store individual gradients in memory, denoted $f'_i(\phi_i^k)$ or $y_i^k$ in the sag(a) papers. The individual gradient is updated for the randomly selected data point $i$. The descent step always uses the approximate average gradient $\frac{1}{S} \sum_i s_i f'_i(\phi_i^k) \sim f'(w)$. So here you apply directly an approximation of the FG at each iteration, which is not biased, whatever the sampling used for the data point selection. The sampling only affects which part of the average gradient is most regularly updated.

This a crucial difference with SG. In the unweighted case for example, to quote the sag paper:

In standard SG methods, it is crucial to sample the functions fi uniformly, at least asymptotically, in order to yield an unbiased gradient estimate and subsequently achieve convergence to the optimal value (alternately, the bias induced by non-uniform sampling would need to be asymptotically corrected). In SAG iterations, however, the weight of each gradient is constant and equal to 1/n, regardless of the frequency at which the corresponding function is sampled. We might thus consider sampling the functions fi non-uniformly, without needing to correct for this bias.

So then a question is which sampling procedure to use ? See above and the sag(a) paper, but I think it involves looking at the Lipschitz smoothness constant.

I don't think removing the zero sample_weight is enough.

If I understood correctly the sag(a) paper the recommended step size is step_size = 1 / L where $L = \max_i L_i$ and $L_i$ is the Lipschitz smoothness constant of $f_i(w)$.

It will depend on how we decompose the weighted sum into individual $f_i(w)$ terms. With the options above:

1. $L_i = s_i \kappa \Vert x_i \Vert^2 + s_i \alpha$.
2. $L_i = s_i \kappa \Vert x_i \Vert^2 + \frac{S}{n} \alpha$.

where $\kappa = 1$ for regression and $\kappa = 1/4$ for classification.

In the end, it depends on how we allocate the regularizer:

• evenly wrt the indices, each $i$ has the same regularizer $\frac{S}{n} \frac{\alpha}{2}\Vert w \Vert^2$
• evenly wrt the sample_weight, each $i$ has a regularizer $s_i \frac{\alpha}{2}\Vert w \Vert^2$ proportional to its sample_weight.

Maybe we should try both strategies and see which one leads to better convergence ?

I think we first need to figure out the step size computation before trying non-uniform sampling, for example by following 5.5 Effect of non-uniform sampling of the sag paper.

Hmmmm so I guess you are suggesting a variable step size only for the special case of sample weights?

Sorry if I wasn't clear, I meant first find/test the correct formula for step_size with sample_weight (see options 1 and 2 above).

Later we can try nonuniform sampling, which I think would be more difficult than just proportional to the sample_weight, see the discussion in 5.5 around the Lipschitz constants.