Setting local model updates (Ditto Federated learning)

I am currently working on implementing the Ditto personalized federated learning algorithm over human activity recognition data (HAR), however I am struggling to update local model parameters/weights. Without the personlization step in my code, local models have ~55% accuracy while after personalization accuracy seems almost random and very low ~7.

Local updates are computed by → v_k = v_k - lr (SGD(v_k) + interpolation(v_k - global_weights) )
If interpolation is 0 → local training is just SGD

I’m currently stuck on line six. I want to want to select new parameter values for my local model based on SGD and the difference between global model weights and original local model weights.

Currently I run one step of SGD and then slightly modify the model parameters by lr*interpolation(local_weights - global weights). I think this might be affecting my optimizer. Is the only way to implement line six of the algorithms to create a custom optimizer?

I also have a picture of the algorithm →

Maybe it too late but, I think you implement it in the wrong way

_, grads, _ = c.solve_sgd(data_batch)
...

for layer in range(len(grads[1])):
    eff_grad = grads[1][layer] + self.lam * (self.local_models[idx][layer] - self.global_model[layer])
    self.local_models[idx][layer] = self.local_models[idx][layer] - self.learning_rate * eff_grad

check Ditto source code if you want to implement it from the ground up: ditto/ditto.py at master · litian96/ditto · GitHub