Enhancing self-supervised representation learning through lightweight learnable data augmentation

Abstract

Data augmentation plays a central role in self-supervised learning, as the quality and diversity of augmented views strongly influence the learned representations. However, most existing self-supervised methods rely on fixed stochastic augmentation pipelines, while more adaptive alternatives often require expensive policy search, adversarial training, or additional optimization procedures. In this paper, we propose a lightweight learnable augmentation framework based on Extreme Learning Machines (ELM) for self-supervised visual representation learning. The proposed module predicts image-dependent transformation parameters and applies them through a differentiable augmentation operator, enabling joint optimization with the representation model while introducing minimal additional computational overhead. The framework is integrated into three representative self-supervised learning methods: SimCLR, BYOL, and SimSiam. Extensive experiments on CIFAR-10, CIFAR-100, and Tiny ImageNet show that the proposed method consistently improves linear evaluation performance relative to reproduced baselines across most settings. In particular, the method yields notable gains on CIFAR datasets and remains effective on the more challenging Tiny ImageNet benchmark. A per-class difficulty analysis further shows that the proposed augmentation strategy substantially improves performance on hard classes, indicating stronger robustness to challenging categories while maintaining competitive overall performance. In general, the results demonstrate that lightweight learnable augmentation can effectively enhance self-supervised representation learning across different frameworks and datasets.