Using supervised learning to train models for image clustering

Approach that uses a hierarchical graph neural network improves F-score by 49% relative to predecessors.

Most machine learning models use supervised learning, meaning they’re trained on annotated data, which is costly and time consuming to acquire.

The chief method for doing unsupervised learning, which doesn’t require annotated data, is clustering, or grouping data points together by salient characteristics. The idea is that each cluster represents some category, such as photos of the same person or the same species of animal.

To decide where to draw boundaries between clusters, clustering algorithms typically rely on heuristics, such as a threshold distance between cluster centers or the shape of the clusters’ distributions. In a paper we’re presenting at the International Conference on Computer Vision (ICCV), we propose, instead, to learn from data how to draw boundaries.

We first represent visual data using a graph, then use a graph neural network (GNN) to produce vector representations of the graph’s nodes. So far, we follow on previous work.

Instead of relying on heuristics, however, we use labeled data to learn how to cluster the vectors and, crucially, to decide how fine-grained those clusters should be. We call the labeled data meta-training data, since the goal is to learn a general clustering technique, not a specific classification model. 

In particular, we propose a hierarchical GNN, meaning that it creates clusters by adding edges between nodes of a graph, then adds edges between the clusters to create still larger clusters, and so on, iterating until it decides that no more edges should be added.

Hierarchical clustering.png
A schematic of our graph-based hierarchical clustering approach. The colors of the image borders and of the graph nodes indicate data types (in this case, photos of the same actor). Our approach is hierarchical, iteratively treating small clusters generated at one level as the units of clustering for the next level. We call our base model LANDER, for link approximation and density estimation refinement, and our hierarchical clustering method Hi-LANDER.

Finally, we apply our hierarchical clustering technique to test sets whose classification categories are disjoint with those of the meta-training data. In our experiments we found that, compared to previous GNN-based supervised and unsupervised approaches, ours increased the F-score — which factors in both false positives and false negatives — by an average of 49% and 47%, respectively.

Constructing the graph

In our paper, we investigate the case in which we are training a model to cluster visual data that is similar to the meta-training data but has no class overlaps with it. For instance, the meta-training data might be faces of movie stars, while the target application is to cluster faces of politicians, athletes, or other public figures.

The first step in our process is to use the meta-training data to build a supervised classifier: if the meta-training data is faces of movie stars, the classifier labels input images with names of movie stars.

The classifier is an encoder-decoder model: the encoder produces a fixed-length vector representation of the input, or feature vector, and the decoder uses that vector to predict a label. Once we’ve trained the classifier, however, we use only the encoder for the rest of the process.

The feature vectors define points in a multidimensional space. On the basis of the vectors’ locations, we construct a graph, in which each node represents an image, and each image’s k nearest neighbors in the feature space are connected to it (share edges with it) in the graph.

This graph will serve as the input to the clustering model, which is also an encoder-decoder model. The encoder is a GNN, which produces a vector representation of each node in the graph, based on that node’s feature vector and those of the nodes it’s connected to. Call this vector the node embedding.

The clustering model

We adopt a hierarchical approach to clustering. Based on the node embeddings, the clustering model predicts edges between nodes. A cluster is defined as a group of nodes each of which shares an edge with at least one other node in the group and none of which shares an edge with any node outside the group.

Note that the goal of the clustering model is not just to reproduce the nearest-neighbor graph but to link nodes that represent data of the same type. The nearest-neighbor linkages are useful for predicting clustering linkages, but they are not identical with them.

After the first pass through the data, we aggregate each cluster into a single, representative “supernode” and repeat the whole process. That is, we create edges between each supernode and its k nearest neighbors, pass the resulting graph through the same GNN, and predict edges based on the supernode embeddings. We repeat this process until the clustering model predicts no edges between nodes.

We train our clustering model on two different objectives. One is to correctly predict links between nodes, where a correct link is one that picks out two representatives of the same data type in the meta-training data (say, two photos of the same actor).

We also train the model to correctly predict the density of a given data type in a given graph neighborhood. That is, for each node, the model should predict the proportion of nearby neighbors of the same data type.

Past research on clustering has shown that factoring in data density improves results. Previously, however, link prediction and data density prediction were handled by separate models. By using a single model to jointly predict both, we significantly increase computational efficiency. We believe that the combination also contributes to our increase in accuracy.

The other novelty of our approach is that, because of our hierarchical processing scheme, we optimize clustering across the entire input graph. Previous approaches would first divide the graph into subgraphs, then perform inference within subgraphs. This prevents natural parallelization, which is runtime efficient, and limits the effectiveness of information flow through the graph. The full graph-wide processing is another reason for our model’s improved efficiency.

In experiments, we considered two different sets of meta-training data. One consisted of closeups of human faces, the other of images of particular animal species. We tested the model trained on human faces on two other datasets, whose data categories had zero or very little overlap with those of the meta-training set — 0% and less than 2%. We tested the model trained on animal species on a dataset of previously unseen species. Across both models and the three test sets, our average improvements over previous GNN-based clustering models and unsupervised clustering methods were 49% and 47%, respectively.

In ongoing work, we are investigating the possibility training a more general clustering model, whose performance at inference time will be more transferrable across different data types — accurately clustering both faces and animal species, for instance.

Acknowledgements: Tianjun Xiao, Yongxin Wang, Yuanjun Xiong, Wei Xia, David Wipf, Zhang Zheng, Stefano Soatto

Related content

US, CA, Pasadena
The Amazon Center for Quantum Computing in Pasadena, CA, is looking to hire an Applied Scientist specializing in Testing of Control Systems hardware. Working alongside other scientists and engineers, you will validate hardware and software systems performing the control and readout functions for Amazon quantum processors. Working effectively within a cross-functional team environment is critical. The ideal candidate will have an established background in test engineering applicable to large mixed-signal systems. Diverse Experiences Amazon values diverse experiences. Even if you do not meet all of the preferred 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. 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 we strive for flexibility as part of our working culture. When we feel supported in the workplace and at home, there’s nothing we can’t achieve in the cloud. Inclusive Team Culture Here at AWS, it’s in our nature to learn and be curious. Our employee-led affinity groups foster a culture of inclusion that empower us to be 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. Mentorship and 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, mentorship and other career-advancing resources here to help you develop into a better-rounded professional. Key job responsibilities Our scientists and engineers collaborate across diverse teams and projects to offer state of the art, cost effective solutions for the control of Amazon quantum processor systems. You’ll bring a passion for innovation and collaboration to: Develop automated test scripts for mid-volume electronics manufacturing, utilizing high-speed test equipment such as Gsps oscilloscopes, logic analyzers, and network analyzers. Design and implement test plans for high-speed, mixed-signal PCAs and instrument assemblies, covering analog/digital interfaces, ADCs/DACs, FPGAs, and power distribution systems. Develop test requirements and coverage matrices with hardware and software stakeholders, including optimization of test coverage vs test time. Analyze test data to identify failure root causes and trends, implement corrective actions, and drive design-for-testability (DFT) enhancements. Drive continuous test improvement to improve test accuracy, improve final product reliability, and adapt to new measurement requirements.
US, WA, Seattle
This role will contribute to developing the Economics and Science products and services in the Fee domain, with specialization in supply chain systems and fees. Through the lens of economics, you will develop causal links for how Amazon, Sellers and Customers interact. You will be a key and senior scientist, advising Amazon leaders how to price our services. You will work on developing frameworks and scalable, repeatable models supporting optimal pricing and policy in the two-sided marketplace that is central to Amazon's business. The pricing for Amazon services is complex. You will partner with science and technology teams across Amazon including Advertising, Supply Chain, Operations, Prime, Consumer Pricing, and Finance. We are looking for an experienced Economist to improve our understanding of seller Economics, enhance our ability to estimate the causal impact of fees, and work with partner teams to design pricing policy changes. In this role, you will provide guidance to scientists to develop econometric models to influence our fee pricing worldwide. You will lead the development of causal models to help isolate the impact of fee and policy changes from other business actions, using experiments when possible, or observational data when not. Key job responsibilities The ideal candidate will have extensive Economics knowledge, demonstrated strength in practical and policy relevant structural econometrics, strong collaboration skills, proven ability to lead highly ambiguous and large projects, and a drive to deliver results. They will work closely with Economists, Data / Applied Scientists, Strategy Analysts, Data Engineers, and Product leads to integrate economic insights into policy and systems production. Familiarity with systems and services that constitute seller supply chains is a plus but not required. About the team The Stores Economics and Sciences team is a central science team that supports Amazon's Retail and Supply Chain leadership. We tackle some of Amazon's most challenging economics and machine learning problems, where our mandate is to impact the business on massive scale.
US, WA, Seattle
WW Amazon Stores Finance Science (ASFS) works to leverage science and economics to drive improved financial results, foster data backed decisions, and embed science within Finance. ASFS is focused on developing products that empower controllership, improve business decisions and financial planning by understanding financial drivers, and innovate science capabilities for efficiency and scale. We are looking for a data scientist to lead high visibility initiatives for forecasting Amazon Stores' financials. You will develop new science-based forecasting methodologies and build scalable models to improve financial decision making and planning for senior leadership up to VP and SVP level. You will build new ML and statistical models from the ground up that aim to transform financial planning for Amazon Stores. We prize creative problem solvers with the ability to draw on an expansive methodological toolkit to transform financial decision-making with science. The ideal candidate combines data-science acumen with strong business judgment. You have versatile modeling skills and are comfortable owning and extracting insights from data. You are excited to learn from and alongside seasoned scientists, engineers, and business leaders. You are an excellent communicator and effectively translate technical findings into business action. Key job responsibilities Demonstrating thorough technical knowledge, effective exploratory data analysis, and model building using industry standard ML models Working with technical and non-technical stakeholders across every step of science project life cycle Collaborating with finance, product, data engineering, and software engineering teams to create production implementations for large-scale ML models Innovating by adapting new modeling techniques and procedures Presenting research results to our internal research community
IN, KA, Bengaluru
RBS (Retail Business Services) Tech team works towards enhancing the customer experience (CX) and their trust in product data by providing technologies to find and fix Amazon CX defects at scale. Our platforms help in improving the CX in all phases of customer journey, including selection, discoverability & fulfilment, buying experience and post-buying experience (product quality and customer returns). The team also develops GenAI platforms for automation of Amazon Stores Operations. As a Sciences team in RBS Tech, we focus on foundational ML research and develop scalable state-of-the-art ML solutions to solve the problems covering customer experience (CX) and Selling partner experience (SPX). We work to solve problems related to multi-modal understanding (text and images), task automation through multi-modal LLM Agents, supervised and unsupervised techniques, multi-task learning, multi-label classification, aspect and topic extraction for Customer Anecdote Mining, image and text similarity and retrieval using NLP and Computer Vision for product groupings and identifying duplicate listings in product search results. Key job responsibilities As an Research Scientist, you will be responsible to design and deploy scalable GenAI, NLP and Computer Vision solutions that will impact the content visible to millions of customer and solve key customer experience issues. You will develop novel LLM, deep learning and statistical techniques for task automation, text processing, image processing, pattern recognition, and anomaly detection problems. You will define the research and experiments strategy with an iterative execution approach to develop AI/ML models and progressively improve the results over time. You will partner with business and engineering teams to identify and solve large and significantly complex problems that require scientific innovation. You will help the team leverage your expertise, by coaching and mentoring. You will contribute to the professional development of colleagues, improving their technical knowledge and the engineering practices. You will independently as well as guide team to file for patents and/or publish research work where opportunities arise. The RBS org deals with problems that are directly related to the selling partners and end customers and the ML team drives resolution to organization level problems. Therefore, the Research Scientist role will impact the large product strategy, identifies new business opportunities and provides strategic direction which is very exciting.
US, WA, Seattle
As part of the AWS Applied AI Solutions Core Services organization, we're advancing the frontier of geospatial intelligence and AI-powered spatial reasoning. Our vision is to be the trusted foundation for transforming every business with Amazon AI teammates. Our mission is to deliver turnkey, enterprise-grade foundational AI capabilities that create delightful AI powered solutions. We're building sophisticated AI systems that enable intelligent agents to understand and operate effectively in the physical world through advanced geospatial optimization. Key job responsibilities - Develop geospatial optimization models that generalize across diverse customer use cases in logistics, transportation, and spatial planning - Scope optimization projects with multiple customers in mind, abstracting away complex science problems to create scalable solutions - Discover, evaluate, and adapt existing optimization models and geospatial tools for customer deployment - Develop semantic enrichment methods to integrate heterogeneous data sources including open geospatial data, multimodal sensor data, images, videos, satellite imagery, and documents - Research novel approaches combining AI agents with geospatial optimization to solve complex spatial problems - Collaborate with engineering teams to integrate science components into production systems - Conduct rigorous experimentation and establish evaluation frameworks to measure solution performance A day in the life A day in the life As an Applied Scientist, you'll develop optimization algorithms and AI-powered geospatial solutions while maintaining a clear path to customer impact. You'll investigate novel approaches to spatial optimization, develop methods for semantic data enrichment, and validate ideas through rigorous experimentation with real customer data. You'll collaborate with other scientists and engineers to transform research insights into scalable solutions, work directly with enterprise customers to understand requirements, and help shape the future direction. Leveraging and advancing generative AI technology will be a big part of your charter. About the team Our Applied AI Solutions Core Services Science team is tackling fundamental challenges in geospatial optimization and AI-powered spatial reasoning. We're investigating novel approaches to how AI systems can solve complex logistics and transportation problems, reason about spatial relationships, and integrate diverse data sources to create enterprise-grade geospatial intelligence. Working at the intersection of optimization, large language models, and geospatial data science, we're developing practical techniques that advance the state-of-the-art in geospatial AI.
US, WA, Bellevue
We are looking for detail-oriented, organized, and responsible individuals who are eager to learn how to apply their causal inference and/or structural econometrics skillsets to solve real world problems. The intern will work in the area of Economics Intelligence in Amazon Returns and Recommerce Technology and Innovation and develop new, data-driven solutions to support the most critical components of this rapidly scaling team. Our PhD Economist Internship Program offers hands-on experience in applied economics, supported by mentorship, structured feedback, and professional development. Interns work on real business and research problems, building skills that prepare them for full-time economist roles at Amazon and beyond. You will learn how to build data sets and perform applied econometric analysis collaborating with economists, scientists, and product managers. These skills will translate well into writing applied chapters in your dissertation and provide you with work experience that may help you with placement. These are full-time positions at 40 hours per week, with compensation being awarded on an hourly basis. About the team The WWRR Economics Intelligence (RREI) team brings together Economists, Data Scientists, and Business Intelligence Engineers experts to delivers economic solutions focused on forecasting, causality, attribution, customer behavior for returns, recommerce, and sustainability domains.
US, WA, Bellevue
We are looking for detail-oriented, organized, and responsible individuals who are eager to learn how to apply their causal inference and/or structural econometrics skillsets to solve real world problems. The intern will work in the area of Economics Intelligence in Amazon Returns and Recommerce Technology and Innovation and develop new, data-driven solutions to support the most critical components of this rapidly scaling team. Our PhD Economist Internship Program offers hands-on experience in applied economics, supported by mentorship, structured feedback, and professional development. Interns work on real business and research problems, building skills that prepare them for full-time economist roles at Amazon and beyond. You will learn how to build data sets and perform applied econometric analysis collaborating with economists, scientists, and product managers. These skills will translate well into writing applied chapters in your dissertation and provide you with work experience that may help you with placement. These are full-time positions at 40 hours per week, with compensation being awarded on an hourly basis. About the team The WWRR Economics Intelligence (RREI) team brings together Economists, Data Scientists, and Business Intelligence Engineers experts to delivers economic solutions focused on forecasting, causality, attribution, customer behavior for returns, recommerce, and sustainability domains.
US, CA, San Francisco
AWS is one of Amazon’s largest and fastest growing businesses, serving millions of customers in more than 190 countries. We use cloud computing to reshape the way global enterprises use information technology. We are looking for entrepreneurial, analytical, creative, flexible leaders to help us redefine the information technology industry. If you want to join a fast-paced, innovative team that is making history, this is the place for you. AWS Central Economics & Science (ACES) drives best practices for objectively applying economics and science in decision making across AWS. The team collaborates with AWS science and business teams to identify, frame, and analyze complex and ambiguous problems of the highest priority. Through data-driven insights and modeling, ACES supports strategic decision-making across the AWS global organization, including sales operations and business performance optimization. The ACES Sales Channels team is hiring an Applied Scientist (Senior or below) to advance our mission of providing rigorous, causal-inference-driven recommendations for AWS sales optimization. This role will focus on building ML systems with a causal modeling foundation, designing seller incentive mechanisms, and developing intervention strategies across the entire sales motion. Key job responsibilities • Causal ML System Development: Build and deploy machine learning models that emphasize causal inference, ensuring recommendations are grounded in valid interventions • Incentive Design: Define and model incentives that drive desirable behaviors across AWS sales channels, partner programs, and reseller ecosystems • Stakeholder Collaboration: Work with business stakeholders to understand requirements, validate approaches, and ensure practical applicability of scientific solutions • Scientific Rigor: Promote findings at internal conferences and contribute to the team's reputation for methodological excellence A day in the life The ACES Sales Channels team works on understanding and optimizing AWS's sales channels, both direct (generalist and specialist sellers) and indirect (partners and Marketplace). Our work falls into three core areas: developing rigorous causal measurement and modeling frameworks using cutting-edge economics and statistical methods; designing programs and incentives to improve customer and business outcomes; and building ML-based recommendation systems for sellers, partners, and other AWS stakeholders. About the team Why AWS? Amazon Web Services (AWS) is the world’s most comprehensive and broadly adopted cloud platform. We pioneered cloud computing and never stopped innovating — that’s why customers from the most successful startups to Global 500 companies trust our robust suite of products and services to power their businesses. Inclusive Team Culture Here at AWS, it’s in our nature to learn and be curious. Our employee-led affinity groups foster a culture of inclusion that empower us to be proud of our differences. Ongoing events and learning experiences, including our Conversations on Race and Ethnicity (CORE) and AmazeCon conferences, inspire us to never stop embracing our uniqueness. Mentorship & 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, mentorship 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 we strive for flexibility as part of our working culture. When we feel supported in the workplace and at home, there’s nothing we can’t achieve in the cloud. Hybrid Work We value innovation and recognize this sometimes requires uninterrupted time to focus on a build. We also value in-person collaboration and time spent face-to-face. Our team affords employees options to work in the office every day or in a flexible, hybrid work model near one of our U.S. Amazon offices.
US, WA, Bellevue
The Central Learning Solutions (CLS) - Science team builds state-of-the-art Artificial Intelligence (AI) solutions for enhancing leadership and associate development within the organization. We develop technology and mechanisms for building personalized learning courses based on the profiles of different learners and asses the post-training performance curves for different learner profiles. As a Data Scientist on the team, you will be driving the data science/ML roadmap for the CLS t Science team. You will leverage your knowledge in statistics and econometrics, estimate the causal impact of training interventions, recommend the right interventions for a given learner profile, and measure the post-launch success of these interventions through A/B weblabs. These insights will help in dynamically changing the training content of Learning & Development courses and unlock opportunities to improve both training effectiveness and learner experience. You will collaborate effectively with internal stakeholders and cross-functional teams for solving business problems, create operational efficiencies, and deliver successfully against high organizational standards. Key job responsibilities - Establish scalable, efficient, automated processes for large scale data analyses, model development, model validation and implementation. - Use advanced causal inference methodologies to estimate the learning curves for different learner profiles and the effectiveness of training content. - Perform statistical analysis and statistical tests including hypothesis testing and A/B testing. - Implement new statistical, machine learning, or other mathematical methodologies to solve specific business problems. - Present deep dives and analysis to both technical and non-technical stakeholders, ensure clarity, and influence the strategy of business partners. About the team We serve North America L&D orgs as the strategic thought leader, looking beyond where other teams are focused to drive transformative solutions that leverage technology and processes to improve learning outcomes and drive down the cost to serve.
US, WA, Bellevue
The Principal Applied Scientist will own the science mission for building next-generation proactive and autonomous agentic experiences across Alexa AI's Personalization, Autonomy and Proactive Intelligence organization. You will technically lead a team of applied scientists to harness state-of-the-art technologies in machine learning, natural language processing, LLM training and application, and agentic AI systems to advance the scientific frontiers of autonomous intelligence and proactive user assistance. The right candidate will be an inventor at heart, provide deep scientific leadership, establish compelling technical direction and vision, and drive ambitious research initiatives that push the boundaries of what's possible with AI agents. You will need to be adept at identifying promising research directions in agentic AI, developing novel autonomous agent solutions, and translating advanced AI research into production-ready agentic systems. You will need to be adept at influencing and collaborating with partner teams, launching AI-powered autonomous agents into production, and building team mechanisms that will foster innovation and execution in the rapidly evolving field of agentic AI. This role represents a unique opportunity to tackle fundamental challenges in how Alexa proactively understands user needs, autonomously takes actions on behalf of users, and delivers intelligent assistance through state-of-the-art agentic AI technologies. As a science leader in Alexa AI, you will shape the technical strategy for making Alexa a truly proactive and autonomous agent that anticipates user needs, takes intelligent actions, and provides seamless assistance without explicit prompting. Your team will be at the forefront of solving complex problems in agentic reasoning, multi-step task planning, autonomous decision-making, proactive intelligence, and context-aware action execution that will fundamentally transform how users interact with Alexa as an intelligent agent. The successful candidate will bring deep technical expertise in machine learning, natural language processing, and agentic AI systems, along with the leadership ability to guide talented scientists in pursuing ambitious research that advances the state of the art in autonomous agents, proactive intelligence, and AI-driven personalization. Experience with multi-agent systems, reinforcement learning, goal-oriented dialogue systems, and production-scale agentic architectures is highly valued. You will lead the development of breakthrough capabilities that enable Alexa to: 1) proactively anticipate user needs through advanced predictive modeling and contextual understanding; 2) autonomously execute complex multi-step tasks with minimal user intervention; 3) reason and plan intelligently across diverse user goals and environmental contexts; 4) learn and adapt continuously from user interactions to improve agentic behaviors; 5) coordinate actions seamlessly across multiple domains and services as a unified intelligent agent. This is a unique opportunity to define the future of conversational AI agents and build technology that will impact hundreds of millions of customers worldwide. Key job responsibilities Technical Leadership - Lead complex research and development projects - Partner closely with the T&C Product and Engineering leaders on the technical strategy and roadmap - Evaluate emerging technologies and methodologies - Make high-level architectural decisions Technical leadership and mentoring: - Mentor and develop technical talent - Set team project goals and metrics - Help with resource allocation and project prioritization from technical side Research & Development - Drive innovation in applied science areas - Translate research into practical business solutions - Author technical papers and patents - Collaborate with academic and industry partners About the team PAPI (Personalization Autonomy and Proactive Intelligence) aims to accelerate personalized and intuitive experiences across Amazon's customer touchpoints through automated, scalable, self-serve AI systems. We leverage customer, device, and ambient signals to deliver conversational, visual, and proactive experiences that delight customers, increase engagement, reduce defects, and enable natural interactions across Amazon touch points including Alexa, FireTV, and Mobile etc. Our systems offer personalized suggestions, comprehend customer inputs, learn from interactions, and propose appropriate actions to serve millions of customers globally.