Re-basin via implicit Sinkhorn differentiation

TL;DR

This paper introduces a differentiable Sinkhorn re-basin network that improves permutation finding in neural networks, enabling better incremental learning and mode connectivity, with competitive results on benchmark datasets.

Contribution

It presents a novel differentiable re-basin method using Sinkhorn optimization, facilitating integration into gradient-based training and enhancing continual learning capabilities.

Findings

Outperforms existing methods in permutation optimization tasks.

Enables incremental learning through a new cost function.

Achieves competitive results on benchmark datasets.

Abstract

The recent emergence of new algorithms for permuting models into functionally equivalent regions of the solution space has shed some light on the complexity of error surfaces, and some promising properties like mode connectivity. However, finding the right permutation is challenging, and current optimization techniques are not differentiable, which makes it difficult to integrate into a gradient-based optimization, and often leads to sub-optimal solutions. In this paper, we propose a Sinkhorn re-basin network with the ability to obtain the transportation plan that better suits a given objective. Unlike the current state-of-art, our method is differentiable and, therefore, easy to adapt to any task within the deep learning domain. Furthermore, we show the advantage of our re-basin method by proposing a new cost function that allows performing incremental learning by exploiting the linear mode connectivity property. The benefit of our method is compared against similar approaches from the literature, under several conditions for both optimal transport finding and linear mode connectivity. The effectiveness of our continual learning method based on re-basin is also shown for several common benchmark datasets, providing experimental results that are competitive with state-of-art results from the literature.