A research prototype that recasts Transformer attention as a continuous-depth process. Instead of a fixed stack of discrete attention layers, CAT integrates a continuous-attention operator over depth using Neural ODEs, with attention kernels parameterized in a spectral basis and regularized for smoothness.
This repository is a single-file, self-contained implementation intended for exploring the idea and reproducing early results on a small character-level language-modeling task. It is exploratory research code, not a production model.
A standard Transformer applies attention in a fixed number of discrete layers. CAT reframes this as evolving the hidden state through continuous depth: each block is a Neural ODE whose dynamics are a continuous-attention layer plus a feed-forward update, integrated from depth 0 to 1.
The attention kernel itself is defined over continuous sequence position. Rather
than computing pairwise token scores, each attention head learns a smooth kernel
W(s, t) built from a truncated spectral basis (Fourier cosines or Chebyshev
polynomials). The kernel is made row-stochastic with a softmax, and the
attention output is computed as an integral over position, approximated with
Gauss-Legendre quadrature. A Sobolev penalty on the spectral coefficients
discourages over-oscillatory kernels and encourages smoothness.
All in CAT.py:
Kernel2D— a 2D attention kernel expanded in a Fourier or Chebyshev basis, softmax-normalized to be row-stochastic, with a Sobolev (biharmonic) penalty on its coefficients.ContinuousAttention— multi-head continuous attention. Values are projected and interpolated at Gauss-Legendre quadrature nodes; the kernel integral is evaluated per head.ODEFunc— one CAT block: continuous attention with a residual connection, layer normalization, and a GELU feed-forward network, used as ODE dynamics.CAT— token embedding,Kstacked ODE blocks, and a linear projection to vocabulary logits, with the Sobolev penalty summed across blocks.Baseline— a standardnn.TransformerEncodertrained alongside CAT for comparison.- Depth integration uses
torchdiffeq(dopri5) on CPU/CUDA, and a manual RK4 integrator on Apple Silicon (MPS).
The demo trains on character-level Tiny-Shakespeare (Karpathy's char-rnn
input.txt, downloaded on first run; a copy is included as tiny.txt). The
first 50,000 characters are used for training and the remainder for validation.
Training runs CAT and the baseline Transformer simultaneously with AdamW and cross-entropy loss; CAT additionally uses the annealed Sobolev penalty. After training, the script reports cross-entropy and perplexity for both models at context lengths of 128, 512, 1024, and 5000 tokens, along with the difference in cross-entropy. Evaluating across a range of context lengths reflects the motivation for the continuous formulation: the operator is defined over continuous position rather than a fixed token grid.
pip install torch numpy torchdiffeqDefaults run on CPU with a single block. Example run with the Fourier basis on Apple Silicon:
python CAT.py \
--device mps \
--basis fourier \
--blocks 3 \
--train_steps 2000 \
--d_model 256 \
--n_heads 8 \
--n_freq 48 \
--sob_lambda 5e-7Key arguments (see CAT.py for defaults):
--device—cpu,cuda, ormps.--basis—fourierorcheb(Chebyshev).--blocks— number of stacked ODE blocks.--d_model,--n_heads,--n_freq— model width, heads, and spectral frequencies per axis.--sob_lambda— weight of the Sobolev smoothness penalty.--train_steps,--lr,--batch,--seq— optimization settings.--eval_windows— maximum number of validation windows per context length.
.
├── CAT.py # Full implementation: model, baseline, training, evaluation
├── tiny.txt # Tiny-Shakespeare corpus (downloaded on first run if absent)
└── LICENSE
Early-stage research code for demonstration and reproducibility of preliminary results. Numbers depend on hyperparameters, hardware, and random seeds.
Released under the MIT License. See LICENSE.