NeurIPS: Why causal-representation learning may be the future of AI

Francesco Locatello on the four NeurIPS papers he coauthored this year, which largely concern generalization to out-of-distribution test data.

In a conversation right before the 2021 Conference on Neural Information Processing Systems (NeurIPS), Amazon vice president and distinguished scientist Bernhard Schölkopf — according to Google Scholar, the most highly cited researcher in the field of causal inference — said that the next frontier in artificial-intelligence research was causal-representation learning.

Where existing approaches to causal inference use machine learning to discover causal relationships between variables — say, the latencies of various interrelated services on a website — causal-representation learning learns the variables themselves. “These kinds of causal representations will also go toward reasoning, which we will ultimately need if we want to move away from this pure pattern recognition view of intelligence,” Schölkopf said.

Francesco.jpg
Senior applied scientist Francesco Locatello.

Francesco Locatello, a senior applied scientist with Amazon Web Services, leads Amazon’s research on causal-representation learning, and he’s a coauthor on four papers at this year’s NeurIPS.

Assaying out-of-distribution generalization in transfer learning” concerns one of the most compelling applications of causal inference in machine learning: generalizing models trained on data with a particular probability distribution to real-world data with a different distribution.

“When you do standard machine learning, you are drawing independent samples from some probability distribution, and then you train a model that's going to generalize to the same distribution,” Locatello explains. “You're describing a physical system using a single probability distribution. Causal models are different because they model every possible state that this physical system can take as a result of an intervention. So instead of having a single probability distribution, you have a set of distributions.

Related content
Amazon Science hosts a conversation with Amazon Scholars Michael I. Jordan and Michael Kearns and Amazon distinguished scientist Bernhard Schölkopf.

“What does it mean that your test data comes from a different distribution? You have the same underlying physical system; the causal structure is the same. It's just a new intervention you have not seen. Your test distribution is different than the training, but now it's not an arbitrary distribution. It’s well posed because it's entailed by the causal structure, and it's a meaningful distribution that may happen in the real world.”

In “Assaying out-of-distribution generalization in transfer learning”, Locatello explains, “what we do is to collect a huge variety of datasets that are constructed for or adapted to this scenario where you have a very narrow data set that you can use for transfer learning, and then you have a wide variety of test data that is all out of distribution. We look at the different approaches that have been studied in the literature and compare them on fair ground.”

Although none of the approaches canvassed in the paper explicitly considers causality, Locatello says, “causal approaches should eventually be able to do better on this benchmark, and this will allow us to evaluate our progress. That's why we built it.”

Neural circuits

Today’s neural networks do representation learning as a matter of course: their inputs are usually raw data, and they learn during training which aspects of the data are most useful for the task at hand. As Schölkopf pointed out in conversation last year, causal-representation learning would simply bring causal machine learning models up to speed with conventional models.

Related content
New method goes beyond Granger causality to identify only the true causes of a target time series, given some graph constraints.

“The important thing to realize is that most machine learning applications don't come structured as a set of well-defined random variables that fully align with the underlying functioning of a physical system,” Locatello explains. “We still want to model these systems in terms of abstract variables, but nobody gives these variables to us. So you may want to learn them in order to be able to perform causal inference.”

Among his and his colleagues’ NeurIPS papers, Locatello says, the one that comes closest to the topic of causal-representation learning is “Neural attentive circuits”. Causal models typically represent causal relationships using graphs, and a neural network, too, can be thought of as an enormous graph. Locatello and his collaborators are trying to make that analogy explicit, by training a neural network to mimic the structure of a causal network.

Neural attentive circuits.png
Visualizations of graph structures learned by neural attentive circuits, from "Neural attentive circuits".

“This is a follow-up on a paper we had last year in NeurIPS,” Locatello says. “The inspiration was to design architectures that behave more similarly to causal models, where you have the noise variables — that's the data — and then you have variables that are being manipulated by functions, and they simply communicate with each other in a graph. And this graph can change dynamically when a distribution changes, for example, because of an intervention.

“In the first paper, we developed an architecture that behaves exactly like that: you have a set of neural functions that can be composed on the fly, depending on the data and the problem. The functions, the routing, and the stitching of the functions are learned. Everything is learned. But it turns out that dynamic stitching is not very scalable.

“In this new work, we essentially compiled the stitching of the functions so that for each sample it's decided beforehand — where it's going to go through the network, how the functions are going to be composed. Instead of doing it on the fly one layer at a time, you decide for the overall forward pass. And we demonstrated that these sparse learned connectivity patterns improve out-of-distribution generalization.”

Success stories

Locatello’s other NeurIPS papers are on more-conventional machine learning topics. “Self supervised amodal video object segmentation” considers the problem of reconstructing the silhouette of an occluded object, which is crucial to robotics applications, including autonomous cars.

Locatello 16_9.png
Segmentations of partially occluded objects, from "Self supervised amodal video object segmentation".

“We exploit the principle that you can build information about an object over time in a video,” Locatello explains. “Perhaps in past frames you've seen parts of the objects that are now occluded. If you can remember that you've seen this object before, and this was its segmentation mask, you can build up your segmentation over time.”

The final paper, “Are two heads the same as one? Identifying disparate treatment in fair neural networks”, considers models whose training objectives are explicitly designed to minimize bias across different types of inputs. Locatello and his colleagues find that frequently, such models — purely through training, without any human intervention — develop two “heads”: that is, they learn two different pathways through the neural network, one for inputs in the sensitive class, and one for all other inputs.

Related content
Amazon ICML paper proposes information-theoretic measurement of quantitative causal contribution.

The researchers argue that, since the network is learning two heads, anyway, it might as well be designed with a two-headed architecture: that would improve performance while meeting the same fairness standard. But this approach hasn’t been adopted, as it runs afoul of rules prohibiting disparate treatment of different groups. In this case, however, disparate treatment could be the best way to ensure fair treatment.

These last two papers are only obliquely related to causality. But, Locatello says, “causal-representation learning is a very young field. So we are trying to identify success stories, and I think these papers are going in that direction.”

“It's clear that causality will have a role in future machine learning,” he adds, “because there are a lot of open problems in machine learning that can at least be partially addressed when you start looking at causal models. And my goal really is to realize the benefits of causal models in mainstream machine learning applications. That's why some of these works are not necessarily about causality, but closer to machine learning. Because ultimately, that's our goal.”

Learn more about Amazon at NeurIPS 2022

For more on the Amazon research being presented at this year's NeurIPS, see our quick guide to Amazon's NeurIPS 2022 papers.

Related content

GB, Cambridge
Our team builds generative AI solutions that will produce some of the future’s most influential voices in media and art. We develop cutting-edge technologies with Amazon Studios, the provider of original content for Prime Video, with Amazon Game Studios and Alexa, the ground-breaking service that powers the audio for Echo. Do you want to be part of the team developing the future technology that impacts the customer experience of ground-breaking products? Then come join us and make history. We are looking for a passionate, talented, and inventive Applied Scientist with a background in Machine Learning to help build industry-leading Speech, Language, Audio and Video technology. As an Applied Scientist at Amazon you will work with talented peers to develop novel algorithms and generative AI models to drive the state of the art in audio (and vocal arts) generation. Position Responsibilities: * Participate in the design, development, evaluation, deployment and updating of data-driven models for digital vocal arts applications. * Participate in research activities including the application and evaluation and digital vocal and video arts techniques for novel applications. * Research and implement novel ML and statistical approaches to add value to the business. * Mentor junior engineers and scientists. We are open to hiring candidates to work out of one of the following locations: Cambridge, GBR | London, GBR
GB, Cambridge
Our team undertakes research together with multiple organizations to advance the state-of-the-art in speech technologies. We not only work on giving Alexa, the ground-breaking service that powers Echo, her voice, but we also develop cutting-edge technologies with Amazon Studios, the provider of original content for Prime Video. Do you want to be part of the team developing the latest technology that impacts the customer experience of ground-breaking products? Then come join us and make history. We are looking for a passionate, talented, and inventive Senior Applied Scientist with a background in Machine Learning to help build industry-leading Speech, Language and Video technology. As a Senior Applied Scientist at Amazon you will work with talented peers to develop novel algorithms and modelling techniques to drive the state of the art in speech and vocal arts synthesis. Position Responsibilities: * Participate in the design, development, evaluation, deployment and updating of data-driven models for digital vocal arts applications. * Participate in research activities including the application and evaluation and digital vocal and video arts techniques for novel applications. * Research and implement novel ML and statistical approaches to add value to the business. * Mentor junior engineers and scientists. We are open to hiring candidates to work out of one of the following locations: Cambridge, GBR | London, GBR
CN, 11, Beijing
Are you interested in applying your strong quantitative analysis and big data skills to world-changing problems? Are you interested in driving the development of methods, models and systems for strategy planning, transportation and fulfillment network? Are you interested to cooperate with Amazonians around the world? If so, then this is the job for you. Our team, ATE(Analytics Technology and Engineering) is looking for an Applied Scientist to join our growing Science Team in Bangalore (India)/ Beijing(China). We are responsible for creating core analytics tech capabilities, quantative models, platforms development, and data engineering. We develop scalable analytics applications and research models to optimize operations processes. We standardize and optimize data sources and visualization efforts across geographies, build up, and maintain the online business intelligence services and data mart. You will work with other scientists, professional data engineers, business intelligence engineers, and product managers using rigorous quantitative approaches to ensure high quality data tech products for our customers around the world, including India, Australia, Brazil, Mexico, Singapore and Middle East. Amazon is growing rapidly and because we are driven by faster delivery to customers, a more efficient supply chain network, and lower cost of operations, our main focus is in the development of strategic models and automation tools fed by our massive amounts of available data. You will be responsible for building these models/tools that improve the economics of Amazon’s worldwide fulfillment networks in different countries as Amazon increases the speed and decreases the cost to deliver products to customers. You will work on large-scale vehicle routing and scheduling problems under complex operational and physical constraints. You will also identify and evaluate opportunities to reduce variable costs by improving fulfillment center processes, transportation operations and scheduling, and the execution of operational plans. Finally, you will help create the metrics to quantify improvements to the fulfillment costs (e.g., transportation and labor costs) resulting from the application of these optimization models and tools. Key job responsibilities - Design and develop complex mathematical, simulation and optimization models and apply them to define strategic and tactical needs and drive the appropriate business and technical solutions in the areas of vehicle routing, inventory management, network flow, supply chain optimization, demand planning. - Apply theories of mathematical optimization, including linear programming, combinatorial optimization, integer programming, dynamic programming, network flows and algorithms to design optimal or near optimal solution methodologies to be used by in-house decision support tools and software. - Translating business questions and concerns into specific analytical questions that can be answered with available data using Statistical and Machine Learning methods. - Prototype models by using modeling and programming languages with efficient data querying and modeling infrastructure. - Communicate proposals and results in a clear manner backed by data and coupled with actionable conclusions to drive business decisions. - Collaborate with colleagues from multidisciplinary science, engineering and business backgrounds. - Manage your own process. Prioritize and execute on high impact projects, triage external requests, and ensure to deliver projects in time. We are open to hiring candidates to work out of one of the following locations: Beijing, 11, CHN
GB, Cambridge
Our team builds generative AI solutions that will produce some of the future’s most influential voices in media and art. We develop cutting-edge technologies with Amazon Studios, the provider of original content for Prime Video, with Amazon Game Studios and Alexa, the ground-breaking service that powers the audio for Echo. Do you want to be part of the team developing the future technology that impacts the customer experience of ground-breaking products? Then come join us and make history. We are looking for a passionate, talented, and inventive Applied Scientist with a background in Machine Learning to help build industry-leading Speech, Language, Audio and Video technology. As an Applied Scientist at Amazon you will work with talented peers to develop novel algorithms and generative AI models to drive the state of the art in audio (and vocal arts) generation. Position Responsibilities: * Participate in the design, development, evaluation, deployment and updating of data-driven models for digital vocal arts applications. * Participate in research activities including the application and evaluation and digital vocal and video arts techniques for novel applications. * Research and implement novel ML and statistical approaches to add value to the business. * Mentor junior engineers and scientists. We are open to hiring candidates to work out of one of the following locations: Cambridge, GBR | London, GBR
US, WA, Bellevue
Do you enjoy solving challenging problems and driving innovations in research? Are you seeking for an environment with a group of motivated and talented scientists like yourself? Do you want to create scalable optimization models and apply machine learning techniques to guide real-world decisions? Do you want to play a key role in the future of Amazon transportation and operations? North America Sort Centers (NASC) are experiencing growth and looking for a skilled, highly motivated Research Scientist in partnership with the Modeling and Optimization (MOP) team. The Sort Center network is the critical Middle-Mile solution in the Amazon Transportation Services (ATS) group, linking Fulfillment Centers to the Last Mile. The experience of our customers is dependent on our ability to efficiently execute volume flow through the middle-mile network. Key job responsibilities A Research Scientist - provides analytical decision support to Amazon planning teams via applying advanced mathematical and statistical techniques. - collaborates effectively with Amazon internal business customers, and is their trusted partner - is proactive and independent in discovering and resolving business pain-points within a given scope - is able to identify a suitable level of sophistication in resolving the different business needs - is confident in leveraging existing solutions to new problems where appropriate and is independent in designing and implementing new solutions where needed - is aware of the limitations of their proposed solutions and is proactive in communicating them to the business, and advances the application of sciences towards Amazon business problems by bringing new methods, ideas, and practices to the team and scientific community. A day in the life - Your will be developing model-based optimization, simulation, and/or predictive tools to identify and evaluate opportunities to improve customer experience, network speed, cost, and efficiency of capital investment. - You will quantify the improvements resulting from the application of these tools and you will evaluate the trade-offs between potentially competing objectives. - You will develop good communication skills and ability to speak at a level appropriate for the audience, will collaborate effectively with fellow scientists, software development engineers, and product managers, and will deliver business value in a close partnership with many stakeholders from operations, finance, IT, and business leadership. About the team - At the Modeling and Optimization (MOP) team, we use mathematical optimization, algorithm design, statistics, and machine learning to improve decision-making capabilities across WW Operations from first mile to last mile. - We focus on transportation topology, labor and resource planning for fulfillment facilities, routing science, visualization research, data science and development, and process optimization. - We create models to simulate, optimize, and control the fulfillment network with the objective of reducing cost while improving speed and reliability. - We support multiple business lanes, therefore maintain a comprehensive and objective view, coordinating solutions across organizational lines where possible. We are open to hiring candidates to work out of one of the following locations: Bellevue, WA, USA
US, MA, Cambridge
The Artificial General Intelligence (AGI) team is looking for a highly skilled and experienced Senior Applied Scientist, to lead the development and implementation of cutting-edge algorithms and models for supervised fine-tuning and reinforcement learning through human feedback; with a focus across text, image, and video modalities. As a Senior Applied Scientist, you will play a critical role in driving the development of Generative AI (GenAI) technologies that can handle Amazon-scale use cases and have a significant impact on our customers' experiences. Key job responsibilities - Collaborate with cross-functional teams of engineers, product managers, and scientists to identify and solve complex problems in GenAI - Design and execute experiments to evaluate the performance of different algorithms and models, and iterate quickly to improve results - Think big about the arc of development of GenAI over a multi-year horizon, and identify new opportunities to apply these technologies to solve real-world problems - Communicate results and insights to both technical and non-technical audiences, including through presentations and written reports - Mentor and guide junior scientists and engineers, and contribute to the overall growth and development of the team We are open to hiring candidates to work out of one of the following locations: Bellevue, WA, USA | Cambridge, MA, USA | New York, NY, USA | Sunnyvale, CA, USA
US, MA, North Reading
We are looking for experienced scientists and engineers to explore new ideas, invent new approaches, and develop new solutions in the areas of Controls, Dynamic modeling and System identification. Are you inspired by invention? Is problem solving through teamwork in your DNA? Do you like the idea of seeing how your work impacts the bigger picture? Answer yes to any of these and you’ll fit right in here at Amazon Robotics. We are a smart team of doers that work passionately to apply cutting edge advances in robotics and software to solve real-world challenges that will transform our customers’ experiences in ways we can’t even imagine yet. We invent new improvements every day. We are Amazon Robotics and we will give you the tools and support you need to invent with us in ways that are rewarding, fulfilling and fun. Key job responsibilities Applied Scientists take on big unanswered questions and guide development team to state-of-the-art solutions. We want to hear from you if you have deep industry experience in the Mechatronics domain and : * the ability to think big and conceive of new ideas and novel solutions; * the insight to correctly identify those worth exploring; * the hands-on skills to quickly develop proofs-of-concept; * the rigor to conduct careful experimental evaluations; * the discipline to fast-fail when data refutes theory; * and the fortitude to continue exploring until your solution is found We are open to hiring candidates to work out of one of the following locations: North Reading, MA, USA | Westborough, MA, USA
IL, Tel Aviv
Come build the future of entertainment with us. Are you interested in helping shape the future of movies and television? Do you want to help define the next generation of how and what Amazon customers are watching? Prime Video is a premium streaming service that offers customers a vast collection of TV shows and movies - all with the ease of finding what they love to watch in one place. We offer customers thousands of popular movies and TV shows including Amazon Originals and exclusive licensed content to exciting live sports events. We also offer our members the opportunity to subscribe to add-on channels which they can cancel at anytime and to rent or buy new release movies and TV box sets on the Prime Video Store. Prime Video is a fast-paced, growth business - available in over 240 countries and territories worldwide. The team works in a dynamic environment where innovating on behalf of our customers is at the heart of everything we do. If this sounds exciting to you, please read on. We are looking for an Applied Scientist to embark on our journey to build a Prime Video Sports tech team in Israel from ground up. Our team will focus on developing products to allow for personalizing the customers’ experience and providing them real-time insights and revolutionary experiences using Computer Vision (CV) and Machine Learning (ML). You will get a chance to work on greenfield, cutting-edge and large-scale engineering and science projects, and a rare opportunity to be one of the founders of the Israel Prime Video Sports tech team in Israel. Key job responsibilities We are looking for an Applied Scientist with domain expertise in Computer Vision or Recommendation Systems to lead development of new algorithms and E2E solutions. You will be part of a team of applied scientists and software development engineers responsible for research, design, development and deployment of algorithms into production pipelines. As a technologist, you will also drive publications of original work in top-tier conferences in Computer Vision and Machine Learning. You will be expected to deal with ambiguity! We're looking for someone with outstanding analytical abilities and someone comfortable working with cross-functional teams and systems. You must be a self-starter and be able to learn on the go. About the team In September 2018 Prime Video launched its first full-scale live streaming experience to world-wide Prime customers with NFL Thursday Night Football. That was just the start. Now Amazon has exclusive broadcasting rights to major leagues like NFL Thursday Night Football, Tennis major like Roland-Garros and English Premium League to list few and are broadcasting live events across 30+ sports world-wide. Prime Video is expanding not just the breadth of live content that it offers, but the depth of the experience. This is a transformative opportunity, the chance to be at the vanguard of a program that will revolutionize Prime Video, and the live streaming experience of customers everywhere. We are open to hiring candidates to work out of one of the following locations: Tel Aviv, ISR
MX, DIF, Mexico City
Are you a data enthusiast? Does the world’s most complex logistic systems inspire your curiosity? Is your passion to navigate through hundreds of systems, processes, and data sources to solve the puzzles and identify the next big opportunity? Are you a creative big thinker who is passionate about using data and optimization tools to direct decision making and solve complex and large-scale challenges? Do you feel like your skills uniquely qualify you to bridge communication between teams with competing priorities? If so, then this position is for you! We are looking for a motivated individual with strong analytic and communication skills to join the effort in evolving the network we have today into the network we need tomorrow. Amazon’s extensive logistics system is comprised of thousands of fixed infrastructure nodes, with millions of possible connections between them. Billions of packages flow through this network on a yearly basis, making the impact of optimal improvements unparalleled. This magnificent challenge is a terrific opportunity to analyze Amazon’s data and generate actionable recommendations using optimization and simulation. Come build with us! In this role, your main focus will be to perform analysis, synthesize information, identify business opportunities, provide project direction, and communicate business and technical requirements within the team and across stakeholder groups. You consider the needs of day-to-day operations and insist on the standards required to build the network of tomorrow. You will assist in defining trade-offs and quantifying opportunities for a variety of projects. You will learn current processes, build metrics, educate diverse stakeholder groups, assist science groups in initial solution design, and audit all model implementation. A successful candidate in this position will have a background in communicating across significant differences, prioritizing competing requests, and quantifying decisions made. The ideal candidate will have a strong ability to model real world data with high complexity and delivery high quality analysis, data products and optimizations models for strategic decision. They are excited to be part of, and learn from, a large science community and are ready to dig into the details to find insights that direct decisions. The successful candidate will have good communication skills and an ability to speak at a level appropriate for the audience, will collaborate effectively with scientists, product managers and business stakeholders. Key job responsibilities Statistical Models (ML, regression, forecasting, ) Optimization models, AB and hypothesis testing, Bayesian models. Communication skills with both tech and non tech stakeholders. Writting skills, capable to create documents for different types of readers (business, science, tech) to communicate results on analysis, testing. A day in the life We are open to hiring candidates to work out of one of the following locations: Mexico City, DIF, MEX
US, WA, Seattle
We are seeking a Senior Applied Scientist to join our AI Security team and use AI to develop foundational services that make security mechanisms more effective and efficient. As a Senior Applied Scientist, you will be responsible for researching, modeling, designing, and implementing state-of-the-art AI-based security solutions at Amazon scale. You will collaborate with applied scientists, security engineers, software engineers, as well as internal stakeholders and partners to develop innovative technologies to solve some of our hardest security problems, and build paved path solutions that support builder teams across Amazon throughout their software development journey, enabling Amazon businesses to accelerate the pace of innovation to delight our customers. Key job responsibilities • Research and develop accurate and scalable methods to solve foundational security problems. • Lead and partner with applied scientists and engineers to drive modeling and technical design for complex problems. • Build security tooling and paved path solutions that support builder teams throughout their software development journey. About the team ABOUT AmSec: Diverse Experiences: Amazon Security values diverse experiences. Even if you do not meet all of the qualifications and skills listed in the job description, we encourage candidates to apply. If your career is just starting, hasn’t followed a traditional path, or includes alternative experiences, don’t let it stop you from applying. Why Amazon Security: At Amazon, security is central to maintaining customer trust and delivering delightful customer experiences. Our organization is responsible for creating and maintaining a high bar for security across all of Amazon’s products and services. We offer talented security professionals the chance to accelerate their careers with opportunities to build experience in a wide variety of areas including cloud, devices, retail, entertainment, healthcare, operations, and physical stores. Inclusive Team Culture: In Amazon Security, it’s in our nature to learn and be curious. Ongoing DEI events and learning experiences inspire us to continue learning and to embrace our uniqueness. Addressing the toughest security challenges requires that we seek out and celebrate a diversity of ideas, perspectives, and voices. Training & Career Growth: 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, training, and other career-advancing resources here to help you develop into a better-rounded professional. Work/Life Balance: We value work-life harmony. Achieving success at work should never come at the expense of sacrifices at home, which is why flexible work hours and arrangements are part of our culture. When we feel supported in the workplace and at home, there’s nothing we can’t achieve. We are open to hiring candidates to work out of one of the following locations: Austin, TX, USA | San Francisco, CA, USA | Seattle, WA, USA