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

IN, HR, Gurugram
Building large-scale forecasting and optimization systems that power Amazon’s global transportation network and directly impact customer experience and cost. Key job responsibilities 1. Guide model and system design across a range of techniques, including tree-based models, deep learning (LSTMs, transformers), LLMs, and reinforcement learning. 2. Ensure models are production-ready, scalable, and robust through close partnership with stakeholders. 3. Partner with Product, Operations, and Engineering leaders to enable proactive decision-making and corrective actions. 4 Own end-to-end business metrics, directly influencing customer experience, cost optimization, and network reliability. 5. Help contribute to the broader ML community through publications, conference submissions, and internal knowledge sharing.
US, WA, Seattle
Estimating the demand response of a pricing decision is genuinely hard. The causal effects are delayed, noisy, and confounded by factors that standard experiment analysis wasn't designed to handle. Most pricing teams default to heuristics not because they don't care about customer responses, but because measuring them rigorously is an unsolved problem. P2OS is building the science to solve it. We're hiring an Economist to own that work — defining how we estimate digital demand response in a pricing context, building the identification strategies that make those estimates credible, and translating outputs into something pricing teams can use to make better decisions. The role sits at the intersection of econometric methodology and production-quality analysis, and requires someone who can operate independently in both. As science lead, you'll own the digital pricing methodology domain, and be the internal authority on causal inference for pricing across P2OS and partner teams. Key job responsibilities * Own the end-to-end digital pricing methodology for pricing — identification strategy, modeling choices, validation approach, and business use cases — and drive adoption across pricing contexts * Deliver high-stakes analyses connecting digital pricing estimates to a concrete pricing decision and strategy change at VP+ level * Apply advanced causal methods to live pricing problems; document approaches so the team can build on and extend them. * Provide causal inference guidance on pricing experiment questions as they arise — being the methodology resource when experiments generate relevant questions * Serve as cross-team economic advisor to Digital Finance, Customer Behavior, and Demand Science on assumptions and causal identification * Actively mentor junior scientists, earn trust of cross-functional tech and product partners. A day in the life In a typical day, you'll move between methodology work and stakeholder-facing analysis. - On the science side, that means reviewing identification assumptions with the Causal AS, validating estimation choices for the LTV framework, and documenting methodology decisions in ways that non-economists can act on. - On the applied side, you'll be in rooms with Finance, Pricing PMs, and other science teams: aligning on LTV definitions, resolving disagreements between competing metrics, and translating causal findings into recommendations that land in strategy reviews. - As tech lead, you need to work to develop the economists and scientists on your scrum: structured reviews, identification strategy feedback, and raising the quality of analyses before they reach stakeholders. The mix shifts, but the through-line is to progress the LTV methodology from open questions to shipped frameworks, and making sure the team's causal work is rigorous enough to hold up when it counts. About the team P2Optimization Science (P2OS) is responsible for the ML models and analytical frameworks that drive pricing decisions at scale. The team spans demand lift modeling, pricing error detection, customer lifetime value, and experimentation. Our small team of specialized applied scientists and economists works closely alongside engineers, and pricing product managers.
US, WA, Seattle
We’re working to improve shopping on Amazon using the conversational capabilities of large language models, and are searching for pioneers who are passionate about technology, innovation, and customer experience, and are ready to make a lasting impact on the industry. You'll be working with talented scientists, engineers, and technical program managers (TPM) to innovate on behalf of our customers. If you're fired up about being part of a dynamic, driven team, then this is your moment to join us on this exciting journey!
US, MA, Boston
Are you interested in how to build AI reasoning systems that give provably correct answers? Are you excited by science at the interface of classical AI reasoning and Large Language Models (LLMs)? Would you like to apply your technology to serve operations customers better? Amazon Robotics is looking for a talented Applied Scientist in Neurosymbolic AI. You will innovate on combining language models (LMs) with classical AI reasoning. You will work with a team of scientists and engineers to achieve this. You will publish your results in papers at leading venues in AI. You will be part of a larger team and have the opportunity to work on problems such as: using LMs to generate plans, using AI reasoning to verify plan correctness, learning efficient reasoning strategies, self-improving models. You will work on basic science and on business problems in robotics, automation and fulfillment across our operations. Key job responsibilities In this role you will: • Work closely with other scientists and engineers, and be part of Amazon’s diverse global science community. • Publish your research in top-tier academic venues and hone your presentation skills. • Be inspired by challenges and opportunities to invent new techniques in your area(s) of expertise. A day in the life You'll meet regularly with your technical lead and your team on your ideas, get guidance and feedback, work together on architectures and algorithms, author papers, build AI systems, all with the aim of delivering results for your operations customers. You'll work closely with other scientists to review your plans and results. You'll meet with engineers to implement your ideas at scale. About the team The Veritas team is a science team working at the boundary between language models and classical AI reasoning. We work across on customer problems in fulfillment, automation and robotics. We focus on high quality research science informed by practical problems.
US, WA, Seattle
Economists in this role partner with business stakeholders to distill complex problems into testable economic questions and generate actionable insights. They collaborate with engineers and scientists to estimate models on large-scale data, design pilots, measure impact, and scale successful prototypes into improved policies and programs. They leverage AI tools to scale economic study for broader business impact. They communicate findings to business leaders, incorporate feedback, and deliver customer-centric solutions at scale.
CA, BC, Vancouver
The Alexa Daily Essentials team delivers experiences critical to how customers interact with Alexa as part of daily life. Alexa users engage with our products across experiences connected to Timers, Alarms, Calendars, Food, and News. Our experiences include critical time saving techniques, ad-supported news audio and video, and in-depth kitchen guidance aimed at serving the needs of the family from sunset to sundown. As a Data Scientist on our team, you'll work with complex data, develop statistical methodologies, and provide critical product insights that shape how we build and optimize our solutions. You will work closely with your Analytics and Applied Science teammates. You will build frameworks and mechanisms to scale data solutions across our organization. If you are passionate about redefining how AI can improves everyone's daily life, we’d love to hear from you. Key job responsibilities Problem-Solving - Analyze complex data to identify patterns, inform product decisions, and understand root causes of anomalies. - Develop analysis and modeling approaches to drive product and engineering actions to identify patterns, insights, and understand root causes of anomalies. Your solutions directly improve the customer experience. - Independently work with product partners to identify problems and opportunities. Apply a range of data science techniques and tools to solve these problems. Use data driven insights to inform product development. Work with cross-disciplinary teams to mechanize your solution into scalable and automated frameworks. Data Infrastructure - Build data pipelines, and identify novel data sources to leverage in analytical work - both from within Alexa and from cross Amazon - Acquire data by building the necessary SQL / ETL queries Communication - Excel at communicating complex ideas to technical and non-technical audiences. - Build relationships with stakeholders and counterparts. Work with stakeholders to translate causal insights into actionable recommendations - Force multiply the work of the team with data visualizations, presentations, and/or dashboards to drive awareness and adoption of data assets and product insights - Collaborate with cross-functional teams. Mentor teammates to foster a culture of continuous learning and development
US, NY, New York
The Ads Measurement Science team in the Measurement, Ad Tech, and Data Science (MADS) team of Amazon Ads serves a centralized role developing solutions for a multitude of performance measurement products. We create solutions which measure the comprehensive impact of advertiser's ad spend, including sales impacts both online and offline and across timescales, and provide actionable insights that enable our advertisers to optimize their media portfolios. We also own the science solutions for AI tools that unlock new insights and automate high-effort customer workflows, such as custom query and report generation based on natural language user requests. We leverage a host of scientific technologies to accomplish this mission, including Generative AI, classical ML, Causal Inference, Natural Language Processing, and Computer Vision. As a Senior Applied Scientist on the team, you will be at the forefront of innovation, developing measurement solutions end-to-end from inception to production. You will set the technical vision and innovate on behalf of our customers. You will propose, design, analyze, and productionize models to provide novel measurement insights to our customers. You will partner with engineering to deploy these solutions into production. You will work with key stakeholders from various business teams to enable advertisers to act upon those metrics. Key job responsibilities * Lead the development of ad measurement models and solutions that address the full spectrum of an advertiser's investment, focusing on scalable and efficient methodologies. * Collaborate closely with cross-functional teams including engineering, product management, and business teams to define and implement measurement solutions. * Use state-of-the-art scientific technologies including Generative AI, Classical Machine Learning, Causal Inference, Natural Language Processing, and Computer Vision to develop state of the art models that measure the impact of ad spend across multiple platforms and timescales. * Drive experimentation and the continuous improvement of ML models through iterative development, testing, and optimization. * Translate complex scientific challenges into clear and impactful solutions for business stakeholders. * Mentor and guide junior scientists, fostering a collaborative and high-performing team culture. * Foster collaborations between scientists to move faster, with broader impact. * Regularly engage with the broader scientific community with presentations, publications, and patents. A day in the life You will solve real-world problems by getting and analyzing large amounts of data, generate business insights and opportunities, design simulations and experiments, and develop statistical and ML models. The team is driven by business needs, which requires collaboration with other Scientists, Engineers, and Product Managers across the advertising organization. You will prepare written and verbal presentations to share insights to audiences of varying levels of technical sophistication. Team video https://advertising.amazon.com/help/G4LNN5YWHP6SM9TJ About the team We are a team of scientists across Applied, Research, Data Science and Economist disciplines. You will work with colleagues with deep expertise in ML, NLP, CV, Gen AI, and Causal Inference with a diverse range of backgrounds. We partner closely with top-notch engineers, product managers, sales leaders, and other scientists with expertise in the ads industry and on building scalable modeling and software solutions.
US, NY, New York
The Ads Measurement Science team in the Measurement, Ad Tech, and Data Science (MADS) team of Amazon Ads serves a centralized role developing solutions for a multitude of performance measurement products. We create solutions which measure the comprehensive impact of advertiser's ad spend, including sales impacts both online and offline and across timescales, and provide actionable insights that enable our advertisers to optimize their media portfolios. We also own the science solutions for AI tools that unlock new insights and automate high-effort customer workflows, such as custom query and report generation based on natural language user requests. We leverage a host of scientific technologies to accomplish this mission, including Generative AI, classical ML, Causal Inference, Natural Language Processing, and Computer Vision. As an Applied Scientist on the team, you will lead measurement solutions end-to-end from inception to production. You will propose, design, analyze, and productionize models to provide novel measurement insights to our customers. Key job responsibilities Leverage deep expertise in one or more scientific disciplines to invent solutions to ambiguous ads measurement problems Disambiguate problems to propose clear evaluation frameworks and success criteria Work autonomously and write high quality technical documents Implement a significant portion of critical-path code, and partner with engineers to directly carry solutions into production Partner closely with other scientists to deliver large, multi-faceted technical projects Share and publish works with the broader scientific community through meetings and conferences Communicate clearly to both technical and non-technical audiences Contribute new ideas that shape the direction of the team's work Mentor more junior scientists and participate in the hiring process About the team We are a team of scientists across Applied, Research, Data Science and Economist disciplines. You will work with colleagues with deep expertise in ML, NLP, CV, Gen AI, and Causal Inference with a diverse range of backgrounds. We partner closely with top-notch engineers, product managers, sales leaders, and other scientists with expertise in the ads industry and on building scalable modeling and software solutions.
US, WA, Seattle
Are you interested in shaping the future of Advertising and B2B Sales? We are a growing team with an exciting AI-first charter and need your passion, innovative thinking, and creativity to help take our products to new heights. 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 businesses driving long term growth. We break fresh ground in product and technical innovations every day! Within the Advertising Sales organization, we are building a central AI/ML team and are seeking top Applied Science talent to help us build new, science-backed services that drive success for our customers. Our goal is to transform the way account teams operate by creating actionable insights and recommendations they can share with their advertising accounts, and ingesting Generative AI throughout their end-to-end workflows to improve their work efficiency. As an Applied Scientist on the team, you will bring deep expertise in modeling dynamic systems using statistical methods and deep learning, and in optimizing those systems using reinforcement learning and operations research. You have the scientific and technical skills to build and refine models that can be implemented in production, and you leverage natural language processing and generative AI to enhance their explainability. You will chart new courses with our ad sales support technologies, and you have the communication skills necessary to explain complex technical approaches to a variety of stakeholders and customers. You will be part of a team of fellow scientists and engineers taking iterative approaches to tackle big, long-term problems. You are fluently able to leverage the latest generative AI systems and services to accelerate and improve your work while maintaining high quality in your outputs. Key job responsibilities Scientific Modeling - Conceptualize and lead state-of-the-art research on new Machine Learning and Generative Artificial Intelligence solutions to optimize all aspects of the Ad Sales business - Lead the technical approach for the design and implementation of successful models and algorithms in support of expert cross-functional teams delivering on demanding projects - Run regular A/B experiments, gather data, and perform statistical analysis - Improve the scalability, efficiency and automation of large-scale data analytics, model training, deployment and serving - Publish scientific findings in reports and papers that can be shared internally and externally Product Development Support - Partner with software engineering and product management teams to support product and service development, define success metrics and measurement approaches, and help drive adoption of innovative new features for our services. - Lead requirements gathering sessions with product teams and business stakeholders - Maintain scientific documentation and knowledge for product initiatives Collaboration & Communication - Work closely with software engineers to deliver end-to-end solutions into production - Translate complex scientific findings into actionable business recommendations for stakeholders and senior management - Provide clear, compelling reports and presentations on a regular basis with respect to your models and services - Communicate with internal teams to showcase results and identify best practices. About the team Sales AI is a central science and engineering organization within Amazon Advertising Sales that powers selling motions and account team workflows via state-of-the-art of AI/ML services. Sales AI is investing in a range of sales intelligence models, including the development of advertiser insights, recommendations and Generative AI-powered applications throughout account team workflows.
US, NY, New York
About Sponsored Products and Brands The Sponsored Products and Brands team at Amazon Ads is re-imagining the advertising landscape through generative AI technologies, revolutionizing how millions of customers discover products and engage with brands across Amazon.com and beyond. We are at the forefront of re-inventing advertising experiences, bridging human creativity with artificial intelligence to transform every aspect of the advertising lifecycle from ad creation and optimization to performance analysis and customer insights. We are a passionate group of innovators dedicated to developing responsible and intelligent AI technologies that balance the needs of advertisers, enhance the shopping experience, and strengthen the marketplace. If you're energized by solving complex challenges and pushing the boundaries of what's possible with AI, join us in shaping the future of advertising. Key job responsibilities As a Machine Learning Applied Scientist, you will: * Conduct deep data analysis to derive insights to the business, and identify gaps and new opportunities * Develop scalable and effective machine-learning models and optimization strategies to solve business problems * Run regular A/B experiments, gather data, and perform statistical analysis * Work closely with software engineers to deliver end-to-end solutions into production * Improve the scalability, efficiency and automation of large-scale data analytics, model training, deployment and serving * Conduct research on new machine-learning modeling and Generative AI solutions to optimize all aspects of Sponsored Products and Brands business About the team The Ad Response Prediction team within Sponsored Products and Brands (SPB) drives personalized shopping experiences for SPB Ads across placements, pages, and devices worldwide. We achieve this through ML and GenAI solutions that include customized shopper response prediction and session-level understanding to optimize every stage of the ad-serving process, from sourcing and bidding to widget discovery and auctions. Our responsibilities include advancing response prediction through model and feature innovations and extending prediction beyond the auction stage to areas such as targeting, sourcing, and bidding.