Assisted Learning for Organizations with Limited Imbalanced Data

1 INTRODUCTION . Modern distributed learning frameworks such as federated learning ( Shokri & Shmatikov , 2015 ; Konecny et al. , 2016 ; McMahan et al. , 2017 ) aim to improve the learning performance for a large number of learners that have limited data and computation/communication resources . These learning frameworks are well suited for cloud systems and IoT systems ( Ray , 2016 ; Gomathi et al. , 2018 ) that manage numerous smart devices through wireless communication . Over the past decade , many organizations , e.g. , government agencies , hospitals , schools , and companies , have integrated machine learning models into their work pipeline to facilitate data analysis and decision making . For example , according to a recent survey ( Financesonline , 2021 ) , 49 % of companies worldwide are exploring or planning to use machine learning , 51 % of organizations claim to be early adopters of machine learning , and the estimated productivity improvement from the learning models is 40 % . However , the performance of their machine learning models critically depends on the quality of the data , which typically is of a limited population and is biased toward certain distributions . Unfortunately , the existing learning frameworks can not help big organizations improve their learning performance due to the following major restrictions . • Unlike smart devices in the conventional federated learning , organizational learners typically cooperate with a single external service provider under a rigorous contract . Moreover , the service provider is often presumed to have more data with better quality than the organization . • Conventional distributed learners achieve the performance goal by frequently exchanging information with other learners . In comparison , each learning round for organizational learners is costly , as they need to pay the provider for assistance and need to exchange a large amount of information with the provider . Hence , organizational learners desire to achieve a significant performance improvement within limited assistance rounds . Therefore , there is an emerging need to develop a modern machine learning framework for organizational learners that can significantly improve the model performance by purchasing limited assistance services from external providers without data sharing . This constitutes the goal of this work . In this work , we develop an assisted learning framework in the ‘ horizontal-splitting ’ setting , where the learner and the service provider possess different datasets that are utilized for training a common model . In our context , the learner ’ s data is assumed to be limited and imbalanced , while the provider ’ s data is supposed to be big and complements the learner ’ s data . Our learning framework nicely suits the organizational learners ’ characteristics : they have a very limited budget for purchasing external assistance services , yet they exchange a large amount of side information with the provider per assistance round to maximize the performance gain . This is opposite to federated learning , where smart devices are equipped with only a limited communication budget but can endlessly learn through interacting with the cloud . We summarize our contributions as follows . 1.1 OUR CONTRIBUTIONS . We identify the need for developing an assisted learning framework for facilitating the deployment of general machine learning in large organizations . This learning framework addresses the unique challenges as explained previously . We first develop an assisted deep learning framework for organizational learners with limited and imbalanced data , and propose a stochastic training algorithm named AssistSGD . Specifically , every assistance round of AssistSGD consists of two phases . In the first phase , the learner performs local SGD training for multiple iterations and sends the generated trajectory of models together with their corresponding local loss values to the service provider . In the second phase , the provider utilizes the learner ’ s information to evaluate the global loss of the received models , and uses the best model with the smallest global loss as an initialization . Then , the provider performs local SGD training for multiple iterations and sends the generated trajectory of models together with their corresponding local loss values to the learner . Finally , the learner utilizes the provider ’ s information to evaluate the global loss of the received models , and outputs the best model with the smallest global loss . Under mild technical assumptions , we formally prove that AssistSGD with full batch gradient updates is guaranteed to find a critical point of the global loss function in general nonconvex optimization . We further generalize the framework to enable assisted reinforcement learning , and develop a policy gradient training algorithm named AssistPG , which has the same training logic as that of AssistSGD . Through extensive experiments with deep learning and reinforcement learning , we demonstrate that the learner can achieve a near-oracle performance with AssistSGD and AssistPG as if all the data were centralized . In particular , as the learner data ’ s level of imbalance increases , AssistSGD can help the learner achieve a higher performance gain . Moreover , data are never exchanged in the assisted learning process for both participants . 1.2 RELATED WORK . Assisted learning . Earlier work on assisted learning ( Xian et al. , 2020 ) considers organizations that collect different features from the same cohort . This is in contrast with our context where organizations hold the same features but imbalanced data distributions or environments . Also , our method applies to general deep learning and reinforcement learning tasks , which are beyond the previously studied regression task . Consequently , our algorithm designs and application scenarios are significantly different from the prior work . Distributed optimization . In conventional distributed optimization , the data is evenly distributed among workers , which collaboratively solve a large-scale problem by exchanging local information ( gradients , models . ) via either decentralized networks ( Xie et al. , 2016 ; Lian et al. , 2017 ; 2018 ) or centralized networks ( Ho et al. , 2013 ; Li et al. , 2014 ; Richtarik & Takavc , 2016 ; Zhou et al. , 2016 ; 2018 ) . As a comparison , our AssistSGD only requires a few transmission rounds between the learner and provider . This is particularly appealing for organizational learners , who can employ a sophisticated optimization process locally while restricting the rounds of assistance . Federated learning . Federated learning is an emerging distributed learning framework ( Shokri & Shmatikov , 2015 ; Konecny et al. , 2016 ; McMahan et al. , 2017 ; Zhao et al. , 2018 ; Li et al. , 2020 ; Diao et al. , 2021 ) , which aims to learn a global model using the average of local models trained by numerous smart devices with heterogeneous data . The existing federated learning algorithms require frequent transmissions of local model parameters . This is different from our solution designed for the organizational learning scenarios , where each learner is an organization that often has unconstrained communication and computation resources , but is restricted to interact with other external service providers . Our solution aims to help the learner improve learning performance within ten rounds , while federated learning needs many more rounds . 2 ASSISTED DEEP LEARNING . In this section , we introduce the assisted deep learning framework . Throughout the paper , L denotes a learner who seeks assistance , and P denotes a service provider who provides assistance to L . 2.1 PROBLEM FORMULATION . We consider the case where the learner L aims to train a machine learning model θ that performs well on its own datasetD ( L ) and generalizes well to unseen data . In general , L can train a machine learning model by solving the empirical risk minimization problem minθ∈Θ f ( θ ; D ( L ) ) , where f ( · ; D ( L ) ) is the loss on D ( L ) and Θ is the parameter space . Standard statistical learning theories show that the obtained model can generalize well to intact test samples under suitable constraints of model parsimoniousness ( Ding et al. , 2018 ) . However , when the user ’ s data D ( L ) contains a limited number of samples that are highly imbalanced , the learned model will suffer from overfitting or deteriorated generalization capability to the unseen test data . To overcome data deficiency , the learner L intends to connect with an external service provider P ( e.g. , a commercialized data company ) , who possesses data D ( P ) that are sufficient or complementary to the learner ’ s data D ( L ) . Ideally , the user L would improve the model by solving the following data-augmented problem , where D ( L , P ) : = D ( L ) ∪ D ( P ) denotes the centralized data . θ ( L , P ) = arg min θ∈Θ f ( θ ; D ( L , P ) ) , where D ( L , P ) = D ( L ) ∪ D ( P ) . ( 1 ) We note that f ( θ ; D ( L , P ) ) = f ( θ , D ( L ) ) + f ( θ , D ( P ) ) . If D ( P ) is generated from a distribution that is close to the underlying data distribution , then it is expected that θ ( L , P ) will achieve significantly improved performance on unseen data . However , it is unrealistic to centralize the data since the interactions between the learner L and the provider P are often restricted by various regulations . Some representative regulations that formally define the assisted learning framework are listed below . Assisted Learning Protocols 1 . No data sharing : Neither the learner L nor the provider P will share data with each other . 2 . Limited assistance : The learner L has a limited budget for purchasing assistance service . The learner desires to maximize the performance gain with only a few assistance rounds . 3 . Unlimited communication bandwidth : In each assistance round , the learner and the provider can exchange unlimited information . For example , the learner ( resp . provider ) can send an employee ( resp . technician ) to deliver a large-capacity hard drive to the other . The above assisted learning framework is different from the existing learning frameworks . For example , in federated learning , many devices collaboratively train a global model via a large number of learning rounds with limited communication bandwidth . In comparison , the organizational learner in assisted learning can only query a few rounds of assistance from the provider , but can exchange unlimited information with it . Hence , we need to develop a training algorithm that can substantially improve the learner ’ s model quality using limited interactions with the provider for assisted learning . Next , we present an assisted stochastic gradient descent ( AssistSGD ) algorithm for this purpose . 2.2 ASSISTSGD FOR ASSISTED DEEP LEARNING . We propose AssistSGD in Algorithm 1 for assisted deep learning . The learning process consists of R rounds , each consisting of the following interactions between the learner L and the provider P. ( 1 ) First , the learner L initiates a local learning process . It initializes a model θ ( L ) 0 and applies SGD with learning rate η to update it for T iterations using the local dataset D ( L ) . Then , the learner evaluates the local loss f ( · ; D ( L ) ) in a subset T of the iterations t = 0 , 1 , ... T − 1 . Lastly , the learner sends this subset of models and their corresponding local loss to the provider P . Algorithm 1 AssistSGD Input : Initialization model θ0 , learning rate η , assistance rounds R , local iterations T . for assistance rounds r = 1 , . . . , R do Learner L : I Initialize θ ( L ) 0 = θ r−1 . I Local SGD training to generate { θ ( L ) t } T−1t=0 . I Send { θ ( L ) t , f ( θ ( L ) t ; D ( L ) ) } t∈T to provider P . ———————————————————— Provider P : I Initialize θ ( P ) 0 = arg minθ∈ { θ ( L ) t } t∈T f ( θ ; D ( L , P ) ) . I Local SGD training to generate { θ ( P ) t } T ′−1 t=0 . I Send { θ ( P ) t , f ( θ ( P ) t ; D ( P ) ) } t∈T ′ to learner L . ———————————————————— Learner L : I Output θr = arg min θ∈ { θ ( P ) t } t∈T ′ f ( θ ; D ( L , P ) ) . end Output : The best model in { θr } Rr=1 . ( 2 ) Upon receiving the information from the learner L , the provider P first evaluates the global loss f ( · ; D ( L , P ) ) of the received set of models { θ ( L ) t , t ∈ T } and picks the best one ( denoted by θ ( P ) 0 ) for initialization . Note that the global loss can be evaluated because the local loss { f ( θ ( L ) t ; D ( L ) ) , t ∈ T } are provided by the learner L , and the provider P just needs to evaluate the local loss { f ( θ ( L ) t ; D ( P ) ) , t ∈ T } . After that , the provider applies SGD with learning rate η to update the model for T ′ iterations on the local dataset D ( P ) . Then , the provider evaluates the local loss f ( · ; D ( P ) ) in a subset T ′ of the iterations t = 0 , 1 , ... T ′ − 1 , and sends the subset of models and their corresponding local loss to the learner L. ( 3 ) Once the learner L receives feedback from the provider P , it evaluates the global loss f ( · ; D ( L , P ) ) of the received set of models { θ ( P ) t , t ∈ T ′ } and picks the best model as the output model of this assistance round . Discussions . The above algorithm works for general deep learning tasks . It does not require data sharing between the learner and the provider . Moreover , in each learning round , the learner and the provider exchange a small number of their local training models . As we show later in the experimental studies , it suffices to sample the iterations in T , T ′ at a low frequency . Such an assisted learning process is very different from , for example , the federated learning process . Particularly , in each round of federated learning , all learners perform a small number of local SGD updates and send their last output models to the cloud due to limited communication bandwidths . Consequently , the global model needs a large number of learning rounds to achieve a desirable performance .

Summary: This paper studies a novel problem setup where a learner has unbalanced data, and a service provider has complementary or sufficient data, and the learner needs to improve accuracy in as few rounds of communication as possible, where the communication in each round is unbounded. An algorithm AssistSGD is proposed and shown to converge to a stationary point. Experiments show that AssistSGD has better performance than baselines and is close to centralized SGD on CIFAR 10 and reinforcement learning tasks,