rope-conformer

A clean, plug-and-play PyTorch Conformer encoder with rotary positional embeddings (RoPE).

import torch
from rope_conformer import RoPEConformer

model = RoPEConformer(dim=256, depth=6, heads=8, dim_head=32)
x = torch.randn(2, 100, 256)        # [B, N, dim]
y = model(x)                        # [B, N, dim]

Install

pip install rope-conformer

Optional Apple-silicon acceleration (experimental):

pip install "rope-conformer[mps-sdpa]"

Usage with a padding mask

mask = torch.zeros(2, 100, dtype=torch.bool)
mask[:, 80:] = True                 # last 20 positions padded
y = model(x, key_padding_mask=mask) # masked positions are ignored in attention

API

The forward signature is model(x, key_padding_mask=None):

• x: [B, N, dim]
• key_padding_mask: [B, N] bool, True for padded positions (optional).

Granular dropout

The four dropout knobs (attn_dropout, proj_dropout, ff_dropout, conv_dropout) target different stages so you can tune each independently. Attention dropout is off by default because zeroing entries before the softmax distorts the resulting probability distribution and creates a training/inference mismatch in the attention pattern; the other three knobs act on unnormalized intermediate features and don't have this issue, so they default to a small 0.1.

Causal use

Set conv_causal=True for a depthwise conv that only sees past frames. This handles the conv path; the self-attention path is not causally masked by this flag — pass your own causal attn_mask (or use a causal-aware downstream stack) if you need full causal behavior.

Optional mps-sdpa integration (experimental)

PyTorch's MPS backend does not currently dispatch scaled_dot_product_attention to Apple's fused MPSGraph.scaledDotProductAttention op; it builds a naive matmul → softmax → matmul graph instead. The mps-sdpa package wraps the fused op directly, giving roughly 5–7× faster inference and 2–2.5× faster training on Apple silicon (M1+, macOS 15+). When installed via the [mps-sdpa] extra, attention can be routed through it per-instance:

model = RoPEConformer(dim=256, depth=6, use_mps_sdpa=True)

If the mps-sdpa package is not installed, the flag silently no-ops and the standard SDPA path is used. The flag also has no effect on non-MPS devices.

Acknowledgements

The architecture mirrors the modernized transformer conventions used in lucidrains/BS-RoFormer (RoPE, RMSNorm, SDPA, GELU FFN, no biases on QKV/output projections), applied to the Conformer block (Gulati et al., 2020).

License

MIT.