Personalized federated learning for a better customer experience

Accounting for data heterogeneity across edge devices enables more useful model updates, both locally and globally.

Federated learning (FL) is a framework that allows edge devices (e.g., Alexa devices) to collaboratively train a global model while keeping customers’ data on-device. A standard FL system involves a cloud server and multiple clients (devices). Each device has its local data and a local copy of the machine learning (ML) model being served.

In each round of FL training, a cloud server sends the current global model to the clients; the clients train their local models using on-device data and send the models to the cloud; and the server aggregates the local models and updates the global model. FL also has a personalization branch, which aims to customize local models to improve their performance on local data.

AmazonScience_FederatedLearning.gif
Illustration of a federated-learning system. Each edge device has its own local model, and it periodically communicates with the cloud server to update a shared global model. A personalization branch aims to customize local models to improve their performance on local data.

In many real-world applications, the local datasets for different clients may have heterogeneous distributions. In a paper we presented at the 36th Conference on Neural Information Processing Systems (NeurIPS), we show that a training procedure that accounts for that heterogeneity improves the efficiency and accuracy of both the local and global models in federated learning.

Related content
The team’s latest research on privacy-preserving machine learning, federated learning, and bias mitigation.

In particular, we consider two new measures of heterogeneity within an FL system: (1) intra-client uncertainty, which profiles the differences across time for a local model, and (2) inter-client uncertainty, which characterizes the differences between different users’ local models in the same time frame.

A larger inter-client uncertainty indicates a more heterogeneous local data distribution among the edge devices, which makes personalization more important. A larger intra-client uncertainty means that the local model (particularly, its learnable parameters) has large variations across different rounds of FL training. Accordingly, our method adjusts the local training configuration and the FL aggregation rule based on the two uncertainty values.

To evaluate our approach, we compared it with seven earlier FL algorithms on seven different datasets, spanning image and audio data. We found that our approach consistently delivered the highest accuracy, on both global and local models.

Accounting for uncertainty

Our intuition was that that when training a local model, selecting a proper initial model and an appropriate number of training steps is critical to minimizing the training loss, thus achieving the desired personalization.

Self-FL intuition.png
Optimization trajectories of a local model. The blue and red curves are the training curves obtained with two different initial models.

Our method, which we call Self-FL, is rooted in a theoretical analysis using Bayesian hierarchical models, in which the intra-client and inter-client uncertainties define different layers of the hierarchy. From the Bayesian analysis, we derive equations relating these two uncertainty measures to three local configuration factors: (1) the local initial model, which is used as the starting point of local model training; (2) the learning rate, which determines how dramatically a single training example can affect network weights; and (3) the early-stop rule, which determines when the training procedure should stop to prevent overfitting.

Related content
Innovative training methods and model compression techniques combine with clever engineering to keep speech processing local.

In practice, we don’t have accurate measurements of the intra-client and inter-client uncertainty quantities. But in the paper, we provide several different methods for estimating them. In our experiments, we use an estimate based on the variance in the parameter optimizations determined over the course of training.

To the best of our knowledge, this is the first work to connect personalized FL with hierarchical modeling and to use uncertainty quantification to drive personalization.

Adaptive FL aggregation rule 

Existing FL algorithms typically update the global model using a weighted sum of local models, where the weight for each local model is proportional to the local dataset size. Our framework uses an adaptive aggregation rule to update the global model for better personalization. Particularly, we derive the aggregation rule from Bayesian hierarchical modeling, where the global model parameters are considered the ”root” of the statistical model.

The idea is, essentially, that the more a local model’s training data deviates from the global averages, the more responsive the model should be to training on that data. Conversely, the more uncertain the optimization of the local models’ parameters appear to be, the less weight they should be given when updating the global model.

Self-FL framework.png
The Self-FL framework for personalized federated learning.

Our method is designed to improve the accuracy of edge devices’ personalized models both directly, by better tailoring them to the types of data they’re likely to see, and indirectly, by making the global models distributed to all clients more accurate. Our empirical results indicate that relative to prior FL schemes, Self-FL improves performance for edge clients. As such, it promises to improve the experience of Amazon customers by making their devices more responsive to their particular needs.

Related content

US, MA, North Reading
Are you excited about developing generative AI and foundation models to revolutionize automation, robotics and computer vision? Are you looking for opportunities to build and deploy them on real problems at truly vast scale? At Amazon Fulfillment Technologies and Robotics we are on a mission to build high-performance autonomous systems that perceive and act to further improve our world-class customer experience - at Amazon scale. We are looking for scientists, engineers and program managers for a variety of roles. The Research team at Amazon Robotics is seeking a passionate, hands-on Sr. Applied Scientist to help create the world’s first foundation model for a many-robot system. The focus of this position is how to predict the future state of our warehouses that feature a thousand or more mobile robots in constant motion making deliveries around the building. It includes designing, training, and deploying large-scale models using data from hundreds of warehouses under different operating conditions. This work spans from research such as alternative state representations of the many-robot system for training, to experimenting using simulation tools, to running large-scale A/B tests on robots in our facilities. Key job responsibilities * Research vision - Where should we be focusing our efforts * Research delivery - Proving/dis-proving strategies in offline data or in simulation * Production studies - Insights from production data or ad-hoc experimentation * Production implementation - Building key parts of deployed algorithms or models About the team You would join our multi-disciplinary science team that includes scientists with backgrounds in planning and scheduling, grasping and manipulation, machine learning, and operations research. We develop novel planning algorithms and machine learning methods and apply them to real-word robotic warehouses, including: - Planning and coordinating the paths of thousands of robots - Dynamic allocation and scheduling of tasks to thousands of robots - Learning how to adapt system behavior to varying operating conditions - Co-design of robotic logistics processes and the algorithms to optimize them Our team also serves as a hub to foster innovation and support scientists across Amazon Robotics. We also coordinate research engagements with academia, such as the Robotics section of the Amazon Research Awards. We are open to hiring candidates to work out of one of the following locations: North Reading, MA, USA | Westborough, MA, USA
US, WA, Bellevue
Are you excited about developing state-of-the-art deep learning foundation models, applied to the automation of labor for the future of Amazon’s Fulfillment network? Are you looking for opportunities to build and deploy them on real problems at truly vast scale? At Amazon Fulfillment Technologies and Robotics we are on a mission to build high-performance autonomous systems that perceive and act to further improve our world-class customer experience - at Amazon scale. To this end, we are looking for an Applied Scientist who will build and deploy models that help automate labor utilizing a wide array of multi-modal signals. Together, we will be pushing beyond the state of the art in optimization of one of the most complex systems in the world: Amazon's Fulfillment Network. Key job responsibilities In this role, you will build models that can identify potential problems with Amazon’s vast inventory, including discrepancies between the physical and virtual manifest and efficient execution of inventory audit operations. You will work with a diverse set of real world structured, unstructured and potentially multimodal datasets to train deep learning models that identify current inventory management problems and anticipate future ones. Datasets include multiple separate inventory management event streams, item images and natural language. You will face a high level of research ambiguity and problems that require creative, ambitious, and inventive solutions. About the team Amazon Fulfillment Technologies (AFT) powers Amazon’s global fulfillment network. We invent and deliver software, hardware, and data science solutions that orchestrate processes, robots, machines, and people. We harmonize the physical and virtual world so Amazon customers can get what they want, when they want it. The AFT AI team has deep expertise developing cutting edge AI solutions at scale and successfully applying them to business problems in the Amazon Fulfillment Network. These solutions typically utilize machine learning and computer vision techniques, applied to text, sequences of events, images or video from existing or new hardware. We influence each stage of innovation from inception to deployment, developing a research plan, creating and testing prototype solutions, and shepherding the production versions to launch. We are open to hiring candidates to work out of one of the following locations: Bellevue, WA, USA
US, CA, Santa Clara
About Amazon Health Amazon Health’s mission is to make it dramatically easier for customers to access the healthcare products and services they need to get and stay healthy. Towards this mission, we (Health Storefront and Shared Tech) are building the technology, products and services, that help customers find, buy, and engage with the healthcare solutions they need. Job summary We are seeking an exceptional Applied Scientist to join a team of experts in the field of machine learning, and work together to break new ground in the world of healthcare to make personalized and empathetic care accessible, convenient, and cost-effective. We leverage and train state-of-the-art large-language-models (LLMs) and develop entirely new experiences to help customers find the right products and services to address their health needs. We work on machine learning problems for intent detection, dialogue systems, and information retrieval. You will work in a highly collaborative environment where you can pursue both near-term productization opportunities to make immediate, meaningful customer impacts while pursuing ambitious, long-term research. You will work on hard science problems that have not been solved before, conduct rapid prototyping to validate your hypothesis, and deploy your algorithmic ideas at scale. You will get the opportunity to pursue work that makes people's lives better and pushes the envelop of science. Key job responsibilities - Translate product and CX requirements into science metrics and rigorous testing methodologies. - Invent and develop scalable methodologies to evaluate LLM outputs against metrics and guardrails. - Design and implement the best-in-class semantic retrieval system by creating high-quality knowledge base and optimizing embedding models and similarity measures. - Conduct tuning, training, and optimization of LLMs to achieve a compelling CX while reducing operational cost to be scalable. A day in the life In a fast-paced innovation environment, you work closely with product, UX, and business teams to understand customer's challenges. You translate product and business requirements into science problems. You dive deep into challenging science problems, enabling entirely new ML and LLM-driven customer experiences. You identify hypothesis and conduct rapid prototyping to learn quickly. You develop and deploy models at scale to pursue productizations. You mentor junior science team members and help influence our org in scientific best practices. About the team We are the ML Science and Engineering team, with a strong focus on Generative AI. The team consists of top-notch ML Scientists with diverse background in healthcare, robotics, customer analytics, and communication. We are committed to building and deploying the most advanced scientific capabilities and solutions for the products and services at Amazon Health. We are open to hiring candidates to work out of one of the following locations: Santa Clara, CA, USA
US, WA, Seattle
We are designing the future. If you are in quest of an iterative fast-paced environment, where you can drive innovation through scientific inquiry, and provide tangible benefit to hundreds of thousands of our associates worldwide, this is your opportunity. Come work on the Amazon Worldwide Fulfillment Design & Engineering Team! We are looking for an experienced and senior Research Scientist with background in Ergonomics and Industrial Human Factors, someone that is excited to work on complex real-world challenges for which a comprehensive scientific approach is necessary to drive solutions. Your investigations will define human factor / ergonomic thresholds resulting in design and implementation of safe and efficient workspaces and processes for our associates. Your role will entail assessment and design of manual material handling tasks throughout the entire Amazon network. You will identify fundamental questions pertaining to the human capabilities and tolerances in a myriad of work environments, and will initiate and lead studies that will drive decision making on an extreme scale. .You will provide definitive human factors/ ergonomics input and participate in design with every single design group in our network, including Amazon Robotics, Engineering R&D, and Operations Engineering. You will work closely with our Worldwide Health and Safety organization to gain feedback on designs and work tenaciously to continuously improve our associate’s experience. Key job responsibilities - Collaborating and designing work processes and workspaces that adhere to human factors / ergonomics standards worldwide. - Producing comprehensive and assessments of workstations and processes covering biomechanical, physiological, and psychophysical demands. - Effectively communicate your design rationale to multiple engineering and operations entities. - Identifying gaps in current human factors standards and guidelines, and lead comprehensive studies to redefine “industry best practices” based on solid scientific foundations. - Continuously strive to gain in-depth knowledge of your profession, as well as branch out to learn about intersecting fields, such as robotics and mechatronics. - Travelling to our various sites to perform thorough assessments and gain in-depth operational feedback, approximately 25%-50% of the time. We are open to hiring candidates to work out of one of the following locations: Seattle, WA, USA
US, CA, Santa Monica
Amazon Advertising is looking for a motivated and analytical self-starter to help pave the way for the next generation of insights and advertising products. You will use large-scale data, advertising effectiveness knowledge and business information needs of our advertising clients to envision new advertising measurement products and tools. You will facilitate innovation on behalf of our customers through end-to-end delivery of measurement solutions leveraging experiments, machine learning and causal inference. You will partner with our engineering teams to develop and scale successful solutions to production. This role requires strong hands-on skills in terms of effectively working with data, coding, and MLOps. However, the ideal candidate will also bring strong interpersonal and communication skills to engage with cross-functional partners, as well as to stay connected to insights needs of account teams and advertisers. This is a truly exciting and versatile position in that it allows you to apply and develop your hands-on data modeling and coding skills, to work with other scientists on research in new measurement solutions while at the same time partner with cross-functional stakeholders to deliver product impact. Key job responsibilities As an Applied Scientist on the Advertising Incrementality Measurement team you will: - Create new analytical products from conception to prototyping and scaling the product end-to-end through to production. - Scope and define new business problems in the realm of advertising effectiveness. Use machine learning and experiments to develop effective and scalable solutions. - Partner closely with the Engineering team. - Partner with Economists, Data Scientists, and other Applied Scientists to conduct research on advertising effectiveness using machine learning and causal inference. Make findings available via white papers. - Act as a liaison to product teams to help productize new measurement solutions. About the team Advertising Incrementality Measurement combines experiments with econometric analysis and machine learning to provide rigorous causal measurement of advertising effectiveness to internal and external customers. We are open to hiring candidates to work out of one of the following locations: Arlington, VA, USA | Boulder, CO, USA | New York, NY, USA | Santa Monica, CA, USA
US, NY, New York
Amazon Advertising is one of Amazon's fastest growing and most profitable businesses, responsible for defining and delivering a collection of advertising products that drive discovery and sales. Our products are strategically important to our Retail and Marketplace businesses driving long term growth. We deliver billions of ad impressions and millions of clicks and break fresh ground in product and technical innovations every day! The Ad Measurement team develops and deploys solutions fueled by machine learning to support Amazon Advertisers in their strategic campaign planning. Leaning on rich data points, we provide measurements, predictions and diagnostics that separate Amazon Advertising from all other media. As a Data Scientist on this team, you will: - Solve real-world problems by getting and analyzing large amounts of data, diving deep to identify business insights and opportunities, design simulations and experiments, developing statistical and ML models by tailoring to business needs, and collaborating with Scientists, Engineers, BIE's, and Product Managers. - Write code (Python, R, Scala, SQL, etc.) to obtain, manipulate, and analyze data - Apply statistical and machine learning knowledge to specific business problems and data. - Build decision-making models and propose solution for the business problem you define. - Retrieve, synthesize, and present critical data in a format that is immediately useful to answering specific questions or improving system performance. - Analyze historical data to identify trends and support optimal decision making. - Formalize assumptions about how our systems are expected to work, create statistical definition of the outlier, and develop methods to systematically identify outliers. Work out why such examples are outliers and define if any actions needed. - Given anecdotes about anomalies or generate automatic scripts to define anomalies, deep dive to explain why they happen, and identify fixes. - Conduct written and verbal presentations to share insights to audiences of varying levels of technical sophistication. Why you will love this opportunity: Amazon has invested heavily in building a world-class advertising business. This team defines and delivers a collection of advertising products that drive discovery and sales. Our solutions generate billions in revenue and drive long-term growth for Amazon’s Retail and Marketplace businesses. We deliver billions of ad impressions, millions of clicks daily, and break fresh ground to create world-class products. We are a highly motivated, collaborative, and fun-loving team with an entrepreneurial spirit - with a broad mandate to experiment and innovate. Impact and Career Growth: You will invent new experiences and influence customer-facing shopping experiences to help suppliers grow their retail business and the auction dynamics that leverage native advertising; this is your opportunity to work within the fastest-growing businesses across all of Amazon! Define a long-term science vision for our advertising business, driven from our customers' needs, translating that direction into specific plans for research and applied scientists, as well as engineering and product teams. This role combines science leadership, organizational ability, technical strength, product focus, and business understanding. We are open to hiring candidates to work out of one of the following locations: New York, NY, USA
US, WA, Bellevue
At AWS, we use Artificial Intelligence to be able to identify every need of a customer across all AWS services before they have to tell us about it, and then find and seamlessly connect them to the most appropriate resolution for their need, eventually fulfilling the vision of a self-healing cloud. We are looking for Data Scientists with unfettered curiosity and drive to help build “best in the world” support (contact center) experience that customers will love! You will have an opportunity to lead, invent, and design tech that will directly impact every customer across all AWS services. We are building industry-leading technology that cuts across a wide range of ML techniques from Natural Language Processing to Deep Learning and Generative Artificial Intelligence. You will be a key driver in taking something from an idea to an experiment to a prototype and finally to a live production system. Our team packs a punch with principal level engineering, science, product, and leadership talent. We are a results focused team and you have the opportunity to lead and establish a culture for the big things to come. We combine the culture of a startup, the innovation and creativity of a R&D Lab, the work-life balance of a mature organization, and technical challenges at the scale of AWS. We offer a playground of opportunities for builders to build, have fun, and make history! Key job responsibilities Deliver real world production systems at AWS scale. Work closely with the business to understand the problem space, identify the opportunities and formulate the problems. Use machine learning, data mining, statistical techniques, Generative AI and others to create actionable, meaningful, and scalable solutions for the business problems. Analyze and extract relevant information from large amounts of data and derive useful insights. Work with software engineering teams to deliver production systems with your ML models Establish scalable, efficient, automated processes for large scale data analyses, model development, model validation and model implementation We are open to hiring candidates to work out of one of the following locations: Bellevue, WA, USA | Seattle, WA, USA
US, CA, Santa Clara
Amazon launched the Generative AI Innovation Center (GAIIC) in Jun 2023 to help AWS customers accelerate the use of Generative AI to solve business and operational problems and promote innovation in their organization (https://press.aboutamazon.com/2023/6/aws-announces-generative-ai-innovation-center). GAIIC provides opportunities to innovate in a fast-paced organization that contributes to game-changing projects and technologies that get deployed on devices and in the cloud. As an Applied Science Manager in GAIIC, you'll partner with technology and business teams to build new GenAI solutions that delight our customers. You will be responsible for directing a team of data/research/applied scientists, deep learning architects, and ML engineers to build generative AI models and pipelines, and deliver state-of-the-art solutions to customer’s business and mission problems. Your team will be working with terabytes of text, images, and other types of data to address real-world problems. The successful candidate will possess both technical and customer-facing skills that will allow you to be the technical “face” of AWS within our solution providers’ ecosystem/environment as well as directly to end customers. You will be able to drive discussions with senior technical and management personnel within customers and partners, as well as the technical background that enables them to interact with and give guidance to data/research/applied scientists and software developers. The ideal candidate will also have a demonstrated ability to think strategically about business, product, and technical issues. Finally, and of critical importance, the candidate will be an excellent technical team manager, someone who knows how to hire, develop, and retain high quality technical talent. About the team Here at AWS, it’s in our nature to learn and be curious about diverse perspectives. Our employee-led affinity groups foster a culture of inclusion that empower employees to feel proud of our differences. Ongoing events and learning experiences, including our Conversations on Race and Ethnicity (CORE) and AmazeCon (gender diversity) conferences, inspire us to never stop embracing our uniqueness. We have a career path for you no matter what stage you’re in when you start here. We’re continuously raising our performance bar as we strive to become Earth’s Best Employer. That’s why you’ll find endless knowledge-sharing, mentorship and other career- advancing resources here to help you develop into a better-rounded professional. We are open to hiring candidates to work out of one of the following locations: San Francisco, CA, USA | San Jose, CA, USA | Santa Clara, CA, USA
GB, London
Amazon Advertising is looking for a Data Scientist to join its brand new initiative that powers Amazon’s contextual advertising products. Advertising at Amazon is a fast-growing multi-billion dollar business that spans across desktop, mobile and connected devices; encompasses ads on Amazon and a vast network of hundreds of thousands of third party publishers; and extends across US, EU and an increasing number of international geographies. The Supply Quality organization has the charter to solve optimization problems for ad-programs in Amazon and ensure high-quality ad-impressions. We develop advanced algorithms and infrastructure systems to optimize performance for our advertisers and publishers. We are focused on solving a wide variety of problems in computational advertising like traffic quality prediction (robot and fraud detection), Security forensics and research, Viewability prediction, Brand Safety, Contextual data processing and classification. Our team includes experts in the areas of distributed computing, machine learning, statistics, optimization, text mining, information theory and big data systems. We are looking for a dynamic, innovative and accomplished Data Scientist to work on data science initiatives for contextual data processing and classification that power our contextual advertising solutions. Are you an experienced user of sophisticated analytical techniques that can be applied to answer business questions and chart a sustainable vision? Are you exited by the prospect of communicating insights and recommendations to audiences of varying levels of technical sophistication? Above all, are you an innovator at heart and have a track record of resolving ambiguity to deliver result? As a data scientist, you help our data science team build cutting edge models and measurement solutions to power our contextual classification technology. As this is a new initiative, you will get an opportunity to act as a thought leader, work backwards from the customer needs, dive deep into data to understand the issues, define metrics, conceptualize and build algorithms and collaborate with multiple cross-functional teams. Key job responsibilities * Define a long-term science vision for contextual-classification tech, driven fundamentally from the needs of our advertisers and publishers, translating that direction into specific plans for the science team. Interpret complex and interrelated data points and anecdotes to build and communicate this vision. * Collaborate with software engineering teams to Identify and implement elegant statistical and machine learning solutions * Oversee the design, development, and implementation of production level code that handles billions of ad requests. Own the full development cycle: idea, design, prototype, impact assessment, A/B testing (including interpretation of results) and production deployment. * Promote the culture of experimentation and applied science at Amazon. * Demonstrated ability to meet deadlines while managing multiple projects. * Excellent communication and presentation skills working with multiple peer groups and different levels of management * Influence and continuously improve a sustainable team culture that exemplifies Amazon’s leadership principles. We are open to hiring candidates to work out of one of the following locations: London, GBR
JP, 13, Tokyo
We are seeking a Principal Economist to be the science leader in Amazon's customer growth and engagement. The wide remit covers Prime, delivery experiences, loyalty program (Amazon Points), and marketing. We look forward to partnering with you to advance our innovation on customers’ behalf. Amazon has a trailblazing track record of working with Ph.D. economists in the tech industry and offers a unique environment for economists to thrive. As an economist at Amazon, you will apply the frontier of econometric and economic methods to Amazon’s terabytes of data and intriguing customer problems. Your expertise in building reduced-form or structural causal inference models is exemplary in Amazon. Your strategic thinking in designing mechanisms and products influences how Amazon evolves. In this role, you will build ground-breaking, state-of-the-art econometric models to guide multi-billion-dollar investment decisions around the global Amazon marketplaces. You will own, execute, and expand a research roadmap that connects science, business, and engineering and contributes to Amazon's long term success. As one of the first economists outside North America/EU, you will make an outsized impact to our international marketplaces and pioneer in expanding Amazon’s economist community in Asia. The ideal candidate will be an experienced economist in empirical industrial organization, labour economics, or related structural/reduced-form causal inference fields. You are a self-starter who enjoys ambiguity in a fast-paced and ever-changing environment. You think big on the next game-changing opportunity but also dive deep into every detail that matters. You insist on the highest standards and are consistent in delivering results. Key job responsibilities - Work with Product, Finance, Data Science, and Data Engineering teams across the globe to deliver data-driven insights and products for regional and world-wide launches. - Innovate on how Amazon can leverage data analytics to better serve our customers through selection and pricing. - Contribute to building a strong data science community in Amazon Asia. We are open to hiring candidates to work out of one of the following locations: Tokyo, 13, JPN