User equipment grouping for federated learning
Various aspects of the present disclosure generally relate to wireless communication. In some aspects, a user equipment (UE) may transmit, to a network node, an indication of a local training data distribution associated with the UE. The UE may transmit, to the network node, local gradient information for a federated learning model that is based at least in part on the local training data distribution associated with the UE. Numerous other aspects are provided.
1 . A user equipment (UE) for wireless communication, comprising:
at least one memory; and
at least one processor communicatively coupled with the at least one memory, the at least one processor configured to cause the UE to:
transmit, to a network node, an indication of a local training data distribution associated with the UE;
receive, from the network node in response to the indication of the local training data distribution associated with the UE, configuration information comprising a federated learning model index associated with a federated learning model of a plurality of federated learning models, wherein the federated learning model index corresponds to a group of UEs, each UE of the group of UEs having a respective local training data distribution within a range of local training data distributions; and
transmit, to the network node, local gradient information in accordance with the federated learning model that is based at least in part on the local training data distribution associated with the UE.
2 . The UE of claim 1 , wherein the indication of the local training data distribution includes an indication of at least one of an input training data distribution of a local dataset associated with the UE or an output training data distribution of the local dataset associated with the UE.
3 . The UE of claim 1 , wherein the at least one processor is further configured to cause the UE to receive, from the network node, an indication of a reporting resource associated with the federated learning model, wherein, to cause the UE to transmit the local gradient information, the at least one processor is configured to cause the UE to:
transmit the local gradient information in the reporting resource associated with the federated learning model.
4 . The UE of claim 1 , wherein the at least one processor is further configured to cause the UE to receive, from the network node, configuration information that indicates respective reporting resources associated with multiple federated learning models, wherein the federated learning model that is based at least in part on the local training data distribution associated with the UE is a first federated learning model of the multiple federated learning models, and wherein, to cause the UE to transmit the local gradient information, the at least one processor is configured to cause the UE to:
transmit the local gradient information in the respective reporting resource associated with the first federated learning model.
5 . The UE of claim 4 , wherein the at least one processor is further configured to cause the UE to:
transmit, to the network node, an indication of an updated local training data distribution associated with the UE; and
transmit other local gradient information in the respective reporting resource associated with a second federated learning model of the multiple federated learning models based at least in part on the updated local training data distribution associated with the UE.
6 . The UE of claim 1 , wherein the configuration information indicates a scheduling condition.
7 . The UE of claim 1 , wherein the at least one processor is further configured to cause the UE to receive configuration information that configures reporting of the local training data distribution associated with the UE, wherein, to cause the UE to transmit the indication of the local training data distribution, the at least one processor is configured to cause the UE to:
transmit the indication of the local training data distribution based at least in part on the configuration information.
8 . The UE of claim 1 , wherein the indication of the local training data distribution indicates one or more statistical properties associated with the local training data distribution.
9 . The UE of claim 1 , wherein the indication of the local training data distribution indicate a P-value associated with the local training data distribution.
10 . The UE of claim 1 , wherein the indication of the local training data distribution indicates a respective input training data distribution for each of multiple inputs in a local dataset associated with the UE.
11 . The UE of claim 1 , wherein the indication of the local training data distribution indicates a plurality of Gaussian components of a Gaussian mixture distribution, and wherein, for each Gaussian component, of the plurality of Gaussian components, the indication includes:
a respective mean,
a respective covariance matrix, and
a respective mixing probability.
12 . The UE of claim 11 , wherein the at least one processor is further configured to cause the UE to:
receive, from the network node, configuration information that indicates the plurality of Gaussian components.
13 . The UE of claim 11 , wherein the at least one processor is further configured to cause the UE to:
estimate the respective mean, the respective covariance matrix, and the respective mixing probability for each Gaussian component of the plurality of Gaussian components based at least in part on the local training data distribution.
14 . The UE of claim 1 , wherein the indication of the local training data distribution indicates the local training data distribution as a mixture distribution including one or more components associated with a base distribution.
15 . The UE of claim 14 , wherein the indication of the local training data distribution indicates, for each component of the one or more components of the mixture distribution:
one or more parameters associated with the base distribution, and
a mixing probability.
16 . The UE of claim 14 , wherein the base distribution includes at least one of a uniform distribution, an exponential distribution, a Gaussian distribution, or an inverse Gaussian distribution.
17 . The UE of claim 14 , wherein the at least one processor is further configured to cause the UE to:
receive, from the network node, configuration information that indicates at least one of the base distribution or a maximum quantity of the one or more components of the mixture distribution.
18 . The UE of claim 17 , wherein the at least one processor is further configured to cause the UE to:
transmit, to the network node, capability information that indicates a capability of the UE for reporting the local training data distribution as the mixture distribution, wherein the capability information indicates at least one of a capability for the base distribution or a capability for the maximum quantity of the one or more components of the mixture distribution.
19 . A network node for wireless communication, comprising:
at least one memory; and
at least one processor communicatively coupled with the at least one memory, the at least one processor configured to cause the network node to:
receive an indication of a local training data distribution associated with a user equipment (UE);
in response to the indication of the local training data distribution associated with the UE, assign the UE to a group of UEs, the group of UEs associated with a federated learning model of a plurality of federated learning models, wherein each UE of the group of UEs has a respective local training data distribution within a range of local training data distributions;
transmit configuration information comprising a federated learning model index associated with the federated learning model of the plurality of federated learning models, wherein the federated learning model index corresponds to the group of UEs; and
receive local gradient information associated with the UE in accordance with the federated learning model that is associated with the group of UEs.
20 . The network node of claim 19 , wherein the at least one processor is further configured to cause the network node to transmit, to the UE, an indication of a reporting resource associated with the federated learning model, wherein, to cause the network node to receive the local gradient information, the at least on processor is configured to cause the network node to:
receive the local gradient information in the reporting resource associated with the federated learning model.
21 . The network node of claim 19 , wherein the at least one processor is further configured to cause the network node to transmit, to the UE, configuration information that indicates respective reporting resources associated with multiple federated learning models, wherein the federated learning model that is based at least in part on the local training data distribution associated with the UE is a first federated learning model of the multiple federated learning models, and wherein, to cause the network node to receive the local gradient information, the at least on processor is configured to cause the network node to:
receive the local gradient information in the respective reporting resource associated with the first federated learning model.
22 . The network node of claim 19 , wherein the at least one processor is further configured to cause the network node to transmit configuration information that configures reporting of the local training data distribution associated with the UE, wherein, to cause the network node to receive the indication of the local training data distribution, the at least on processor is configured to cause the network node to:
receive the indication of the local training data distribution based at least in part on the configuration information.
23 . The network node of claim 19 , wherein the indication of the local training data distribution indicates one or more statistical properties associated with the local training data distribution.
24 . The network node of claim 19 , wherein the indication of the local training data distribution indicates a plurality of Gaussian components of a Gaussian mixture distribution, and wherein, for each Gaussian component, of the plurality of Gaussian components, the indication includes:
a respective mean,
a respective covariance matrix, and
a respective mixing probability.
25 . The network node of claim 19 , wherein the indication of the local training data distribution indicates the local training data distribution as a mixture distribution including one or more components associated with a base distribution.
26 . The network node of claim 19 , wherein the configuration information indicates a scheduling condition.
27 . A method of wireless communication performed by a user equipment (UE), comprising:
transmitting, to a network node, an indication of a local training data distribution associated with the UE;
receiving, from the network node in response to the indication of the local training data distribution associated with the UE, configuration information comprising a federated learning model index associated with a federated learning model of a plurality of federated learning models, wherein the federated learning model index corresponds to a group of UEs, each UE of the group of UEs having a respective local training data distribution within a range of local training data distributions; and
transmitting, to the network node, local gradient information in accordance with the federated learning model that is based at least in part on the local training data distribution associated with the UE.
28 . The method of claim 27 , wherein the configuration information indicates a scheduling condition.
29 . A method of wireless communication performed by a network node, comprising:
receiving an indication of a local training data distribution associated with a user equipment (UE);
in response to the indication of the local training data distribution associated with the UE, assigning the UE to a group of UEs, the group of UEs associated with a federated learning model of a plurality of federated learning models, wherein each UE of the group of UEs has a respective local training data distribution within a range of local training data distributions;
transmitting configuration information comprising a federated learning model index associated with the federated learning model of the plurality of federated learning models, wherein the federated learning model index corresponds to the group of UEs; and
receiving local gradient information associated with the UE in accordance with the federated learning model that is associated with the group of UEs.
30 . The method of claim 29 , wherein the configuration information indicates a scheduling condition.