rope-conformer

A clean, plug-and-play PyTorch Conformer encoder with rotary positional embeddings (RoPE). The block layout follows the original Conformer architecture from [1] (Gulati et al., 2020) — Macaron half-FFN → multi-head self-attention → depthwise conv module → Macaron half-FFN — with the classic Shaw relative positional embedding swapped for rotary embeddings [2] (Su et al., 2021) and the rest of the transformer-side conventions modernized to match the lucidrains/BS-RoFormer family used in current music source separation work [3] (RMSNorm pre-norm everywhere, fused F.scaled_dot_product_attention, GELU FFN, no biases on QKV / output projections).

Quick Start

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]

Installation

Install from PyPI

pip install rope-conformer

Optional Apple-silicon acceleration (experimental — see Optional mps-sdpa integration):

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

Install from GitHub

pip install git+https://github.com/crlandsc/rope-conformer.git

Or, you can clone the repository and install it in editable mode for development:

git clone https://github.com/crlandsc/rope-conformer.git
cd rope-conformer
pip install -e ".[dev]"

Dependencies

• PyTorch (>=2.3)
• einops (>=0.8)
• rotary-embedding-torch (>=0.6)

Optional:

• mps-sdpa (>=0.2.0) — fused MPSGraph SDPA on Apple silicon

Usage

Padding mask

Pad-mask any positions that should be ignored by attention:

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

Stack your own blocks

If you want to compose blocks manually rather than use the top-level stack, the public API also exposes ConformerBlock and the ConformerConvModule sublayer:

from rope_conformer import ConformerBlock, RotaryEmbedding

rotary_embed = RotaryEmbedding(dim=32)              # one shared instance per stack
block = ConformerBlock(dim=256, rotary_embed=rotary_embed, heads=8, dim_head=32)
y = block(torch.randn(2, 100, 256))                 # [B, N, dim]

API

RoPEConformer constructor arguments:

Argument	Default	Description
dim	—	Model channels (input and output, unless output_dim is set).
depth	—	Number of stacked Conformer blocks.
dim_head	64	Per-head dimension.
heads	8	Number of attention heads.
ff_mult	4	Feedforward expansion factor.
conv_expansion_factor	2	Pointwise expansion in the conv module.
conv_kernel_size	31	Depthwise conv kernel (Conformer paper default).
attn_dropout	0.0	Dropout inside attention (off by default).
proj_dropout	0.1	Dropout after attention output projection.
ff_dropout	0.1	Dropout in feedforward sublayers.
conv_dropout	0.1	Dropout at the end of the conv module.
conv_causal	False	Left-pad the depthwise conv (no future-frame leakage).
conv_norm_type	"rms"	"rms" (per-token, causal-safe), "group" or "batch" (cross-time stats; both fall back to Identity in causal mode).
use_attn_gates	False	Add per-head sigmoid output gates on attention.
flash_attn	True	Use F.scaled_dot_product_attention; False falls back to einsum.
use_mps_sdpa	False	Route attention through mps-sdpa (experimental, MPS only).
norm_output	True	Apply final RMSNorm before projection.
output_dim	None	Optional output projection to a different dimension.

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.

Contributing

Contributions are welcome! Please open an issue or submit a pull request if you have any bug fixes, improvements, or new features to suggest.

License

This project is licensed under the MIT License - see the LICENSE file for details.

Acknowledgments

The Conformer architecture is from Gulati et al., 2020 [1]. Rotary positional embeddings are from Su et al., 2021 [2], and use the implementation from lucidrains/rotary-embedding-torch. The modernized transformer conventions (RMSNorm, fused SDPA, GELU FFN, no QKV/output biases) follow the patterns established in lucidrains/BS-RoFormer [3]. Thank you for your work!

References

[1] A. Gulati, J. Qin, C.-C. Chiu, N. Parmar, Y. Zhang, J. Yu, W. Han, S. Wang, Z. Zhang, Y. Wu, and R. Pang, "Conformer: Convolution-augmented Transformer for Speech Recognition," Proceedings of Interspeech 2020. arXiv:2005.08100

[2] J. Su, Y. Lu, S. Pan, A. Murtadha, B. Wen, and Y. Liu, "RoFormer: Enhanced Transformer with Rotary Position Embedding," 2021. arXiv:2104.09864

[3] W.-T. Lu, J.-C. Wang, Q. Kong, and Y.-N. Hung, "Music Source Separation with Band-split RoPE Transformer," ICASSP 2024. arXiv:2309.02612