nn(sequence-vision): add convolutional, recurrent, attention families, and masterclass docs/examples
- Dominant language
- Python
- Stars
- 1
- Forks
- 0
- Avg merge
- 19h 2m
- Merged PRs (30d)
- 16
Description
## Problem / Request
`pyrox.nn` ships dense Bayesian layer families and a Multi-Head Attention BatchEnsemble layer (`MultiHeadAttentionBE`), but it does **not** ship convolutional / recurrent / full-attention Bayesian layer families. Sequence and vision models that want to compose `pyrox`'s uncertainty primitives currently have to hand-wire `numpyro.sample` calls around each `eqx.nn.Conv2d` / `equinox.nn.GRU` / `equinox.nn.MultiheadAttention` — losing the immutable-module discipline and breaking SVI batching.
## User Story
> As a user fitting a Bayesian CNN for image classification (or a Bayesian Transformer for time-series), I want `BayesianConv2d`, `BayesianGRU`, and `BayesianMultiheadAttention` so I compose them like dense layers and get the same `numpyro` prior wiring and SVI batching for free.
## Motivation
Concrete model patterns that need these:
- **Bayesian CNN for image regression / classification** with calibrated uncertainty — currently no Bayesian Conv2d in pyrox.
- **Bayesian time-series forecasting** — `MultiScaleSIREN` (#91) handles INR-style; `BayesianGRU` / `BayesianLSTM` is the recurrent counterpart for variable-length sequences.
- **Bayesian Transformer for tabular / text** — full-attention Bayesian layer extends `MultiHeadAttentionBE` (which is BatchEnsemble-specific) to the standard variational reparameterisation pattern.
- **Sequence-labelling with `LinearChainCRF` head** (#51) — the upstream feature extractor is most naturally a `BayesianGRU` or `BayesianTransformer`.
Without these, sequence / vision modelling is a second-class citizen in `pyrox`.
## Mathematics
The math for each layer is standard — Bayesian convolution / recurrence / attention is the same as the deterministic version with priors on the weight tensors. The complexity is in the **JAX wiring**:
**Bayesian Conv2d.** Weight tensor $W \in \mathbb{R}^{C_{\text{out}} \times C_{\text{in}} \times K \times K}$ with a Gaussian prior; reparameterised forward,
$$
\hat W \;=\; \mu_W + \sigma_W \odot \epsilon, \quad \epsilon \sim \mathcal{N}(0, I),
$$
then `lax.conv_general_dilated(x, W_hat, …)`. KL divergence to prior contributes one term per spatial location $\times C_{\text{out}} \times C_{\text{in}}$.
**Bayesian GRU.** Reparameterised input + recurrent kernels; the hidden-state recurrence
$$
h_t \;=\; (1 - z_t) \odot h_{t-1} + z_t \odot \tilde h_t
$$
is unchanged; only the weight matrices that build $z_t$ and $\tilde h_t$ are Bayesian.
**Bayesian MultiheadAttention.** Reparameterised Q/K/V/O projections; attention scores
$$
\mathrm{attn}(Q, K, V) \;=\; \mathrm{softmax}\!\left(\frac{Q K^{\top}}{\sqrt{d_k}}\right) V
$$
is unchanged.
## Proposed API
```python
class BayesianConv2d(PyroxModule):
in_channels: int
out_channels: int
kernel_size: int | tuple[int, int]
stride: int | tuple[int, int]
padding: str | int
prior_scale: float # static
posterior_loc: Float[Array, "Cout Cin K K"]
posterior_scale: Float[Array, "Cout Cin K K"]
def __call__(self, x, *, key) -> Float[Array, "Cout H' W'"]: ...
class BayesianGRU(PyroxModule):
in_features: int
hidden: int
prior_scale: float
# ... reparameterised kernel weights ...
def __call__(
self, x_seq: Float[Array, "T in"], *, key,
) -> Float[Array, "T hidden"]: ...
class BayesianMultiheadAttention(PyroxModule):
embed_dim: int
num_heads: int
prior_scale: float
# ... reparameterised Q, K, V, O ...
def __call__(
self, x: Float[Array, "T embed"], *, key,
) -> Float[Array, "T embed"]: ...
```
All three reuse the `DenseReparameterization` pattern for prior / posterior wiring — the only new code is the operator-specific forward (`conv_general_dilated`, `lax.scan`, `dot_product_attention`).
## Example usage
```python
import jax, pyrox.nn as pnn, equinox as eqx
# Bayesian CNN for image classification
key = jax.random.PRNGKey(0)
keys = jax.random.split(key, 3)
cnn = eqx.nn.Sequential([
pnn.BayesianConv2d(in_channels=3, out_channels=32, kernel_size=3, key=keys[0]),
eqx.nn.Lambda(jax.nn.relu),
pnn.BayesianConv2d(in_channels=32, out_channels=64, kernel_size=3, key=keys[1]),
eqx.nn.Lambda(jax.nn.relu),
eqx.nn.Lambda(lambda x: x.reshape(-1)),
pnn.DenseReparameterization(in_features=..., out_features=10, key=keys[2]),
])
# Bayesian GRU for sequence regression
rnn = pnn.BayesianGRU(in_features=8, hidden=64, key=key)
```
## Tasks / sub-tasks
- [ ] **Spike** — confirm `lax.conv_general_dilated` with a reparameterised kernel is `vmap`-friendly under the BatchEnsemble pattern (one independent `epsilon` per ensemble member).
- [ ] **`BayesianConv2d`** — `src/pyrox/nn/_conv.py`. Reuses the `DenseReparameterization` posterior pattern.
- [ ] **`BayesianGRU` / `BayesianLSTM`** — `src/pyrox/nn/_recurrent.py`. `lax.scan`-based forward; reparameterised input + recurrent kernels.
- [ ] **`BayesianMultiheadAttention`** — `src/pyrox/nn/_attention.py`. Reparameterised Q/K/V/O; complements existing `MultiHeadAttentionBE`.
- [ ] **Composition tests** with `LinearChainCRF` (#51) and `MixtureOfGaussiansDenseFA` (#51) — confirm sequence / mixture heads compose with each new layer family.
- [ ] **Tests** — (i) each layer matches its deterministic counterpart when prior scale → 0; (ii) `numpyro.handlers.seed` + `trace` recovers the expected number of sample sites; (iii) `jit` / `vmap` / `grad` compatibility; (iv) BatchEnsemble interop check.
- [ ] **Notebook example** — small Bayesian CNN on a toy image classification dataset with reliability-curve calibration.
- [ ] **Docs** — `pyrox.nn` reference entries; "Bayesian sequence / vision" tutorial page.
## References
- Kingma, D. P., Salimans, T. & Welling, M. (2015). *Variational dropout and the local reparameterization trick*. NeurIPS — covers the BatchEnsemble + reparameterised-conv pattern.
- Gal, Y. & Ghahramani, Z. (2015). *Bayesian convolutional neural networks with Bernoulli approximate variational inference*. arXiv:1506.02158.
- Existing layers: [`src/pyrox/nn/_layers.py`](https://github.com/jejjohnson/pyrox/blob/main/src/pyrox/nn/_layers.py), [`src/pyrox/nn/_ensemble.py`](https://github.com/jejjohnson/pyrox/blob/main/src/pyrox/nn/_ensemble.py) (`MultiHeadAttentionBE`).
## Relationships
- Parent (theme epic): #46
- Blocked by: #
- Blocks: #
- Related: #51, #91
Contributor guide
Assessment
This issue has not been assessed yet.