flwr-nightly 1.21.0.dev20250902__py3-none-any.whl → 1.21.0.dev20250904__py3-none-any.whl

flwr/serverapp/strategy/dp_fixed_clipping.py

@@ -0,0 +1,352 @@
+# Copyright 2025 Flower Labs GmbH. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Message-based Central differential privacy with fixed clipping.
+
+Papers: https://arxiv.org/abs/1712.07557, https://arxiv.org/abs/1710.06963
+"""
+
+from abc import ABC
+from collections import OrderedDict
+from collections.abc import Iterable
+from logging import INFO, WARNING
+from typing import Optional
+
+from flwr.common import Array, ArrayRecord, ConfigRecord, Message, MetricRecord, log
+from flwr.common.differential_privacy import (
+    add_gaussian_noise_inplace,
+    compute_clip_model_update,
+    compute_stdv,
+)
+from flwr.common.differential_privacy_constants import (
+    CLIENTS_DISCREPANCY_WARNING,
+    KEY_CLIPPING_NORM,
+)
+from flwr.server import Grid
+
+from .strategy import Strategy
+
+
+class DifferentialPrivacyFixedClippingBase(Strategy, ABC):
+    """Base class for DP strategies with fixed clipping.
+
+    This class contains common functionality shared between server-side and
+    client-side fixed clipping implementations.
+
+    Parameters
+    ----------
+    strategy : Strategy
+        The strategy to which DP functionalities will be added by this wrapper.
+    noise_multiplier : float
+        The noise multiplier for the Gaussian mechanism for model updates.
+        A value of 1.0 or higher is recommended for strong privacy.
+    clipping_norm : float
+        The value of the clipping norm.
+    num_sampled_clients : int
+        The number of clients that are sampled on each round.
+    """
+
+    # pylint: disable=too-many-arguments,too-many-instance-attributes
+    def __init__(
+        self,
+        strategy: Strategy,
+        noise_multiplier: float,
+        clipping_norm: float,
+        num_sampled_clients: int,
+    ) -> None:
+        super().__init__()
+
+        self.strategy = strategy
+
+        if noise_multiplier < 0:
+            raise ValueError("The noise multiplier should be a non-negative value.")
+
+        if clipping_norm <= 0:
+            raise ValueError("The clipping norm should be a positive value.")
+
+        if num_sampled_clients <= 0:
+            raise ValueError(
+                "The number of sampled clients should be a positive value."
+            )
+
+        self.noise_multiplier = noise_multiplier
+        self.clipping_norm = clipping_norm
+        self.num_sampled_clients = num_sampled_clients
+
+    def _validate_replies(self, replies: Iterable[Message]) -> bool:
+        """Validate replies and log errors/warnings.
+
+        Returns
+        -------
+        bool
+            True if replies are valid for aggregation, False otherwise.
+        """
+        num_errors = 0
+        num_replies_with_content = 0
+        for msg in replies:
+            if msg.has_error():
+                log(
+                    INFO,
+                    "Received error in reply from node %d: %s",
+                    msg.metadata.src_node_id,
+                    msg.error,
+                )
+                num_errors += 1
+            else:
+                num_replies_with_content += 1
+
+        # Errors are not allowed
+        if num_errors:
+            log(
+                INFO,
+                "aggregate_train: Some clients reported errors. Skipping aggregation.",
+            )
+            return False
+
+        log(
+            INFO,
+            "aggregate_train: Received %s results and %s failures",
+            num_replies_with_content,
+            num_errors,
+        )
+
+        if num_replies_with_content != self.num_sampled_clients:
+            log(
+                WARNING,
+                CLIENTS_DISCREPANCY_WARNING,
+                num_replies_with_content,
+                self.num_sampled_clients,
+            )
+
+        return True
+
+    def _add_noise_to_aggregated_arrays(
+        self, aggregated_arrays: ArrayRecord
+    ) -> ArrayRecord:
+        """Add Gaussian noise to aggregated arrays.
+
+        Parameters
+        ----------
+        aggregated_arrays : ArrayRecord
+            The aggregated arrays to add noise to.
+
+        Returns
+        -------
+        ArrayRecord
+            The aggregated arrays with noise added.
+        """
+        aggregated_ndarrays = aggregated_arrays.to_numpy_ndarrays()
+        stdv = compute_stdv(
+            self.noise_multiplier, self.clipping_norm, self.num_sampled_clients
+        )
+        add_gaussian_noise_inplace(aggregated_ndarrays, stdv)
+
+        log(
+            INFO,
+            "aggregate_fit: central DP noise with %.4f stdev added",
+            stdv,
+        )
+
+        return ArrayRecord(
+            OrderedDict(
+                {
+                    k: Array(v)
+                    for k, v in zip(aggregated_arrays.keys(), aggregated_ndarrays)
+                }
+            )
+        )
+
+    def configure_evaluate(
+        self, server_round: int, arrays: ArrayRecord, config: ConfigRecord, grid: Grid
+    ) -> Iterable[Message]:
+        """Configure the next round of federated evaluation."""
+        return self.strategy.configure_evaluate(server_round, arrays, config, grid)
+
+    def aggregate_evaluate(
+        self,
+        server_round: int,
+        replies: Iterable[Message],
+    ) -> Optional[MetricRecord]:
+        """Aggregate MetricRecords in the received Messages."""
+        return self.strategy.aggregate_evaluate(server_round, replies)
+
+    def summary(self) -> None:
+        """Log summary configuration of the strategy."""
+        self.strategy.summary()
+
+
+class DifferentialPrivacyServerSideFixedClipping(DifferentialPrivacyFixedClippingBase):
+    """Strategy wrapper for central DP with server-side fixed clipping.
+
+    Parameters
+    ----------
+    strategy : Strategy
+        The strategy to which DP functionalities will be added by this wrapper.
+    noise_multiplier : float
+        The noise multiplier for the Gaussian mechanism for model updates.
+        A value of 1.0 or higher is recommended for strong privacy.
+    clipping_norm : float
+        The value of the clipping norm.
+    num_sampled_clients : int
+        The number of clients that are sampled on each round.
+
+    Examples
+    --------
+    Create a strategy::
+
+        strategy = fl.serverapp.FedAvg( ... )
+
+    Wrap the strategy with the `DifferentialPrivacyServerSideFixedClipping` wrapper::
+
+        dp_strategy = DifferentialPrivacyServerSideFixedClipping(
+            strategy, cfg.noise_multiplier, cfg.clipping_norm, cfg.num_sampled_clients
+        )
+    """
+
+    def __init__(
+        self,
+        strategy: Strategy,
+        noise_multiplier: float,
+        clipping_norm: float,
+        num_sampled_clients: int,
+    ) -> None:
+        super().__init__(strategy, noise_multiplier, clipping_norm, num_sampled_clients)
+        self.current_arrays: ArrayRecord = ArrayRecord()
+
+    def __repr__(self) -> str:
+        """Compute a string representation of the strategy."""
+        return "Differential Privacy Strategy Wrapper (Server-Side Fixed Clipping)"
+
+    def configure_train(
+        self, server_round: int, arrays: ArrayRecord, config: ConfigRecord, grid: Grid
+    ) -> Iterable[Message]:
+        """Configure the next round of training."""
+        self.current_arrays = arrays
+        return self.strategy.configure_train(server_round, arrays, config, grid)
+
+    def aggregate_train(
+        self,
+        server_round: int,
+        replies: Iterable[Message],
+    ) -> tuple[Optional[ArrayRecord], Optional[MetricRecord]]:
+        """Aggregate ArrayRecords and MetricRecords in the received Messages."""
+        if not self._validate_replies(replies):
+            return None, None
+
+        # Clip arrays in replies
+        current_ndarrays = self.current_arrays.to_numpy_ndarrays()
+        for reply in replies:
+            for arr_name, record in reply.content.array_records.items():
+                # Clip
+                reply_ndarrays = record.to_numpy_ndarrays()
+                compute_clip_model_update(
+                    param1=reply_ndarrays,
+                    param2=current_ndarrays,
+                    clipping_norm=self.clipping_norm,
+                )
+                # Replace content while preserving keys
+                reply.content[arr_name] = ArrayRecord(
+                    OrderedDict(
+                        {k: Array(v) for k, v in zip(record.keys(), reply_ndarrays)}
+                    )
+                )
+            log(
+                INFO,
+                "aggregate_fit: parameters are clipped by value: %.4f.",
+                self.clipping_norm,
+            )
+
+        # Pass the new parameters for aggregation
+        aggregated_arrays, aggregated_metrics = self.strategy.aggregate_train(
+            server_round, replies
+        )
+
+        # Add Gaussian noise to the aggregated arrays
+        if aggregated_arrays:
+            aggregated_arrays = self._add_noise_to_aggregated_arrays(aggregated_arrays)
+
+        return aggregated_arrays, aggregated_metrics
+
+
+class DifferentialPrivacyClientSideFixedClipping(DifferentialPrivacyFixedClippingBase):
+    """Strategy wrapper for central DP with client-side fixed clipping.
+
+    Use `fixedclipping_mod` modifier at the client side.
+
+    In comparison to `DifferentialPrivacyServerSideFixedClipping`,
+    which performs clipping on the server-side,
+    `DifferentialPrivacyClientSideFixedClipping` expects clipping to happen
+    on the client-side, usually by using the built-in `fixedclipping_mod`.
+
+    Parameters
+    ----------
+    strategy : Strategy
+        The strategy to which DP functionalities will be added by this wrapper.
+    noise_multiplier : float
+        The noise multiplier for the Gaussian mechanism for model updates.
+        A value of 1.0 or higher is recommended for strong privacy.
+    clipping_norm : float
+        The value of the clipping norm.
+    num_sampled_clients : int
+        The number of clients that are sampled on each round.
+
+    Examples
+    --------
+    Create a strategy::
+
+        strategy = fl.serverapp.FedAvg(...)
+
+    Wrap the strategy with the `DifferentialPrivacyClientSideFixedClipping` wrapper::
+
+        dp_strategy = DifferentialPrivacyClientSideFixedClipping(
+            strategy, cfg.noise_multiplier, cfg.clipping_norm, cfg.num_sampled_clients
+        )
+
+    On the client, add the `fixedclipping_mod` to the client-side mods::
+
+        app = fl.client.ClientApp(mods=[fixedclipping_mod])
+    """
+
+    def __repr__(self) -> str:
+        """Compute a string representation of the strategy."""
+        return "Differential Privacy Strategy Wrapper (Client-Side Fixed Clipping)"
+
+    def configure_train(
+        self, server_round: int, arrays: ArrayRecord, config: ConfigRecord, grid: Grid
+    ) -> Iterable[Message]:
+        """Configure the next round of training."""
+        # Inject clipping norm in config
+        config[KEY_CLIPPING_NORM] = self.clipping_norm
+        # Call parent method
+        return self.strategy.configure_train(server_round, arrays, config, grid)
+
+    def aggregate_train(
+        self,
+        server_round: int,
+        replies: Iterable[Message],
+    ) -> tuple[Optional[ArrayRecord], Optional[MetricRecord]]:
+        """Aggregate ArrayRecords and MetricRecords in the received Messages."""
+        if not self._validate_replies(replies):
+            return None, None
+
+        # Aggregate
+        aggregated_arrays, aggregated_metrics = self.strategy.aggregate_train(
+            server_round, replies
+        )
+
+        # Add Gaussian noise to the aggregated arrays
+        if aggregated_arrays:
+            aggregated_arrays = self._add_noise_to_aggregated_arrays(aggregated_arrays)
+
+        return aggregated_arrays, aggregated_metrics

flwr/serverapp/strategy/fedadagrad.py

@@ -0,0 +1,162 @@
+# Copyright 2025 Flower Labs GmbH. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""FedAdagrad [Reddi et al., 2020] strategy.
+
+Adaptive Federated Optimization using Adagrad.
+
+Paper: arxiv.org/abs/2003.00295
+"""
+
+from collections import OrderedDict
+from collections.abc import Iterable
+from typing import Callable, Optional
+
+import numpy as np
+
+from flwr.common import Array, ArrayRecord, Message, MetricRecord, RecordDict
+
+from .fedopt import FedOpt
+from .strategy_utils import AggregationError
+
+
+# pylint: disable=line-too-long
+class FedAdagrad(FedOpt):
+    """FedAdagrad strategy - Adaptive Federated Optimization using Adagrad.
+
+    Implementation based on https://arxiv.org/abs/2003.00295v5
+
+    Parameters
+    ----------
+    fraction_train : float (default: 1.0)
+        Fraction of nodes used during training. In case `min_train_nodes`
+        is larger than `fraction_train * total_connected_nodes`, `min_train_nodes`
+        will still be sampled.
+    fraction_evaluate : float (default: 1.0)
+        Fraction of nodes used during validation. In case `min_evaluate_nodes`
+        is larger than `fraction_evaluate * total_connected_nodes`,
+        `min_evaluate_nodes` will still be sampled.
+    min_train_nodes : int (default: 2)
+        Minimum number of nodes used during training.
+    min_evaluate_nodes : int (default: 2)
+        Minimum number of nodes used during validation.
+    min_available_nodes : int (default: 2)
+        Minimum number of total nodes in the system.
+    weighted_by_key : str (default: "num-examples")
+        The key within each MetricRecord whose value is used as the weight when
+        computing weighted averages for both ArrayRecords and MetricRecords.
+    arrayrecord_key : str (default: "arrays")
+        Key used to store the ArrayRecord when constructing Messages.
+    configrecord_key : str (default: "config")
+         Key used to store the ConfigRecord when constructing Messages.
+    train_metrics_aggr_fn : Optional[callable] (default: None)
+        Function with signature (list[RecordDict], str) -> MetricRecord,
+        used to aggregate MetricRecords from training round replies.
+        If `None`, defaults to `aggregate_metricrecords`, which performs a weighted
+        average using the provided weight factor key.
+    evaluate_metrics_aggr_fn : Optional[callable] (default: None)
+        Function with signature (list[RecordDict], str) -> MetricRecord,
+        used to aggregate MetricRecords from training round replies.
+        If `None`, defaults to `aggregate_metricrecords`, which performs a weighted
+        average using the provided weight factor key.
+    eta : float, optional
+        Server-side learning rate. Defaults to 1e-1.
+    eta_l : float, optional
+        Client-side learning rate. Defaults to 1e-1.
+    tau : float, optional
+        Controls the algorithm's degree of adaptability. Defaults to 1e-3.
+    """
+
+    # pylint: disable=too-many-arguments
+    def __init__(
+        self,
+        *,
+        fraction_train: float = 1.0,
+        fraction_evaluate: float = 1.0,
+        min_train_nodes: int = 2,
+        min_evaluate_nodes: int = 2,
+        min_available_nodes: int = 2,
+        weighted_by_key: str = "num-examples",
+        arrayrecord_key: str = "arrays",
+        configrecord_key: str = "config",
+        train_metrics_aggr_fn: Optional[
+            Callable[[list[RecordDict], str], MetricRecord]
+        ] = None,
+        evaluate_metrics_aggr_fn: Optional[
+            Callable[[list[RecordDict], str], MetricRecord]
+        ] = None,
+        eta: float = 1e-1,
+        eta_l: float = 1e-1,
+        tau: float = 1e-3,
+    ) -> None:
+        super().__init__(
+            fraction_train=fraction_train,
+            fraction_evaluate=fraction_evaluate,
+            min_train_nodes=min_train_nodes,
+            min_evaluate_nodes=min_evaluate_nodes,
+            min_available_nodes=min_available_nodes,
+            weighted_by_key=weighted_by_key,
+            arrayrecord_key=arrayrecord_key,
+            configrecord_key=configrecord_key,
+            train_metrics_aggr_fn=train_metrics_aggr_fn,
+            evaluate_metrics_aggr_fn=evaluate_metrics_aggr_fn,
+            eta=eta,
+            eta_l=eta_l,
+            beta_1=0.0,
+            beta_2=0.0,
+            tau=tau,
+        )
+
+    def aggregate_train(
+        self,
+        server_round: int,
+        replies: Iterable[Message],
+    ) -> tuple[Optional[ArrayRecord], Optional[MetricRecord]]:
+        """Aggregate ArrayRecords and MetricRecords in the received Messages."""
+        aggregated_arrayrecord, aggregated_metrics = super().aggregate_train(
+            server_round, replies
+        )
+
+        if aggregated_arrayrecord is None:
+            return aggregated_arrayrecord, aggregated_metrics
+
+        if self.current_arrays is None:
+            reason = (
+                "Current arrays not set. Ensure that `configure_train` has been "
+                "called before aggregation."
+            )
+            raise AggregationError(reason=reason)
+
+        # Compute intermediate variables
+        delta_t, m_t, aggregated_ndarrays = self._compute_deltat_and_mt(
+            aggregated_arrayrecord
+        )
+
+        # v_t
+        if not self.v_t:
+            self.v_t = {k: np.zeros_like(v) for k, v in aggregated_ndarrays.items()}
+        self.v_t = {k: v + (delta_t[k] ** 2) for k, v in self.v_t.items()}
+
+        new_arrays = {
+            k: x + self.eta * m_t[k] / (np.sqrt(self.v_t[k]) + self.tau)
+            for k, x in self.current_arrays.items()
+        }
+
+        # Update current arrays
+        self.current_arrays = new_arrays
+
+        return (
+            ArrayRecord(OrderedDict({k: Array(v) for k, v in new_arrays.items()})),
+            aggregated_metrics,
+        )

flwr/serverapp/strategy/fedadam.py

@@ -0,0 +1,181 @@
+# Copyright 2025 Flower Labs GmbH. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Adaptive Federated Optimization using Adam (FedAdam) strategy.
+
+[Reddi et al., 2020]
+
+Paper: arxiv.org/abs/2003.00295
+"""
+
+from collections import OrderedDict
+from collections.abc import Iterable
+from typing import Callable, Optional
+
+import numpy as np
+
+from flwr.common import Array, ArrayRecord, Message, MetricRecord, RecordDict
+
+from .fedopt import FedOpt
+from .strategy_utils import AggregationError
+
+
+# pylint: disable=line-too-long
+class FedAdam(FedOpt):
+    """FedAdam - Adaptive Federated Optimization using Adam.
+
+    Implementation based on https://arxiv.org/abs/2003.00295v5
+
+    Parameters
+    ----------
+    fraction_train : float (default: 1.0)
+        Fraction of nodes used during training. In case `min_train_nodes`
+        is larger than `fraction_train * total_connected_nodes`, `min_train_nodes`
+        will still be sampled.
+    fraction_evaluate : float (default: 1.0)
+        Fraction of nodes used during validation. In case `min_evaluate_nodes`
+        is larger than `fraction_evaluate * total_connected_nodes`,
+        `min_evaluate_nodes` will still be sampled.
+    min_train_nodes : int (default: 2)
+        Minimum number of nodes used during training.
+    min_evaluate_nodes : int (default: 2)
+        Minimum number of nodes used during validation.
+    min_available_nodes : int (default: 2)
+        Minimum number of total nodes in the system.
+    weighted_by_key : str (default: "num-examples")
+        The key within each MetricRecord whose value is used as the weight when
+        computing weighted averages for both ArrayRecords and MetricRecords.
+    arrayrecord_key : str (default: "arrays")
+        Key used to store the ArrayRecord when constructing Messages.
+    configrecord_key : str (default: "config")
+         Key used to store the ConfigRecord when constructing Messages.
+    train_metrics_aggr_fn : Optional[callable] (default: None)
+        Function with signature (list[RecordDict], str) -> MetricRecord,
+        used to aggregate MetricRecords from training round replies.
+        If `None`, defaults to `aggregate_metricrecords`, which performs a weighted
+        average using the provided weight factor key.
+    evaluate_metrics_aggr_fn : Optional[callable] (default: None)
+        Function with signature (list[RecordDict], str) -> MetricRecord,
+        used to aggregate MetricRecords from training round replies.
+        If `None`, defaults to `aggregate_metricrecords`, which performs a weighted
+        average using the provided weight factor key.
+    eta : float, optional
+        Server-side learning rate. Defaults to 1e-1.
+    eta_l : float, optional
+        Client-side learning rate. Defaults to 1e-1.
+    beta_1 : float, optional
+        Momentum parameter. Defaults to 0.9.
+    beta_2 : float, optional
+        Second moment parameter. Defaults to 0.99.
+    tau : float, optional
+        Controls the algorithm's degree of adaptability. Defaults to 1e-3.
+    """
+
+    # pylint: disable=too-many-arguments, too-many-locals
+    def __init__(
+        self,
+        *,
+        fraction_train: float = 1.0,
+        fraction_evaluate: float = 1.0,
+        min_train_nodes: int = 2,
+        min_evaluate_nodes: int = 2,
+        min_available_nodes: int = 2,
+        weighted_by_key: str = "num-examples",
+        arrayrecord_key: str = "arrays",
+        configrecord_key: str = "config",
+        train_metrics_aggr_fn: Optional[
+            Callable[[list[RecordDict], str], MetricRecord]
+        ] = None,
+        evaluate_metrics_aggr_fn: Optional[
+            Callable[[list[RecordDict], str], MetricRecord]
+        ] = None,
+        eta: float = 1e-1,
+        eta_l: float = 1e-1,
+        beta_1: float = 0.9,
+        beta_2: float = 0.99,
+        tau: float = 1e-3,
+    ) -> None:
+        super().__init__(
+            fraction_train=fraction_train,
+            fraction_evaluate=fraction_evaluate,
+            min_train_nodes=min_train_nodes,
+            min_evaluate_nodes=min_evaluate_nodes,
+            min_available_nodes=min_available_nodes,
+            weighted_by_key=weighted_by_key,
+            arrayrecord_key=arrayrecord_key,
+            configrecord_key=configrecord_key,
+            train_metrics_aggr_fn=train_metrics_aggr_fn,
+            evaluate_metrics_aggr_fn=evaluate_metrics_aggr_fn,
+            eta=eta,
+            eta_l=eta_l,
+            beta_1=beta_1,
+            beta_2=beta_2,
+            tau=tau,
+        )
+
+    def aggregate_train(
+        self,
+        server_round: int,
+        replies: Iterable[Message],
+    ) -> tuple[Optional[ArrayRecord], Optional[MetricRecord]]:
+        """Aggregate ArrayRecords and MetricRecords in the received Messages."""
+        aggregated_arrayrecord, aggregated_metrics = super().aggregate_train(
+            server_round, replies
+        )
+
+        if aggregated_arrayrecord is None:
+            return aggregated_arrayrecord, aggregated_metrics
+
+        if self.current_arrays is None:
+            reason = (
+                "Current arrays not set. Ensure that `configure_train` has been "
+                "called before aggregation."
+            )
+            raise AggregationError(reason=reason)
+
+        # Compute intermediate variables
+        delta_t, m_t, aggregated_ndarrays = self._compute_deltat_and_mt(
+            aggregated_arrayrecord
+        )
+
+        # v_t
+        if not self.v_t:
+            self.v_t = {k: np.zeros_like(v) for k, v in aggregated_ndarrays.items()}
+        self.v_t = {
+            k: self.beta_2 * v + (1 - self.beta_2) * (delta_t[k] ** 2)
+            for k, v in self.v_t.items()
+        }
+
+        # Compute the bias-corrected learning rate, `eta_norm` for improving convergence
+        # in the early rounds of FL training. This `eta_norm` is `\alpha_t` in Kingma &
+        # Ba, 2014 (http://arxiv.org/abs/1412.6980) "Adam: A Method for Stochastic
+        # Optimization" in the formula line right before Section 2.1.
+        eta_norm = (
+            self.eta
+            * np.sqrt(1 - np.power(self.beta_2, server_round + 1.0))
+            / (1 - np.power(self.beta_1, server_round + 1.0))
+        )
+
+        new_arrays = {
+            k: x + eta_norm * m_t[k] / (np.sqrt(self.v_t[k]) + self.tau)
+            for k, x in self.current_arrays.items()
+        }
+
+        # Update current arrays
+        self.current_arrays = new_arrays
+
+        return (
+            ArrayRecord(OrderedDict({k: Array(v) for k, v in new_arrays.items()})),
+            aggregated_metrics,
+        )