TFF 中的隨機雜訊生成

本教學課程將討論 TFF 中隨機雜訊生成的建議最佳做法。隨機雜訊生成是 Federated Learning 演算法中許多隱私保護技術的重要組成部分,例如差分隱私。

在 TensorFlow.org 上檢視 在 Google Colab 中執行 在 GitHub 上檢視原始碼 下載筆記本

開始之前

首先,請確定筆記本已連線至具有相關組件的後端。

pip install --quiet --upgrade tensorflow-federated
import numpy as np
import tensorflow as tf
import tensorflow_federated as tff

執行下列「Hello World」範例,以確保 TFF 環境設定正確。如果無法運作,請參閱安裝指南以取得操作說明。

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

hello_world()
b'Hello, World!'

用戶端上的隨機雜訊

用戶端上對雜訊的需求通常分為兩種情況:相同的雜訊和 i.i.d. 雜訊。

  • 對於相同的雜訊,建議的模式是在伺服器上維護種子值,將其廣播至用戶端,並使用 tf.random.stateless 函式產生雜訊。
  • 對於 i.i.d. 雜訊,請使用在用戶端上以 from_non_deterministic_state 初始化的 tf.random.Generator,這與 TF 避免使用 tf.random.<distribution> 函式的建議一致。

用戶端行為與伺服器不同 (不會遇到稍後討論的陷阱),因為每個用戶端都會建構自己的運算圖並初始化自己的預設種子值。

用戶端上相同的雜訊

# Set to use 10 clients.
tff.backends.native.set_sync_local_cpp_execution_context(default_num_clients=10)

@tff.tensorflow.computation
def noise_from_seed(seed):
  return tf.random.stateless_normal((), seed=seed)

seed_type_at_server = tff.FederatedType(tff.to_type((np.int64, [2])), tff.SERVER)

@tff.federated_computation(seed_type_at_server)
def get_random_min_and_max_deterministic(seed):
  # Broadcast seed to all clients.
  seed_on_clients = tff.federated_broadcast(seed)

  # Clients generate noise from seed deterministicly.
  noise_on_clients = tff.federated_map(noise_from_seed, seed_on_clients)

  # Aggregate and return the min and max of the values generated on clients.
  min = tff.federated_min(noise_on_clients)
  max = tff.federated_max(noise_on_clients)
  return min, max

seed = tf.constant([1, 1], dtype=tf.int64)
min, max = get_random_min_and_max_deterministic(seed)
assert min == max
print(f'Seed: {seed.numpy()}. All clients sampled value {min:8.3f}.')

seed += 1
min, max = get_random_min_and_max_deterministic(seed)
assert min == max
print(f'Seed: {seed.numpy()}. All clients sampled value {min:8.3f}.')
Seed: [1 1]. All clients sampled value    1.665.
Seed: [2 2]. All clients sampled value   -0.219.

用戶端上獨立的雜訊

@tff.tensorflow.computation
def nondeterministic_noise():
  gen = tf.random.Generator.from_non_deterministic_state()
  return gen.normal(())

@tff.federated_computation
def get_random_min_and_max_nondeterministic():
  noise_on_clients = tff.federated_eval(nondeterministic_noise, tff.CLIENTS)
  min = tff.federated_min(noise_on_clients)
  max = tff.federated_max(noise_on_clients)
  return min, max

min, max = get_random_min_and_max_nondeterministic()
assert min != max
print(f'Values differ across clients. {min:8.3f},{max:8.3f}.')

new_min, new_max = get_random_min_and_max_nondeterministic()
assert new_min != new_max
assert new_min != min and new_max != max
print(f'Values differ across rounds.  {new_min:8.3f},{new_max:8.3f}.')
Values differ across clients.   -1.490,   1.172.
Values differ across rounds.    -1.358,   1.208.

用戶端上的模型初始化器

def _keras_model():
  inputs = tf.keras.Input(shape=(1,))
  outputs = tf.keras.layers.Dense(1)(inputs)
  return tf.keras.Model(inputs=inputs, outputs=outputs)

@tff.tensorflow.computation
def tff_return_model_init():
  model = _keras_model()
  # return the initialized single weight value of the dense layer
  return tf.reshape(
      tff.learning.models.ModelWeights.from_model(model).trainable[0], [-1])[0]

@tff.federated_computation
def get_random_min_and_max_nondeterministic():
  noise_on_clients = tff.federated_eval(tff_return_model_init, tff.CLIENTS)
  min = tff.federated_min(noise_on_clients)
  max = tff.federated_max(noise_on_clients)
  return min, max

min, max = get_random_min_and_max_nondeterministic()
assert min != max
print(f'Values differ across clients. {min:8.3f},{max:8.3f}.')

new_min, new_max = get_random_min_and_max_nondeterministic()
assert new_min != new_max
assert new_min != min and new_max != max
print(f'Values differ across rounds.  {new_min:8.3f},{new_max:8.3f}.')
Values differ across clients.   -1.022,   1.567.
Values differ across rounds.    -1.675,   1.550.

伺服器上的隨機雜訊

不建議的使用方式:直接使用 tf.random.normal

根據 TF 中的隨機雜訊生成教學課程,TF2 強烈建議不要使用 TF1.x 類似的 API (例如 tf.random.normal) 來產生隨機雜訊。當這些 API 與 tf.functiontf.random.set_seed 一起使用時,可能會發生令人意外的行為。例如,下列程式碼每次呼叫都會產生相同的值。這種令人意外的行為在 TF 中是預期的,您可以在 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

在 TFF 中,情況略有不同。如果我們將雜訊生成包裝為 tff.tensorflow.computation 而不是 tf.function,則會產生非決定性的隨機雜訊。但是,如果我們多次執行此程式碼片段,則每次都會產生不同的 (n1, n2) 組合。沒有簡單的方法可以為 TFF 設定全域隨機種子值。

tf.random.set_seed(1)

@tff.tensorflow.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)
0.11990704 1.9185987

此外,即使未明確設定種子值,也可以在 TFF 中產生決定性雜訊。下列程式碼片段中的函式 return_two_noise 會傳回兩個相同的雜訊值。這是預期的行為,因為 TFF 會在執行前預先建構運算圖。但是,這表示使用者必須注意在 TFF 中使用 tf.random.normal 的方式。

謹慎使用:tf.random.Generator

我們可以依照 TF 教學課程中的建議使用 tf.random.Generator

@tff.tensorflow.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.38260335

但是,使用者可能必須謹慎使用

一般而言,TFF 偏好函數式運算,我們將在以下章節中展示 tf.random.stateless_* 函式的使用方式。

在用於 Federated Learning 的 TFF 中,我們通常使用巢狀結構而不是純量,而先前的程式碼片段可以自然地擴展到巢狀結構。

@tff.tensorflow.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.38260335, -0.4780486 ],
       [-0.5187485 , -1.8471988 ]], dtype=float32), array([-0.77835274], dtype=float32)]

TFF 中的一般建議是使用函數式 tf.random.stateless_* 函式來產生隨機雜訊。這些函式採用種子值 (形狀為 [2] 的張量或兩個純量張量的 tuple) 作為明確的輸入引數,以產生隨機雜訊。我們先定義一個輔助程式類別,以將種子值維護為虛擬狀態。輔助程式 RandomSeedGenerator 具有以狀態輸入-狀態輸出的方式運作的函數式運算子。使用計數器作為 tf.random.stateless_* 的虛擬狀態是合理的,因為這些函式會在 scramble 種子值,然後再使用它來使相關種子值產生的雜訊在統計上不相關。

def timestamp_seed():
  # tf.timestamp returns microseconds as decimal places, thus scaling by 1e6.
  return tf.cast(tf.timestamp() * 1e6, tf.int64)

class RandomSeedGenerator():

  def initialize(self, seed=None):
    if seed is None:
      return tf.stack([timestamp_seed(), 0])
    else:
      return tf.constant(self.seed, dtype=tf.int64, shape=(2,))

  def next(self, state):
    return state + tf.constant([0, 1], tf.int64)

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

現在讓我們使用輔助程式類別和 tf.random.stateless_normal 在 TFF 中產生 (巢狀結構的) 隨機雜訊。下列程式碼片段看起來很像 TFF 迭代程序,請參閱 simple_fedavg,以取得將 Federated Learning 演算法表示為 TFF 迭代程序的範例。此處用於隨機雜訊生成的虛擬種子值狀態是 tf.Tensor,可以在 TFF 和 TF 函式中輕鬆傳輸。

@tff.tensorflow.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.tensorflow.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.86828816,  0.8535084 ],
       [ 1.0053564 , -0.42096713]], dtype=float32), array([0.18048067], dtype=float32)]
n2 [array([[-1.1973879 , -0.2974589 ],
       [ 1.8309833 ,  0.17024393]], dtype=float32), array([0.68991095], dtype=float32)]