How to manipulate client gradients in tensorflow federated sgd

Question

I'm following this tutorial to get started with tensorflow federated. My aim is to run federated sgd (not federated avg) with some manipulations on client gradient values before they are sent to the server.

Before moving forward, to briefly reiterate the federated sgd process, for each turn clients will send their computed gradients (not updated weights) to the server, the server aggregates them and broadcasts the updated model to the clients.

Now from what I've gathered so far, I can use the function build_federated_sgd_process instead of build_federated_averaging_process in the mentioned tutorial to perform federated sgd the way described above.

Where I'm lost is, I need to clip the client gradients and add some noise to them (independently generated for each gradient value) before sending the gradients to the server and I'm not sure how to do it. Generating the noise is straightforward enough, but which function should I modify/implement to be able to do apply the noise to the gradients?

Answer 1

build_federated_sgd_process is fully-canned; it is really designed to serve as a reference implementation, not as a point of extensibility.

I believe what you are looking for is the function that build_federated_sgd_process calls under the hoos, tff.learning.framework.build_model_delta_optimizer_process. This function allows you to supply your own mapping from a model function (IE, a zero-arg callable that returns a tff.learning.Model) to a tff.learning.framework.ClientDeltaFn.

Your ClientDeltaFn would look something like:

@tf.function
def _clip_and_noise(grads):
  return ...

class ClippedGradClientDeltaFn(tff.learning.framework.ClientDeltaFn)

def __init__(self, model, ...):
  self._model = model
  ...

@tf.function
def __call__(dataset, weights):
  # Compute gradients grads
  return _clip_and_noise(grads)

And you would be able to construct a tff.templates.IterativeProcess by calling:

def clipped_sgd(model_fn: Callable[[], model_lib.Model]) -> ClippedGradClientDeltaFn:
    return ClippedGradClientDeltaFn(
        model_fn(),
        ...)

iterproc = optimizer_utils.build_model_delta_optimizer_process(
      model_fn, model_to_client_delta_fn=clipped_sgd, ...)

as more or less in the body of build_federated_sgd_process.

It sounds to me like you are interested in differential privacy; TFF is actually designed to compose with differential privacy generally through the aggregation processes rather than writing different client updates, though this is certainly one approach. See the pointers from the TFF for research documentation for idiomatic ways to wire differential privacy in to TFF.