State-space models through the lens of ensemble control

TL;DR

This paper presents a control-theoretic framework for understanding and optimizing the training of state-space models by formulating it as an ensemble optimal control problem and deriving necessary optimality conditions.

Contribution

It introduces a novel ensemble control formulation for SSM training, derives Pontryagin's maximum principle, and proposes a convergent iterative algorithm with global optimality conditions.

Findings

Derived PMP for ensemble control of SSMs

Proposed an iterative approximation algorithm with convergence guarantees

Provided a control-theoretic perspective on SSM training

Abstract

State-space models (SSMs) are effective architectures for sequential modeling, but a rigorous theoretical understanding of their training dynamics is still lacking. In this work, we formulate the training of SSMs as an ensemble optimal control problem, where a shared control law governs a population of input-dependent dynamical systems. We derive Pontryagin's maximum principle (PMP) for this ensemble control formulation, providing necessary conditions for optimality. Motivated by these conditions, we introduce an algorithm based on the method of successive approximations. We prove convergence of this iterative scheme along a subsequence and establish sufficient conditions for global optimality. The resulting framework provides a control-theoretic perspective on SSM training.