Save the date! Google I/O returns May 18-20 Register now

Random noise generation in TFF

This tutorial will discuss the recommendation for random noise generation in TFF. Random noise generation is an important component of many privacy protection techniques in federated learning algorithms, e.g., differential privacy.

View on TensorFlow.org Run in Google Colab View source on GitHub Download notebook

Before we begin

First, let us make sure the notebook is connected to a backend that has the relevant components compiled.

!pip install --quiet --upgrade tensorflow_federated_nightly
!pip install --quiet --upgrade nest_asyncio

import nest_asyncio
nest_asyncio.apply()
import numpy as np
import tensorflow as tf
import tensorflow_federated as tff

Run the following "Hello World" example to make sure the TFF environment is correctly setup. If it doesn't work, please refer to the Installation guide for instructions.

@tff.federated_computation
def hello_world():
  return 'Hello, World!'

hello_world()
b'Hello, World!'

Discouraged usage: directly using tf.random.normal

TF1.x like APIs tf.random.normal for random noise generation are strongly discouraged in TF2 according to the random noise generation tutorial in TF. Surprising behavior may happen when these APIs are used together with tf.function and tf.random.set_seed. For example, the following code will generate the same value with each call. This surprising behavior is expected for TF, and explanation can be found in the documentation of tf.random.set_seed.

tf.random.set_seed(1)

@tf.function
def return_one_noise(_):
  return tf.random.normal([])

n1=return_one_noise(1)
n2=return_one_noise(2) 
assert n1 == n2
print(n1.numpy(), n2.numpy())
0.3052047 0.3052047

In TFF, things are slightly different. If we wrap the noise generation as tff.tf_computation instead of tf.function, non-deterministic random noise will be generated. However, if we run this code snippet multiple times, different set of (n1, n2) will be generated each time. There is no easy way to set a global random seed for TFF.

tf.random.set_seed(1)

@tff.tf_computation
def return_one_noise(_):
  return tf.random.normal([])

n1=return_one_noise(1)
n2=return_one_noise(2) 
assert n1 != n2
print(n1, n2)
1.3283143 0.45740178

Moreover, deterministic noise can be generated in TFF without explicitly setting a seed. The function return_two_noise in the following code snippet returns two identical noise values. This is expected behavior because TFF will build computation graph in advance before execution. However, this suggests users have to pay attention on the usage of tf.random.normal in TFF.

@tff.tf_computation
def tff_return_one_noise():
  return tf.random.normal([])

@tff.federated_computation
def return_two_noise():
  return (tff_return_one_noise(), tff_return_one_noise())

n1, n2=return_two_noise() 
assert n1 == n2
print(n1, n2)
-0.15665223 -0.15665223

Usage with care: tf.random.Generator

We can use tf.random.Generator as suggested in the TF tutorial.

@tff.tf_computation
def tff_return_one_noise(i):
  g=tf.random.Generator.from_seed(i)
  @tf.function
  def tf_return_one_noise():
    return g.normal([])
  return tf_return_one_noise()

@tff.federated_computation
def return_two_noise():
  return (tff_return_one_noise(1), tff_return_one_noise(2))

n1, n2 = return_two_noise() 
assert n1 != n2
print(n1, n2)
0.3052047 -0.38260338

However, users may have to be careful on its usage

In general, TFF prefers functional operations and we will showcase the usage of tf.random.stateless_* functions in the following sections.

In TFF for federated learning, we often work with nested structures instead of scalars and the previous code snippet can be naturally extended to nested structures.

@tff.tf_computation
def tff_return_one_noise(i):
  g=tf.random.Generator.from_seed(i)
  weights = [
         tf.ones([2, 2], dtype=tf.float32),
         tf.constant([2], dtype=tf.float32)
     ]
  @tf.function
  def tf_return_one_noise():
    return tf.nest.map_structure(lambda x: g.normal(tf.shape(x)), weights)
  return tf_return_one_noise()

@tff.federated_computation
def return_two_noise():
  return (tff_return_one_noise(1), tff_return_one_noise(2))

n1, n2 = return_two_noise() 
assert n1[1] != n2[1]
print('n1', n1)
print('n2', n2)
n1 [array([[0.3052047 , 0.5671378 ],
       [0.41852272, 0.2326421 ]], dtype=float32), array([1.1675092], dtype=float32)]
n2 [array([[-0.38260338, -0.47804865],
       [-0.5187485 , -1.8471988 ]], dtype=float32), array([-0.77835274], dtype=float32)]

A general recommendation in TFF is to use the functional tf.random.stateless_* functions for random noise generation. These functions take seed as an explicit input argument to generate random noise. We first define a helper class to maintain the seed as pseudo state. The helper RandomSeedGenerator has functional operators in a state-in-state-out fashion. It is reasonable to use a counter as pseudo state for tf.random.stateless_* as these functions scramble the seed before using it to make noises generated by correlated seeds statistically uncorrelated.

class RandomSeedGenerator():

  def initialize(self, seed=None):
    if seed is None:
      return tf.cast(tf.stack(
          [tf.math.floor(tf.timestamp()*1e6),
            tf.math.floor(tf.math.log(tf.timestamp()*1e6))]), dtype=tf.int64)
    else:
      return tf.constant(self.seed, dtype=tf.int64, shape=(2,))

  def next(self, state):
    return state + 1

  def structure_next(self, state, nest_structure):
    "Returns seed in nexted structure and the next state seed."
    flat_structure = tf.nest.flatten(nest_structure)
    flat_seeds = [state+i for i in range(len(flat_structure))]
    nest_seeds = tf.nest.pack_sequence_as(nest_structure, flat_seeds)
    return nest_seeds, flat_seeds[-1] + 1

Now let us use the helper class and tf.random.stateless_normal to generate (nested structure of) random noise in TFF. The following code snippet looks a lot like a TFF iterative process, see simple_fedavg as an example of expressing federated learning algorithm as TFF iterative process. The pseudo seed state here for random noise generation is tf.Tensor that can be easily transported in TFF and TF functions.

@tff.tf_computation
def tff_return_one_noise(seed_state):
  g=RandomSeedGenerator()
  weights = [
         tf.ones([2, 2], dtype=tf.float32),
         tf.constant([2], dtype=tf.float32)
     ]
  @tf.function
  def tf_return_one_noise():
    nest_seeds, updated_state = g.structure_next(seed_state, weights)
    nest_noise = tf.nest.map_structure(lambda x,s: tf.random.stateless_normal(
        shape=tf.shape(x), seed=s), weights, nest_seeds)
    return nest_noise, updated_state
  return tf_return_one_noise()

@tff.tf_computation
def tff_init_state():
  g=RandomSeedGenerator()
  return g.initialize()

@tff.federated_computation
def return_two_noise():
  seed_state = tff_init_state()
  n1, seed_state = tff_return_one_noise(seed_state)
  n2, seed_state = tff_return_one_noise(seed_state)
  return (n1, n2)

n1, n2 = return_two_noise() 
assert n1[1] != n2[1]
print('n1', n1)
print('n2', n2)
n1 [array([[-0.21598858, -0.30700883],
       [ 0.7562299 , -0.21218438]], dtype=float32), array([-1.0359321], dtype=float32)]
n2 [array([[ 1.0722181 ,  0.81287116],
       [-0.7140338 ,  0.5896157 ]], dtype=float32), array([0.44190162], dtype=float32)]