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$.

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.