An Information-Geometric Distance on the Space of Tasks
Yansong Gao, Pratik Chaudhari
Abstract
This paper prescribes a distance between learning tasks modeled as joint distributions on data and labels. Using tools in information geometry, the distance is defined to be the length of the shortest weight trajectory on a Riemannian manifold as a classifier is fitted on an interpolated task. The interpolated task evolves from the source to the target task using an optimal transport formulation. This distance, which we call the "coupled transfer distance" can be compared across different classifier architectures. We develop an algorithm to compute the distance which iteratively transports the marginal on the data of the source task to that of the target task while updating the weights of the classifier to track this evolving data distribution. We develop theory to show that our distance captures the intuitive idea that a good transfer trajectory is the one that keeps the generalization gap small during transfer, in particular at the end on the target task. We perform thorough empirical validation and analysis across diverse image classification datasets to show that the coupled transfer distance correlates strongly with the difficulty of fine-tuning.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Cited by top-tier papers5
- A Picture of the Space of Typical Learnable TasksRahul Ramesh, Jialin Mao, Itay Griniasty, Rubing Yang et al.ICML 2023 · 7 citations
- Time-Varying Propensity Score to Bridge the Gap between the Past and PresentRasool Fakoor, Jonas Mueller, Zachary Chase Lipton, Pratik Chaudhari et al.ICLR 2024 · 4 citations
- Interpolation for Robust Learning: Data Augmentation on Wasserstein GeodesicsJiacheng Zhu, Jielin Qiu, Aritra Guha, Zhuolin Yang et al.ICML 2023 · 4 citations
- The Geometry of Updates: Fisher Alignment at Vocabulary ScaleJohn SweeneyICML 2026 · 1 citation
- Lightspeed Geometric Dataset Distance via Sliced Optimal TransportKhai Nguyen, Hai Nguyen, Tuan Pham, Nhat HoICML 2025
Builds on5
- A Baseline for Few-Shot Image ClassificationGuneet Singh Dhillon, Pratik Chaudhari, Avinash Ravichandran, Stefano SoattoICLR 2020 · 640 citations
- Task2Vec: Task Embedding for Meta-LearningAlessandro Achille, Michael Lam, Rahul Tewari, Avinash Ravichandran et al.ICCV 2019 · 359 citations
- Geometric Dataset Distances via Optimal TransportDavid Alvarez-Melis, Nicolò FusiNeurIPS 2020 · 267 citations
- Rethinking the Hyperparameters for Fine-tuningHao Li, Pratik Chaudhari, Hao Yang, Michael Lam et al.ICLR 2020 · 142 citations
- A Free-Energy Principle for Representation LearningYansong Gao, Pratik ChaudhariICML 2020 · 11 citations
Related papers
- Curriculum Reinforcement Learning using Optimal Transport via Gradual Domain AdaptationPeide Huang, Mengdi Xu, Jiacheng Zhu, Laixi Shi et al.NeurIPS 2022 · 44 citations
- Geometrically Aligned Transfer Encoder for Inductive Transfer in Regression TasksSung Moon Ko, Sumin Lee, Dae-Woong Jeong, Woohyung Lim et al.ICLR 2024 · 6 citations
- Learning Structured Representations by Embedding Class Hierarchy with Fast Optimal TransportSiqi Zeng, Sixian Du, Makoto Yamada, Han ZhaoICLR 2025
- Flowing Datasets with Wasserstein over Wasserstein Gradient FlowsClément Bonet, Christophe Vauthier, Anna KorbaICML 2025
- Riemannian Metric Learning via Optimal TransportChristopher Scarvelis, Justin SolomonICLR 2023 · 2 citations
