COMPUTATION-EFFICIENT FEDERATED LEARNING FOR SYSTEMS WITH RESOURCE HETEROGENEITY

A computer-implemented method for training a global model on a central server in a federated learning system comprised of a plurality of nodes includes splitting the global model along a width and a depth via two-dimensional uniform downscaling of the global model. A plurality of local models are created based on the splitting of the global model. Selected ones of the plurality of local models are trained on respective selected ones of a plurality of clients based on computational constraints of each of the plurality of clients.

BACKGROUND

Technical Field

The present disclosure generally relates to federated learning systems, and more particularly, to a computer-implemented method, a computer system, and a computer program product for computation-efficient federated learning that can handle system heterogeneity by using early exits, two-dimensional model downscaling and optimization with self-distillation.

Description of the Related Art

Federated learning (FL) enables training a deep learning model across multiple clients, such as internet of things (IoT) devices, security cameras, laptops, smartphones, data centers, and the like, with decentralized data. With typical federated learning architectures, at each round, a global model is distributed to clients for local training and the central server aggregates the local updates received from each client. In most FL architectures, all clients are assumed to have similar computational capabilities and be able to finetune/train the global model.

SUMMARY

In one embodiment, a system and method are provided that can provide computationally overhead-adjusted models to clients and can perform self-distillation on the client side, protecting data and reducing network traffic involved for the training of FL models across various clients.

In one embodiment, a computer-implemented method for training a global model on a central server in a federated learning system comprised of a plurality of nodes includes splitting the global model along a width and a depth via two-dimensional uniform downscaling of the global model. A plurality of local models can be created based on the splitting of the global model. Selected ones of the plurality of local models can be trained on respective selected ones of a plurality of clients based on computational constraints of each of the plurality of clients.

In some embodiments, the method further includes receiving, at the central server, local model parameters from each of the plurality of clients. In some embodiments, the method further includes aggregating, at the central server, the local model parameters across the plurality of clients into global model parameters and sending the global model parameters to each of the plurality of clients to update respective local models at each of the plurality of clients.

In some embodiments, the method further includes waiting, by the central server, until the aggregated local model parameters are received from all of the plurality of clients before updating the global model parameters.

In some embodiments, the method further includes obtaining a number of complexity levels for the federated learning system and a target computational overhead reduction ratio for each of the complexity levels. A computational overhead of each of the plurality of local models at each of the complexity levels can then be computed.

In some embodiments, the method further includes determining early exits of the global model to generate a local model for each of the complexity levels.

In another embodiment, a computer-implemented method for training a global model on a central server in a federated learning system having a plurality of nodes includes obtaining a number of complexity levels for the federated learning system and a target computational overhead reduction ratio for each of the complexity levels. The global model can be split along a width and a depth via two-dimensional uniform downscaling of the global model to create a plurality of local models, wherein one of the plurality of local models corresponds to each of the number of complexity levels. A computational overhead of each of a plurality of local models at each of the complexity levels can be computed and an assigned one of the plurality of local models can be sent to each of a plurality of clients based on an available computational overhead budget at each of the plurality of clients, wherein the computational overhead of the assigned one is less than the available computational overhead budget at each of the plurality of clients. The assigned ones of the plurality of local models can be trained on respective ones of the plurality of clients.

The above method can be performed on non-transitory computer readable storage medium tangibly embodying a computer readable program code having computer readable instructions that, when executed, causes a computer device to provide computation-efficient federated learning that can handle system heterogeneity by using early exits, two-dimensional model downscaling and optimization with self-distillation.

By virtue of the concepts discussed herein, systems and methods are provided for providing computation-efficient federated learning that can handle system heterogeneity by using early exits, two-dimensional model downscaling and optimization with self-distillation. As discussed in greater detail below, such a system and method can reduce computational overhead/complexity by providing scaled models for different clients depending on the computational resources of the various clients.

DETAILED DESCRIPTION

Broadly, aspects of the present disclosure provide systems and methods that provide an improved FL framework. As described in greater detail below, the systems and methods can provide efficient small subnetworks for constrained clients via two-dimensional uniform downscaling through model splitting along a width (hidden size) and a depth (number of layers) using early exits. Resulting local models provide the best balance between preserving low-level basic and high-level complex feature extraction capabilities. Local models at lower complexity levels preserve high performance for constrained clients during inference. Early exits bring the adaptive inference capability if inference-time constraints are dynamic. The systems and methods provide local optimization with self-distillation over early exit predictions (early exits as students and final exit as teacher) to improve the knowledge transfer among subnetworks within the global model. With the systems and methods described herein, neither additional training on clients or the central server over shared data nor sharing of the intermediate layer outputs are required.

Although the operational/functional descriptions described herein may be understandable by the human mind, they are not abstract ideas of the operations/functions divorced from computational implementation of those operations/functions. Rather, the operations/functions represent a specification for an appropriately configured computing device. As discussed in detail below, the operational/functional language is to be read in its proper technological context, i.e., as concrete specifications for physical implementations.

Accordingly, one or more of the methodologies discussed herein may provide federated learning systems and methods that can provide efficient small subnetworks for constrained clients via two-dimensional uniform downscaling through model splitting along a width (hidden size) and a depth (number of layers) using early exits. This may have the technical effect of permitting clients with limited computing resources to use a consolidated model, or an early exit therefrom, for their local data, thus providing a model that can be executed on the client while not requiring additional computing resources beyond those dedicated to the model by the client.

It should be appreciated that aspects of the teachings herein are beyond the capability of a human mind. It should also be appreciated that the various embodiments of the subject disclosure described herein can include information that is impossible to obtain manually by an entity, such as a human user. For example, the type, amount, and/or variety of information included in performing the process discussed herein can be more complex than information that could be reasonably be processed manually by a human user.

Framework

A training dataset can have N sets of input and target pairs={(Xi, yi)}i=1Ndistributed to K clients with {k}k=1K(the set of data indices at each client). A goal is training a model in a scenario where clients can have different computational constraints. The configuration of constraints is given, which contains the complexity level 1k∈{1, . . . . L} of each client k and the target computational overhead reduction ratio for each level n. As used herein, the computational overhead and/or the target computational overhead can be based on, for example, model number of parameters (model size), random access memory (RAM) usage, number of floating-point operations (#FLOPs), latency, power consumption, or the like.

Algorithm 1, below, explains the procedure followed in the present disclosure. The system architecture is provided inFIG.1for three levels. Given the constraint configuration for each level j, aspects of the present disclosure can determine the horizontal and vertical split ratios (shl, svl) using Equation (2), illustrated below, which determines what computed downscaling ratios of the model will be kept for each complexity level along the model's depth and width dimensions respectively. Since the initial model architecture M has only one output in the final layer, early exit classifiers can be injected to the layers based on the computed horizontal split ratio values. This multi-exit model architecture can be denoted as ML, which is considered as the global model in the framework. InFIG.1, local models can be trained using a combination of cross-entropy (CE) and KL-divergence losses as given in Equation (5), below. Updates can be aggregated back using Equation (1) in the central server for the next round.

Algorithm 1: FrameworkInputs: Dataset= {(Xi, yi)}i=1Ndistributed to K clients with indexes {k}k=1K, numberof complexity levels L, complexity level of each client {lk}k=1K, target overhead reduc-tion ratios for each level {rl}l=1L, client availability rate s, model architecture M.Parameters: number of learning rounds T, number of local training epochs E, batch sizeB, learning rate η.Outputs: Trained global model MLwith weights θ.1:   ML← M2:   for level 1−1, ... L-1 do3:     Compute split ratio pair (sh(l), sv(l)) using Equation (2)4:     Add early exit classifier to MLat [sv(l)N]-th layer5:   end for6:   Initialize global model MLwith θ07:   for round t = 0, ... T-1 do8:     St← random subset of max (1, sK) clients9:     for client k ∈ Stin parallel do10:       Split Mlk← split (ML; sh(lk), sv(lk))11:       for epoch e = 1, ... E in client k do12:         for batch b ⊂kdo13:L=1B⁢∑i∈bℒ⁡(Mlk(Xi;θtk),yi)⁢with⁢Equation⁢(5)14:θtk←θtk-η⁡(∂L∂θtk)15:         end for16:       end for17:     end for18:     Aggregate, and obtain θt+1using Equation (1)19:   end for20:   return MLwith θT

At each communication round t, a set of available clients Stare sampled and the global model MLwith weights θtis scaled down to local models with architecture Mlkand weights θtkfor each client k∈Stbased on their complexity level 1k. This procedure is detailed in Algorithm 2, below. In this operation, the layers after the corresponding exit at the └sh(lk)N┘-th layer are removed. Here, index function takes a tensor size “size (W)” and split ratio value sv(lk)as inputs and returns the Boolean index tensor Z, which is used to split weights. In general, for hidden layers, this operation results in accessing the first └sv(lk)i┘ elements of the tensor W (in a hidden layer) along every dimension with sizei.

After receiving the local model, each client performs training for E epochs by minimizing the loss L with self-distillation defined in Equation (5) on its local dataset and sends back the updated weights. In the conventional federated learning model, FedAVG, all models are assumed to have the same architecture, hence aggregation is directly done by averaging. In the present disclosure, the aggregation procedure can be described as follows for every W∈θt:

for 1∈{1, . . . . L}, where Stl={k|k∈St, 1k≥1}. Here, Z1=index(size(W), svl) for 1∈{1′, . . . . L} and 0 for 1<1′, where 1′ is the minimum level that W exists (e.g., 1′=L if W is after (L−1)-th exit). Lastly,Wkis the zero padded local weight Wksuch thatWk[Z1]=WkandWk[1−Z1]=0. In other words, overlapping weights are aggregated after being scaled by the number of contributing clients.

Split Configuration

In this section, it is explained how the split ratios used to downscale the global into local models is computed at each complexity level. Let C=overhead(M) denote the computational overhead of the initial model M. For a given target computational overhead reduction ratio r1for the complexity level 1∈{1, . . . . L−1}, it can be determined at what ratio of the model to keep at each level horizontally and vertically (sh(l), sv(l)) as follows:

In other words, the most uniform split ratio pair can be found that satisfies the target computational overhead through a grid search within the window defined by ϵ. For the highest complexity level, sh{L}=sv(L)=1 is considered, i.e., local models at level L are the same as the global model during training. The overhead function can be defined by the user depending on the computational constraints of the application scenario. Two cases are considered in the present disclosure, (1) a spatial constraint: overhead(M)=#PARAM S(M) is the number of parameters in the model; and (2) a temporal constraint: overhead(M)=#FLOPs(M) is the number of floating point operations (FLOPs) in one forward pass of the model.

Optimization with Self-Distillation

The subnetworks of the global model MLcan be denoted as illustrated inFIG.2, where subnetwork structures for local models at levels 1, 2 and L as shown. For the local model Mjat level j, fiis the i-th core subnetwork with weights ωi,j(f). Likewise, gi is the i-th exit classifier subnetwork with weights ωi,j(g)is the output at the i-th exit of Mj. The forward pass of local models can be formulated as follows:

for 1≤i≤j≤L where j is the level of the local model and H0,j=X. After obtaining the prediction logits at each exit for M1at level 1, the loss with self-distillation is calculated as follows:

whereKL(ŷs·ŷl; τ)=sum (σ(ŷt/τ) log ((σ(ŷt/τ)/σ(ŷs/τ))τ2) is Kullback-Leibler divergence with temperature τ>0,CE(ŷ, y)=−log σ(ŷ) [y] is cross-entropy loss for target class y, σ is softmax function and β∈[0, 1) is the hyperparameter that controls the self-distillation effect.

Inference

Finally, the pseudo-code is provided for inferencing in Algorithm 3, below. Based on the complexity level of the client 1, global model MLis split and inference is performed on the local model M1. If an adaptive inference flag a is enabled, this procedure also continuously outputs early exit predictions.

Example Process

It may be helpful now to consider a high-level discussion of an example process. To that end,FIGS.3A and3Bpresent an illustrative process300related to the method for training deep neural networks in a federated learning system. Process300is illustrated as a collection of blocks, in a logical flowchart, which represents a sequence of operations that can be implemented in hardware, software, or a combination thereof. In the context of software, the blocks represent computer-executable instructions that, when executed by one or more processors, perform the recited operations. Generally, computer-executable instructions may include routines, programs, objects, components, data structures, and the like that perform functions or implement abstract data types. In each process, the order in which the operations are described is not intended to be construed as a limitation, and any number of the described blocks can be combined in any order and/or performed in parallel to implement the process.

Referring toFIGS.3A and3B, a process300for training deep neural networks in a federated learning system starts at block302where the system can obtain the number of complexity levels L and target overhead reduction ratios for each level. At block304, a model MLis initialized with, at block306, an initialized complexity level 1=1. If 1 is not equal to L, then, at block308, the process can determine the most uniform 2-dimensional downscaling ratios (shl, Sv) through a grid search while satisfying the target overhead reduction ratio using Equation (2). At block310, early exit classifiers are injected into the └Nshl┘-th layer of ML. At block312, MLis split based on the split ratios (shl, Sv) using Algorithm 2 to obtain a local model M1for level 1. The increment level 1 can be increased by one at block314and the process can again determine whether 1=L. If 1 is not equal to L, then blocks308through314are repeated.

Once 1=L, the process continues to block316, where shland svlare set to 1. At block318, the computational overheads of models M1, M2, . . . . MLare computed and stored. At block320, the federated learning process is initiated over K clients for T rounds. The training round t is set to 1 at block322. If t>T, then the output of the trained global model MLis output at block324.

If t is not greater than T, then, at block326, sK available clients are identified and k is initialized to 1. If k is not greater than sK, then, at block328, the complexity level 1 is assigned for client k such that the overhead of M1does not exceed the budget of the client. At block330, the global model M1is split based on the split ratios (shl, Svl) using Algorithm 2 for level 1 and the local model M1can be obtained. The local model M1can be sent to the client device and k is incremented by one at block332. The process can return to comparison block338to check to determine if k is greater than sK. If k is not greater than sK, then blocks328through332are repeated.

Once k is greater than sK, then, at block340, the process can wait until all updated weights are received back from the clients. At block342, then local model weights are aggregated using Algorithm 3 and the global model ML is updated. At block344, the training round t is incremented by one and the process continues back to decision block346.

At block334, local training can be performed with self-distillation to minimize Equation (5) for e epochs. At block336, updated model weights can be sent back to the central server. It should be noted that blocks334and336can be performed in parallel at all K clients.

Results

The datasets CIFAR-10 and CIFAR-100 (Canadian Institute For Advanced Research) were used, where each dataset has train size of 50,000 images, a test size of 10,000 images, a resolution of 32 and a number of classes of either 10 (CIFAR-10) or 100 (CIFAR-100). The system topology includes 100 clients with 10% availability at each round and four complexity levels with target overhead reduction ratios of 12.5%, 25%, 50%, 100%, where the client level distribution is uniform (each level contains 25% of the clients).

Data was obtained comparing top-1 accuracy to either inference time per sample or number of parameters for the process of the present disclosure and conventional federated learning baselines, including FedAVG, which is a level-1 subnetwork trained using federated averaging algorithm and Decoupled, where one model for each complexity level is trained in a decoupled way. Further, existing methods, such as HeteroFL, which employs vertical model splitting along width, and FedDF, which uses ensemble distillation on central server over an additional dataset after each round, are compared to the process of the present disclosure.

FIGS.4A and4Billustrates the inference time per sample and number of parameters, respectively, for the conventional process as compared to the process of the present disclosure, on CIFAR-10.FIGS.5A and5Billustrates the inference time per sample and number of parameters, respectively, for the conventional process as compared to the process of the present disclosure, on CIFAR-100. In the results, it can be seen how the process of the present disclosure (labeled as “Pres. Disc.”) provides the highest accuracy at each inference time per sample measured as well as the best accuracy at each measured number of parameters.

Example Computing Platform

Referring toFIG.6, computing environment600includes an example of an environment for the execution of at least some of the computer code involved in performing the inventive methods, including a federated learning system deep neural network training engine block700. In addition to block700, computing environment600includes, for example, computer601, wide area network (WAN)602, end user device (EUD)603, remote server604, public cloud605, and private cloud606. In this embodiment, computer601includes processor set610(including processing circuitry620and cache621), communication fabric611, volatile memory612, persistent storage613(including operating system622and block700, as identified above), peripheral device set614(including user interface (UI) device set623, storage624, and Internet of Things (IoT) sensor set625), and network module615. Remote server604includes remote database630. Public cloud605includes gateway640, cloud orchestration module641, host physical machine set642, virtual machine set643, and container set644.

PROCESSOR SET610includes one, or more, computer processors of any type now known or to be developed in the future. Processing circuitry620may be distributed over multiple packages, for example, multiple, coordinated integrated circuit chips. Processing circuitry620may implement multiple processor threads and/or multiple processor cores. Cache621is memory that is located in the processor chip package(s) and is typically used for data or code that should be available for rapid access by the threads or cores running on processor set610. Cache memories are typically organized into multiple levels depending upon relative proximity to the processing circuitry. Alternatively, some, or all, of the cache for the processor set may be located “off chip.” In some computing environments, processor set610may be designed for working with qubits and performing quantum computing.

VOLATILE MEMORY612is any type of volatile memory now known or to be developed in the future. Examples include dynamic type random access memory (RAM) or static type RAM. Typically, volatile memory612is characterized by random access, but this is not required unless affirmatively indicated. In computer601, the volatile memory612is located in a single package and is internal to computer601, but, alternatively or additionally, the volatile memory may be distributed over multiple packages and/or located externally with respect to computer601.

END USER DEVICE (EUD)603is any computer system that is used and controlled by an end user (for example, a customer of an enterprise that operates computer601), and may take any of the forms discussed above in connection with computer601. EUD603typically receives helpful and useful data from the operations of computer601. For example, in a hypothetical case where computer601is designed to provide a recommendation to an end user, this recommendation would typically be communicated from network module615of computer601through WAN602to EUD603. In this way, EUD603can display, or otherwise present, the recommendation to an end user. In some embodiments, EUD603may be a client device, such as thin client, heavy client, mainframe computer, desktop computer and so on.

REMOTE SERVER604is any computer system that serves at least some data and/or functionality to computer601. Remote server604may be controlled and used by the same entity that operates computer601. Remote server604represents the machine(s) that collect and store helpful and useful data for use by other computers, such as computer601. For example, in a hypothetical case where computer601is designed and programmed to provide a recommendation based on historical data, then this historical data may be provided to computer601from remote database630of remote server604.

PRIVATE CLOUD606is similar to public cloud605, except that the computing resources are only available for use by a single enterprise. While private cloud606is depicted as being in communication with WAN602, in other embodiments a private cloud may be disconnected from the internet entirely and only accessible through a local/private network. A hybrid cloud is a composition of multiple clouds of different types (for example, private, community or public cloud types), often respectively implemented by different vendors. Each of the multiple clouds remains a separate and discrete entity, but the larger hybrid cloud architecture is bound together by standardized or proprietary technology that enables orchestration, management, and/or data/application portability between the multiple constituent clouds. In this embodiment, public cloud605and private cloud606are both part of a larger hybrid cloud.

CONCLUSION

Aspects of the present disclosure are described herein with reference to a flowchart illustration and/or block diagram of a method, apparatus (systems), and computer program products according to embodiments of the present disclosure. It will be understood that each block of the flowchart illustrations and/or block diagrams, and combinations of blocks in the flowchart illustrations and/or block diagrams, can be implemented by computer readable program instructions.