OT-Transformer: A Continuous-time Transformer Architecture with Optimal Transport Regularization

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Kan, Kelvin, Li, Xingjian, Osher, Stanley
Format: Preprint
Published: 2025
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866909470908481536
author Kan, Kelvin
Li, Xingjian
Osher, Stanley
author_facet Kan, Kelvin
Li, Xingjian
Osher, Stanley
contents Transformers have achieved state-of-the-art performance in numerous tasks. In this paper, we propose a continuous-time formulation of transformers. Specifically, we consider a dynamical system whose governing equation is parametrized by transformer blocks. We leverage optimal transport theory to regularize the training problem, which enhances stability in training and improves generalization of the resulting model. Moreover, we demonstrate in theory that this regularization is necessary as it promotes uniqueness and regularity of solutions. Our model is flexible in that almost any existing transformer architectures can be adopted to construct the dynamical system with only slight modifications to the existing code. We perform extensive numerical experiments on tasks motivated by natural language processing, image classification, and point cloud classification. Our experimental results show that the proposed method improves the performance of its discrete counterpart and outperforms relevant comparing models.
format Preprint
id arxiv_https___arxiv_org_abs_2501_18793
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle OT-Transformer: A Continuous-time Transformer Architecture with Optimal Transport Regularization
Kan, Kelvin
Li, Xingjian
Osher, Stanley
Machine Learning
Artificial Intelligence
Transformers have achieved state-of-the-art performance in numerous tasks. In this paper, we propose a continuous-time formulation of transformers. Specifically, we consider a dynamical system whose governing equation is parametrized by transformer blocks. We leverage optimal transport theory to regularize the training problem, which enhances stability in training and improves generalization of the resulting model. Moreover, we demonstrate in theory that this regularization is necessary as it promotes uniqueness and regularity of solutions. Our model is flexible in that almost any existing transformer architectures can be adopted to construct the dynamical system with only slight modifications to the existing code. We perform extensive numerical experiments on tasks motivated by natural language processing, image classification, and point cloud classification. Our experimental results show that the proposed method improves the performance of its discrete counterpart and outperforms relevant comparing models.
title OT-Transformer: A Continuous-time Transformer Architecture with Optimal Transport Regularization
topic Machine Learning
Artificial Intelligence
url https://arxiv.org/abs/2501.18793