dp_ftrl/blogpost_supplemental_privacy_accounting.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vJoJh5iC95j4"
      },
      "source": [
        "This colab is a supplement to the [Federated Learning with Formal Differential Privacy Guarantees](http://ai.googleblog.com/2022/02/federated-learning-with-formal.html) blogpost, providing details on the parameters we used for training the Gboard Spanish-language next-word-prediction model, and the \n",
        "corresponding $\\rho$-zCDP (zero concentrated differential privacy) and $(\\varepsilon,\\delta)$-DP (differential privacy) guarantee for the DP-FTRL algorithm. "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AGSWIM1Qmb8a"
      },
      "source": [
        "# Code Locations\n",
        "---\n",
        "\n",
        "Our core algorithm is available in open-source; the routines for [estimating cumulative sums using tree aggregation](https://github.com/tensorflow/privacy/blob/master/tensorflow_privacy/privacy/dp_query/tree_aggregation_query.py) with DP are released in TensorFlow Privacy. These are then integrated with TensorFlow Federated's [aggregation libraries](https://www.tensorflow.org/federated/api_docs/python/tff/aggregators/DifferentiallyPrivateFactory?version=nightly#tree_aggregation), which allow them to be plugged into different learning algorithms, in particular [Federated Averaging](https://www.tensorflow.org/federated/api_docs/python/tff/learning/build_federated_averaging_process).\n",
        "\n",
        "The [code to perform privacy accounting](https://github.com/tensorflow/privacy/blob/master/tensorflow_privacy/privacy/analysis/tree_aggregation_accountant.py) (which we reproduce below to make this colab stand alone) can also be found in TensorFlow Privacy."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xg--bb-NtzI5"
      },
      "source": [
        "# Notion of differential privacy (DP)\n",
        "---\n",
        "The DP analysis covers the use of the Gboard training data cache for the training of a production model from random initialization using DP-FTRL.\n",
        "\n",
        "In this work, we conform to the add/remove notion of DP (where neighboring data sets differ by addition/removal of a single client). In the absence of a client at any training step, we assume that the client's model update gets replaced with the all zeros vector. This assumption enforces a subtle modification to the traditional definition of the add/remove notion of DP which allows neighboring data sets to have the same number of records. The formal definition is provided in Definition 1.1 in [the paper](https://arxiv.org/abs/2103.00039). It is a special instantiation of Definition II.3 (removal DP) in [Erlingsson et al.](https://arxiv.org/abs/2001.03618).\n",
        "\n",
        "Our privacy guarantee holds for all well-behaved clients (that is, clients that faithfully follow the algorithm including participation limits).  One thing to emphasize is that due to the design of the algorithm, a mis-behaved client does not adversely affect the DP guarantee of any well-behaved clients.\n",
        "\n",
        "The notion of adjacency is with respect to arbitrary training datasets on each client device (i.e., the device removed might have an aribtrarily large local dataset containing arbitrary training examples). For user's with a single device, this corresponds directly to user-level DP; for devices shared with multiple users, this provides a stronger notion of DP than user-level; for a user with say 2 devices that happen to both participate in training the model, the notion is weaker, but group privacy can be used to obtain a user-level guarantee.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IG0xshkrP4X4"
      },
      "source": [
        "# Differential privacy analysis for the DP-FTRL algorithm \n",
        "---\n",
        "Here, we instantiate the DP-FTRL algorithm without tree-restarts from the [arxiv version v3](https://arxiv.org/abs/2103.00039) of our ICML 2021 paper. The privacy analysis is from Appendix D.2 from the arxiv version. \n",
        "\n",
        "The main idea in the algorithm is the following:\n",
        "\n",
        "Let $T$ be the number of rounds (`total_steps` in the code below) of the server-side training process, and in each round $t \\in [T]$, we compute $\\nabla_t$, which is the sum of gradients over `COHORT_SIZE` distinct clients who contributed at round $t$, at the current model state $\\theta_t$. Now, we create a forest of binary trees as follows: \n",
        "\n",
        "   1. Create a complete binary tree with $T'$ leaf nodes, where $T'$ equals to smallest power of two $\\geq T$. \n",
        "   1. Set the first $T$ leaf nodes (from left) to be $\\nabla_1,\\ldots,\\nabla_T$ respectively, and delete the remaining leaf nodes. \n",
        "   1. Delete all the internal nodes $z$ which are not complete binary trees rooted at $z$ (that is, delete all $x$ where any leaf of the subtree rooted at $x$ is not present/has been deleted). \n",
        "   1. Set each node $z$ of this forest to be the sum of all the leaves in the subtree rooted at $z$. Call the resulting forest $\\cal T$.\n",
        "\n",
        "In this formulation, each client is allowed to participate *only once* in the computation of any $\\nabla_t$, and a maximum of $E$  times (`max_participation` in the code below). Furthermore, each client is allowed to participate *once in every 24-hours*. We encode this constraint in the privacy accounting via $\\xi$ (`min_separation` in the code below), which is the minimum number of rounds in between two successive participation of any clients. Recall, the following: \n",
        "\n",
        "  * Each node $z\\in\\cal T$ stores a sum of all the $\\nabla_t$'s in the leaves of the sub-tree rooted at $z$. \n",
        "  * Each leaf node in $\\cal T$ only affects values at the nodes in the path path from it to the root of $\\cal T$. \n",
        "\n",
        "Let $L$ be the clipping norm of each individual client gradient participating in any of the $\\nabla_t$'s. For `noise multiplier` $\\sigma$, the noise that gets added to each node in the tree $\\cal T$ is sampled from ${\\cal N}\\left(0, \\sigma^2 L^2\\right)^{p}$, where $p$ is the dimensionality of the model update. Let $\\zeta^*$ bound the squared $\\ell_2$-sensitivity of the tree $\\cal T$: Consider an empty forest with the same structure of $\\cal T$, call it $\\widehat{\\cal T}$. Let $c\\in\\{0,1\\}^T$ be any bit vector with $\\|c\\|_0\\leq E$ (with $E$ being the maximum number of participations per client), and any two successive ones have at least $\\xi$ zeros in between them. Let $\\widehat{\\cal T}(c)$ be an instantiation of $\\widehat{\\cal T}$ with the leaf nodes being $c$, and each internal node being the sum of all the leaves in the sub-tree rooted at that node. We define $\\zeta^*=\\max\\limits_c\\sum_{z\\in\\widehat{\\cal T}(c)}z^2$. In Theorem D.3 [in the paper](https://arxiv.org/abs/2103.00039) (and in function `_tree_sensitivity_square_sum`), we provide a dynamic programming based approach to calculate an upper bound on $\\zeta^*$. \n",
        "\n",
        "Since our algorithm is primarily based on the Gaussian mechanism, we can easily provide $\\rho$-zCDP guarantee, along with $(\\epsilon,\\delta)$-DP guarantee (with $\\delta$ being `target_delta` in the code below). The bound for zCDP is immediate, and taken care of by function `compute_zcdp`. It is a direct consequence of Lemma 2.4 of [Bun and Steinke'16](https://arxiv.org/pdf/1605.02065.pdf) and Theorem D.3 [in our paper](https://arxiv.org/abs/2103.00039)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mpqb2PciimnG"
      },
      "source": [
        "# General purposing accounting code\n",
        "\n",
        "The code in this section is a replica of the open source version of privacy accounting\n",
        "[here](https://github.com/tensorflow/privacy/blob/master/tensorflow_privacy/privacy/analysis/tree_aggregation_accountant.py), with the purpose of ease of verifiability. The `_tree_sensitivity_square_sum` function here computes the $\\zeta^*$ mentioned earlier."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EhGIMJJVP7qm"
      },
      "outputs": [],
      "source": [
        "!pip install tensorflow_privacy\n",
        "import itertools\n",
        "import time\n",
        "import warnings\n",
        "\n",
        "from absl import app\n",
        "from absl import flags\n",
        "\n",
        "import six\n",
        "from typing import Collection, Union, Dict, Tuple\n",
        "import numpy as np\n",
        "import math\n",
        "import collections\n",
        "import tensorflow_privacy as tfp\n",
        "\n",
        "# RDP orders to consider\n",
        "ORDERS = [1 + x / 10. for x in range(1, 100)] + list(range(12, 64))\n",
        "\n",
        "def _check_nonnegative(value: Union[int, float], name: str):\n",
        "  if value \u003c 0:\n",
        "    raise ValueError(f\"Provided {name} must be non-negative, got {value}\")\n",
        "\n",
        "\n",
        "def _check_possible_tree_participation(num_participation: int,\n",
        "                                       min_separation: int, start: int,\n",
        "                                       end: int, steps: int) -\u003e bool:\n",
        "  \"\"\"Check if participation is possible with `min_separation` in `steps`.\n",
        "\n",
        "  This function checks if it is possible for a sample to appear\n",
        "  `num_participation` in `steps`, assuming there are at least `min_separation`\n",
        "  nodes between the appearance of the same sample in the streaming data (leaf\n",
        "  nodes in tree aggregation). The first appearance of the sample is after\n",
        "  `start` steps, and the sample won't appear in the `end` steps after the given\n",
        "  `steps`.\n",
        "\n",
        "  Args:\n",
        "    num_participation: The number of times a sample will appear.\n",
        "    min_separation: The minimum number of nodes between two appearance of a\n",
        "      sample. If a sample appears in consecutive x, y steps in a streaming\n",
        "      setting, then `min_separation=y-x-1`.\n",
        "    start:  The first appearance of the sample is after `start` steps.\n",
        "    end: The sample won't appear in the `end` steps after the given `steps`.\n",
        "    steps: Total number of steps (leaf nodes in tree aggregation).\n",
        "\n",
        "  Returns:\n",
        "    True if a sample can appear `num_participation` with given conditions.\n",
        "  \"\"\"\n",
        "  return start + (min_separation + 1) * num_participation \u003c= steps + end\n",
        "\n",
        "\n",
        "def _tree_sensitivity_square_sum(\n",
        "    num_participation: int, min_separation: int, start: int, end: int,\n",
        "    steps: int, hist_buffer: Dict[Tuple[int, int, int, int], float]) -\u003e float:\n",
        "  \"\"\"Compute the worst-case sum of sensitivtiy square for `num_participation`.\n",
        "\n",
        "  This is the key algorithm for DP accounting for DP-FTRL tree aggregation\n",
        "  without restart, which recurrently counts the worst-case occurence of a sample\n",
        "  in all the nodes in a tree. This implements a dynamic programming algorithm\n",
        "  that exhausts the possible `num_participation` appearance of a sample in\n",
        "  `steps` leaf nodes. See Appendix D of\n",
        "  \"Practical and Private (Deep) Learning without Sampling or Shuffling\"\n",
        "  https://arxiv.org/abs/2103.00039.\n",
        "\n",
        "  Args:\n",
        "    num_participation: The number of times a sample will appear.\n",
        "    min_separation: The minimum number of nodes between two appearance of a\n",
        "      sample. If a sample appears in consecutive x, y steps in a streaming\n",
        "      setting, then `min_separation=y-x-1`.\n",
        "    start:  The first appearance of the sample is after `start` steps.\n",
        "    end: The sample won't appear in the `end` steps after the given `steps`.\n",
        "    steps: Total number of steps (leaf nodes in tree aggregation).\n",
        "    hist_buffer: A dictionary stores the worst-case sum of sesentivity square\n",
        "      keyed by (num_participation, start, end, steps).\n",
        "\n",
        "  Returns:\n",
        "    The worst-case sum of sesentivity square for the given input.\n",
        "  \"\"\"\n",
        "  key_tuple = (num_participation, start, end, steps)\n",
        "  if key_tuple in hist_buffer:\n",
        "    return hist_buffer[key_tuple]\n",
        "  if not _check_possible_tree_participation(num_participation, min_separation,\n",
        "                                            start, end, steps):\n",
        "    sum_value = -np.inf\n",
        "  elif num_participation == 0:\n",
        "    sum_value = 0.\n",
        "  elif num_participation == 1 and steps == 1:\n",
        "    sum_value = 1.\n",
        "  else:\n",
        "    steps_log2 = math.log2(steps)\n",
        "    max_2power = math.floor(steps_log2)\n",
        "    if max_2power == steps_log2:\n",
        "      sum_value = num_participation**2\n",
        "      max_2power -= 1\n",
        "    else:\n",
        "      sum_value = 0.\n",
        "    candidate_sum = []\n",
        "    for right_part in range(num_participation + 1):\n",
        "      for right_start in range(min_separation + 1):\n",
        "        left_sum = _tree_sensitivity_square_sum(\n",
        "            num_participation=num_participation - right_part,\n",
        "            min_separation=min_separation,\n",
        "            start=start,\n",
        "            end=right_start,\n",
        "            steps=2**max_2power,\n",
        "            hist_buffer=hist_buffer)\n",
        "        if np.isinf(left_sum):\n",
        "          candidate_sum.append(-np.inf)\n",
        "          continue  # Early pruning for dynamic programming\n",
        "        right_sum = _tree_sensitivity_square_sum(\n",
        "            num_participation=right_part,\n",
        "            min_separation=min_separation,\n",
        "            start=right_start,\n",
        "            end=end,\n",
        "            steps=steps - 2**max_2power,\n",
        "            hist_buffer=hist_buffer)\n",
        "        candidate_sum.append(left_sum + right_sum)\n",
        "    sum_value += max(candidate_sum)\n",
        "  hist_buffer[key_tuple] = sum_value\n",
        "  return sum_value\n",
        "\n",
        "\n",
        "def _max_tree_sensitivity_square_sum(max_participation: int,\n",
        "                                     min_separation: int, steps: int) -\u003e float:\n",
        "  \"\"\"Compute the worst-case sum of sensitivtiy square in tree aggregation.\n",
        "\n",
        "  See Appendix D of\n",
        "  \"Practical and Private (Deep) Learning without Sampling or Shuffling\"\n",
        "  https://arxiv.org/abs/2103.00039.\n",
        "\n",
        "  Args:\n",
        "    max_participation: The maximum number of times a sample will appear.\n",
        "    min_separation: The minimum number of nodes between two appearance of a\n",
        "      sample. If a sample appears in consecutive x, y steps in a streaming\n",
        "      setting, then `min_separation=y-x-1`.\n",
        "    steps: Total number of steps (leaf nodes in tree aggregation).\n",
        "\n",
        "  Returns:\n",
        "    The worst-case sum of sesentivity square for the given input.\n",
        "  \"\"\"\n",
        "  num_participation = max_participation\n",
        "  while not _check_possible_tree_participation(\n",
        "      num_participation, min_separation, 0, min_separation, steps):\n",
        "    num_participation -= 1\n",
        "  candidate_sum, hist_buffer = [], collections.OrderedDict()\n",
        "  for num_part in range(1, num_participation + 1):\n",
        "    candidate_sum.append(\n",
        "        _tree_sensitivity_square_sum(num_part, min_separation, 0,\n",
        "                                     min_separation, steps, hist_buffer))\n",
        "  return max(candidate_sum)\n",
        "\n",
        "\n",
        "def _compute_gaussian_rdp(sigma: float, sum_sensitivity_square: float,\n",
        "                          alpha: float) -\u003e float:\n",
        "  \"\"\"Computes RDP of Gaussian mechanism.\"\"\"\n",
        "  if np.isinf(alpha):\n",
        "    return np.inf\n",
        "  return alpha * sum_sensitivity_square / (2 * sigma**2)\n",
        "\n",
        "\n",
        "def compute_rdp_single_tree(\n",
        "    noise_multiplier: float, total_steps: int, max_participation: int,\n",
        "    min_separation: int,\n",
        "    orders: Union[float, Collection[float]]=ORDERS) -\u003e Union[float, Collection[float]]:\n",
        "  \"\"\"Computes RDP of the Tree Aggregation Protocol for a single tree.\n",
        "\n",
        "  The accounting assume a single tree is constructed for `total_steps` leaf\n",
        "  nodes, where the same sample will appear at most `max_participation` times,\n",
        "  and there are at least `min_separation` nodes between two appearance. The key\n",
        "  idea is to (recurrently) count the worst-case occurence of a sample\n",
        "  in all the nodes in a tree, which implements a dynamic programming algorithm\n",
        "  that exhausts the possible `num_participation` appearance of a sample in\n",
        "  `steps` leaf nodes.\n",
        "\n",
        "  See Appendix D of\n",
        "  \"Practical and Private (Deep) Learning without Sampling or Shuffling\"\n",
        "  https://arxiv.org/abs/2103.00039.\n",
        "\n",
        "  Args:\n",
        "    noise_multiplier: A non-negative float representing the ratio of the\n",
        "      standard deviation of the Gaussian noise to the l2-sensitivity of a single\n",
        "      contribution (a leaf node), which is usually set in\n",
        "      `TreeCumulativeSumQuery` and `TreeResidualSumQuery` from\n",
        "      `dp_query.tree_aggregation_query`.\n",
        "    total_steps: Total number of steps (leaf nodes in tree aggregation).\n",
        "    max_participation: The maximum number of times a sample can appear.\n",
        "    min_separation: The minimum number of nodes between two appearance of a\n",
        "      sample. If a sample appears in consecutive x, y steps in a streaming\n",
        "      setting, then `min_separation=y-x-1`.\n",
        "    orders: An array (or a scalar) of RDP orders.\n",
        "\n",
        "  Returns: The RDPs at all orders. Can be `np.inf`.\n",
        "  \"\"\"\n",
        "  _check_nonnegative(noise_multiplier, \"noise_multiplier\")\n",
        "  if noise_multiplier == 0:\n",
        "    return np.inf\n",
        "  _check_nonnegative(total_steps, \"total_steps\")\n",
        "  _check_nonnegative(max_participation, \"max_participation\")\n",
        "  _check_nonnegative(min_separation, \"min_separation\")\n",
        "  sum_sensitivity_square = _max_tree_sensitivity_square_sum(\n",
        "      max_participation, min_separation, total_steps)\n",
        "  if np.isscalar(orders):\n",
        "    rdp = _compute_gaussian_rdp(noise_multiplier, sum_sensitivity_square,\n",
        "                                orders)\n",
        "  else:\n",
        "    rdp = np.array([\n",
        "        _compute_gaussian_rdp(noise_multiplier, sum_sensitivity_square, alpha)\n",
        "        for alpha in orders\n",
        "    ])\n",
        "  return rdp, sum_sensitivity_square\n",
        "\n",
        "def compute_zcdp(noise_multiplier, sensitivity_sq):\n",
        "  \"\"\"Computes zCDP of the Tree Aggregation Protocol for a single tree, using\n",
        "     Lemma 2.4 from https://arxiv.org/pdf/1605.02065.pdf.\n",
        "  Args:\n",
        "    noise_multiplier: A non-negative float representing the ratio of the\n",
        "      standard deviation of the Gaussian noise to the l2-sensitivity of a single\n",
        "      contribution (a leaf node), which is usually set in\n",
        "      `TreeCumulativeSumQuery` and `TreeResidualSumQuery` from\n",
        "      `dp_query.tree_aggregation_query`.\n",
        "    sensitivity_sq: The sum of squared sensitivity of the nodes in the binary tree, assuming \n",
        "      each minibatch has \\ell_2 sensitivity of one.\n",
        "  Returns: zCDP parameter.\n",
        "  \"\"\"\n",
        "  return sensitivity_sq / (2 * pow(noise_multiplier, 2))\n",
        "\n",
        "def eps_from_rdp(rdp, target_delta):\n",
        "  \"\"\"Compute epsilon for (eps, target_delta)-DP from rdp.\"\"\"\n",
        "  return tfp.get_privacy_spent(\n",
        "          ORDERS, rdp, target_delta=target_delta)[0]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CvvO7Y16QB9w"
      },
      "source": [
        "# Application to the training of a production Gboard language model\n",
        "\n",
        "\n",
        "We can now plug the parameters used in the production training run into this privacy accounting code. The actual DP-training procedure via DP-FTRL aggregates client updates via the following code path: \n",
        "\n",
        "```python\n",
        "NOISE_MULTIPLIER = 7.0\n",
        "COHORT_SIZE = 6500\n",
        "tff.aggregators.DifferentiallyPrivateFactory.tree_aggregation(\n",
        "        noise_multiplier=NOISE_MULTIPLIER,\n",
        "        clients_per_round=COHORT_SIZE,       \n",
        "        use_efficient=True,\n",
        "        ...)\n",
        "```\n",
        "\n",
        "Under the hood, this uses the [`TreeResidualSumQuery.tree_aggregation`](https://github.com/tensorflow/federated/blob/139a123d6918631fc6604df67a0a2dc58c971d40/tensorflow_federated/python/aggregators/differential_privacy.py#L222) which estimates per-round model updates. This is achieved via post-processing of the tree aggregation method: we estimate the sum of updates on round $t$ by taking the difference between the private estimates of the cumulative sums after round $t$ and $t-1$. Finally, for optimization we want an average update, which is provided (again via post-processing) by a [` tfp.NormalizedQuery`](https://www.tensorflow.org/responsible_ai/privacy/api_docs/python/tf_privacy/NormalizedQuery) parameterized by `COHORT_SIZE`. The privacy calculations are independent of the constant `COHORT_SIZE`, but utility depends on (roughly) `COHORT_SIZE` clients contributing to each round.\n",
        "\n",
        "Finally, the parameters `min_separation`, `max_participation`, and `total_steps` are computed post-facto after the training completes. As each device is configured to participate in training at most once per 24 hours, `min_separation` is computed by taking the smallest number of training rounds completed in any 24 hour period, and `max_participation` is the maximum number of times any device can participate in `total_steps` training rounds while satisfying the participation constraint. \n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pjgDuJDnP-pI",
        "outputId": "dda9ab78-197c-4228-fe53-d836d49030db"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Accounting time 347.85 secs\n",
            "zCDP = 0.81\n",
            "(epsilon, delta) = (8.90, 1e-10)\n"
          ]
        }
      ],
      "source": [
        "noise_multiplier=7.0\n",
        "min_separation=313\n",
        "max_participation=6\n",
        "total_steps=2000\n",
        "target_delta=1e-10\n",
        "\n",
        "start_time = time.time()\n",
        "\n",
        "rdp, sensitivity_sq = compute_rdp_single_tree(\n",
        "    noise_multiplier, total_steps, \n",
        "    max_participation, min_separation)\n",
        "eps = eps_from_rdp(rdp, target_delta)\n",
        "zcdp = compute_zcdp(noise_multiplier, sensitivity_sq)\n",
        "print(f'Accounting time {time.time()-start_time:.2f} secs')\n",
        "print(f'zCDP = {zcdp:.2f}')\n",
        "print(f'(epsilon, delta) = ({eps:.2f}, {target_delta:.2g})') "
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "blogpost_supplemental_privacy_accounting.ipynb",
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}