Self-Supervised Pre-Training with Equilibrium Constraints

2025-08-27

AI Summary

This paper addresses the weak downstream adaptability of self-supervised pre-trained models under heterogeneous data settings (e.g., multi-domain or multilingual scenarios). To tackle this, we propose a bilevel optimization framework with equilibrium constraints. At the lower level, we model local-optimal learning on each source domain as an inner optimization problem; at the upper level, model parameters are updated under equilibrium constraints that enforce balanced representation learning across domains. Efficient solving is achieved via first-order approximation and K-step gradient descent. Theoretically, we establish a formal connection to Model-Agnostic Meta-Learning (MAML), revealing the meta-optimization nature of our approach. Experiments on multi-domain and multilingual benchmarks demonstrate substantial improvements in generalization and performance under downstream supervised fine-tuning—particularly enhancing robustness and adaptation capability for low-resource domains.

Technology Category

• Machine Learning: Transfer, Domain Adaptation, Multi-Task Learning
• Search and Optimization: Learning to Search
• Natural Language Processing: Learning & Optimization for NLP

Application Category

• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Web Mining and Content Analysis: Large pretrained models with web data

Abstract

Self-supervised pre-training using unlabeled data is widely used in machine learning. In this paper, we propose a new self-supervised pre-training approach to dealing with heterogeneous data. Instead of mixing all the data and minimizing the averaged global loss in the conventional way, we impose additional equilibrium constraints to ensure that the models optimizes each source of heterogeneous data to its local optima after $K$-step gradient descent initialized from the model. We formulate this as a bilevel optimization problem, and use the first-order approximation method to solve the problem. We discuss its connection to model-agnostic meta learning (MAML). Experiments are carried out on self-supervised pre-training using multi-domain and multilingual datasets, demonstrating that the proposed approach can significantly improve the adaptivity of the self-supervised pre-trained model for the downstream supervised fine-tuning tasks.

Problem

• Proposes self-supervised pre-training for heterogeneous data
• Uses equilibrium constraints to optimize each data source locally
• Formulates as bilevel optimization to enhance model adaptivity

Innovation

• Equilibrium constraints for heterogeneous data optimization
• Bilevel optimization with first-order approximation
• Connection to model-agnostic meta learning (MAML)