torch-slices 0.3.3

Structured Linear CDE (SLiCE) layers for sequence modelling in PyTorch

SLiCEs

slices is a PyTorch implementation of Structured Linear CDEs (SLiCEs) with parallel-in-time computation.

It provides expressive, scalable sequence layers based on the methods developed in Structured Linear CDEs: Maximally Expressive and Parallel-in-Time Sequence Models.

Mathematical form

Given an input sequence $x_i \in \mathbb{R}^D$ for $i=1,\dots,T$, a SLiCE computes hidden states $y_i \in \mathbb{R}^H$ via

$$ y_i = y_{i-1} + A(X_i)y_{i-1} + B(X_i), $$

where $A(\cdot): \mathbb{R}^D \rightarrow \mathbb{R}^{H \times H}$ and $B(\cdot): \mathbb{R}^D \rightarrow \mathbb{R}^H$ are learned linear maps, the initial state $y_0$ is either a function of $X_0$ or a learnt vector, and the driving path is augmented with an extra channel:

• inc = a constant “increment” channel (all ones)

such that

$$ X_i = [inc_i, x_i] \in \mathbb{R}^{D+1}. $$

By default, SLiCE treats the provided sequence as path values and internally computes first differences with torch.diff(..., prepend=zeros). Set path_mode="increments" to treat the provided sequence as increments instead. By default, each step uses the first-order Euler update I + A(X_i). Set transition_mode="matrix_exp" to use exp(A(X_i)) instead.

Installation

pip install torch-slices

Or install from source:

pip install git+https://github.com/datasig-ac-uk/slices.git

Components

• SLiCE: the core structured recurrence
• SLiCELayer: a residual sequence layer built around SLiCE, with RMSNorm and a GELU MLP by default
• StackedSLiCE: a full sequence model that stacks SLiCELayers with input and output projections for token or continuous inputs

SLiCE supports both:

• Recurrent execution for step-by-step updates
• Parallel chunked scan execution via torch.associative_scan

Structured transition matrices

SLiCE supports different $A(X_i)$ structures:

1) Diagonal (elementwise update)

Set:

• diagonal_dense=False
• block_size=1

Then $A(X_i)$ is diagonal, matching the diagonal state-transition setting used by Mamba (as discussed in Structured Linear CDEs: Maximally Expressive and Parallel-in-Time Sequence Models).

2) Block-diagonal

Set:

• diagonal_dense=False
• block_size > 1

Then $A(X_i)$ is block-diagonal with blocks of shape (block_size × block_size).

3) Diagonal + dense tail block

Set:

• diagonal_dense=True
• block_size > 1

Then the first (hidden_dim - block_size) dimensions are diagonal, and the final block_size dimensions are updated via a dense (block_size × block_size) matrix.

Quickstart

Use SLiCE directly

import torch
from slices import SLiCE

x = torch.randn(8, 128, 32)  # (batch, seq, input_dim)

layer = SLiCE(
    input_dim=32,
    hidden_dim=64,
    block_size=4,
    diagonal_dense=False,
    bias=True,
    use_parallel=True,
    chunk_size=256,
)

y = layer(x)  # (8, 128, 64)
print(y.shape)

Execution mode is controlled by use_parallel and chunk_size. path_mode determines whether the input sequence is interpreted as path values or increments. transition_mode selects between the default Euler step and a matrix-exponential transition.

Use SLiCELayer as a residual sequence layer

import torch
from slices import SLiCELayer

x = torch.randn(4, 256, 64)  # (batch, seq, input_dim)

layer = SLiCELayer(
    input_dim=64,
    block_size=4,
    diagonal_dense=True,
)

y = layer(x)  # (4, 256, 64)

SLiCELayer uses the following default structure:

• RMSNorm -> SLiCE -> residual
• RMSNorm -> Linear -> GELU -> Linear -> residual

Optional configuration options include:

• prenorm=False
• second_norm=False
• norm_type="layernorm"
• ff_style="single"
• ff_mult=1
• ff_activation="glu" or ff_activation="tanh"
• dropout_position="output"

Build a full model with stacked layers

Token sequence mode (tokens=True)

Uses an nn.Embedding(data_dim, hidden_dim) front-end.

import torch
from slices import StackedSLiCE

batch, seq_len = 2, 128
vocab_size = 5000

x = torch.randint(0, vocab_size, (batch, seq_len))

model = StackedSLiCE(
    num_layers=4,
    data_dim=vocab_size,
    hidden_dim=256,
    label_dim=vocab_size,
    tokens=True,
    block_size=4,
    diagonal_dense=False,
)

logits = model(x)  # (batch, seq_len, vocab_size)

Continuous time-series mode (tokens=False)

Uses an nn.Linear(data_dim, hidden_dim) front-end.

import torch
from slices import StackedSLiCE

x = torch.randn(16, 100, 12)  # (batch, seq, input_dim)

model = StackedSLiCE(
    num_layers=3,
    data_dim=12,
    hidden_dim=64,
    label_dim=10,
    tokens=False,
    block_size=4,
    diagonal_dense=True,
)

y = model(x)  # (16, 100, 10)

Training Example

examples/language_disambiguation.py trains a StackedSLiCE model end-to-end on a real dataset.

The example:

• uses a real dataset (wikimedia/wikipedia, English/French subset) for character-level language disambiguation
• trains a compact token-mode StackedSLiCE end-to-end
• evaluates validation accuracy every --eval-every training steps
• prints sample predictions for quick qualitative inspection

To run it, first install the example dependencies:

uv sync --group examples

and then execute the script:

uv run python examples/language_disambiguation.py

Useful flags:

• --device cpu|cuda
• --train-steps 3000
• --eval-every 300
• --train-size 12000 --val-size 3000
• --max-seq-len 192

Self-Supervised Example

examples/language_selfsupervised.py uses the same streaming Wikipedia data loader, trains a self-supervised token model, exports hidden states to a memmap, and can optionally plot UMAP projections from those exports.

Generated artifacts are written under examples/outputs/language_selfsupervised/ by default.

Install the example dependencies:

uv sync --group examples

Train, export, and plot with the default CPU UMAP backend:

uv run python examples/language_selfsupervised.py --plot-umap

To use the CUDA UMAP backend in the project environment:

uv sync --group examples --group cuda-umap
uv run python examples/language_selfsupervised.py --device cuda --plot-umap --umap-backend cuda

Useful flags:

• --horizon single|multi
• --export-mode all|final
• --plot-umap
• --umap-backend auto|cpu|cuda
• --plot-time-indices 1,20,100,-5,-1,0 where 0 means H_T, the hidden state after consuming EOS and aligned with the first PAD target; negative indices are relative to that point, and positive indices are one-indexed absolute timesteps
• --n-per-lang 4000

The CUDA UMAP path is optional. --umap-backend auto falls back to CPU UMAP if the CUDA packages are unavailable.

Benchmarking

To compare recurrent and parallel throughput across sequence lengths and hidden dimensions:

uv run python examples/benchmark_parallel_vs_recurrent.py

This script:

• benchmarks all four SLiCE matrix modes (diagonal, block_diagonal, diagonal_dense, dense)
• uses the default value-path semantics unless path_mode="increments" is set in code
• prints timing and speedup tables
• saves a combined 3D plot to examples/images/parallel_vs_recurrent_speedup_3d_all_modes.png

For plotting in development, install development dependencies (includes matplotlib):

uv sync --dev

Requirements

• Python ≥ 3.11
• PyTorch ≥ 2.8.0
• NumPy ≥ 2

License

MIT License. See LICENSE.