ET-KAN: an energy-based transformer model with Kolmogorov–Arnold network for image reconstruction

Transformers models have significantly changed many areas of Machine Learning, due to their structure and extensive number of parameters, which enables them to capture complex patterns in data. Building on this foundation, an energy-based version, the Energy Transformer (ET), has emerged as a powerful variation, achieving parameter efficiency without sacrificing performances. In this work, we introduce an innovative evolution of the ET model: the ET-KAN architecture. By integrating a Kolmogorov–Arnold Network (KAN) in place of the energy function, our model generalizes the ET as structural design, unlocking enhanced learning capabilities. We demonstrate the potential of this new architecture through an image reconstruction task, where it achieves comparable or higher results respect to the ET ($$Loss\approx 0.08$$ when covering more than half of the image), outperforming the standard ET architecture while using fewer parameters. The model hereby constructed paves the way to a deeper investigation of the interplay between KAN and energy-based models, for addressing some of the key limitations of traditional transformers.