Continuous-Depth Transformers with Learned Control Dynamics

TL;DR

This paper introduces a hybrid transformer that replaces a portion of discrete layers with a continuous-depth neural ODE block driven by a learnable control signal $u$, enabling inference-time steering of generation attributes. Stability is ensured by a learned output scale $\alpha$, and the model demonstrates zero gradient pathologies, continuous steerable behavior, and efficient latency comparable to standard transformers. Through four experiments, the approach achieves high semantic steering accuracy, shows negligible trajectory divergence between fixed and adaptive solvers, and reveals geometry of the learned dynamics via solver probing. The work offers a practical, interpretable framework for steerable language generation with potential for dynamic computation and multi-dimensional control, while acknowledging scale and generalization limits to larger models as future work.

Abstract

We present a hybrid transformer architecture that replaces discrete middle layers with a continuous-depth Neural Ordinary Differential Equation (ODE) block, enabling inference-time control over generation attributes via a learned steering signal. Unlike standard transformers that process representations through fixed discrete layers, our approach treats depth as a continuous variable governed by a learned vector field $F_θ(H, τ, u)$, where $u$ is a low-dimensional control signal injected via explicit concatenation. We validate the architecture through four experiments: (1) gradient flow stability with zero exploding/vanishing gradient events, (2) semantic steering achieving 98\%/88\% accuracy for positive/negative sentiment control, (3) continuous interpolation validated by a negligible 0.068\% trajectory divergence between fixed and adaptive solvers, and (4) efficiency benchmarking demonstrating latency parity with standard discrete baselines. Additionally, we show that adaptive ODE solvers reveal geometric structure in the learned dynamics: the control signal partitions the vector field into distinct dynamical regimes with different curvature characteristics. The adjoint method enables $O(1)$ memory training regardless of integration depth. Our results demonstrate that continuous-depth dynamics with learned control signals provide a viable, efficient mechanism for steerable language generation.

Figures (4)

• Figure 1: Training dynamics. Left: Loss curves converge similarly. Center: Total gradient norms are comparable. Right: ODE block gradients remain healthy throughout training.
• Figure 2: Control signal sweep. Smooth sigmoid curves demonstrate continuous steering---intermediate values produce mixed sentiment states.
• Figure 3: Solver effort reveals two dynamical regimes. The control signal partitions the vector field into regions of different curvature, with the transition aligned to the semantic crossover.
• Figure 4: Resolution scaling toward continuous limit. NFE increases sublinearly with tighter tolerance, confirming smooth learned dynamics.