stackformers 4.2.0__tar.gz → 4.3.0__tar.gz

{stackformers-4.2.0 → stackformers-4.3.0}/CHANGELOG.md

@@ -6,6 +6,71 @@ The format follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), and
 adheres to [Semantic Versioning](https://semver.org/): MAJOR for breaking public API changes,
 MINOR for backwards-compatible features, PATCH for bug fixes and internal changes.
 
+## [4.3.0] — 2026-07-17
+
+### Added
+
+- **`MLMWrapper`** (`stackformers/mlm/`) — a masked-token-reconstruction auxiliary loss over any
+  encoder satisfying `EncoderLike`, domain-agnostic about what a token represents. Takes the
+  encoder at call time (`forward(input, encoder)`) rather than owning it, so nothing is
+  registered as a submodule and `mlm_wrapper.parameters()` never includes the encoder's weights.
+- `MLMWrapperConfig`, `MaskingStrategy`/`RandomMasking`, `ReconstructionHead`/`RegressionHead`,
+  and `MLMOutput` (`out`, `mlm_loss`).
+
+### Notes
+
+- **`out` is always the encoder's clean, unmasked output; only `mlm_loss` ever reflects
+  masking.** In training, `forward` runs the encoder twice — once clean (returned as `out`) and
+  once on a separately-masked copy (used only to compute `mlm_loss`) — so the main pipeline sees
+  byte-identical output whether or not this loss is being trained alongside it. In eval, only the
+  clean pass runs and `mlm_loss` is a constant zero, gated on `self.training` the same way
+  `nn.Dropout` and `nn.BatchNorm` already are — so callers can invoke it unconditionally in both
+  modes with no `if training` branch of their own.
+
+- **The reconstruction target is always `input.x.detach()`.** This severs the gradient path from
+  the loss back to whatever produced `input.x`, removing the collapse shortcut a trainable
+  tokenizer would otherwise have (drive every token toward one constant vector to make
+  reconstruction trivial).
+
+- **`RandomMasking` needs no packed-sequence boundary awareness** — each token's masking decision
+  is independent of every other's, so document identity (`cu_seqlens`) never enters it. That only
+  becomes necessary for a contiguous-span (blockwise) strategy, which could otherwise straddle
+  two packed documents.
+
+## [4.2.1] — 2026-07-14
+
+### Removed
+
+- **`RoPENDConfig.headroom`.** It was redundant with `r_max`: the ladder only ever depended on
+  the product `headroom · r_max`, so the two knobs shared one degree of freedom, and anyone
+  wanting more reach at the slow end could get it by raising `r_max` — the parameter that has a
+  measurement procedure attached. It was also the last *tuned* number in a module whose premise
+  is that the band range is measured, not tuned.
+
+  Its default has a derivation, so it need not be a parameter at all. Attention sees *signed*
+  offsets spanning `[−r_max, +r_max]`, a width of `2·r_max`; asking that the slowest band turn
+  through at most half a circle across that width gives `ω_lo = π / (2·r_max)` directly. That is
+  bit-identical to the old `headroom=4.0` default at every `r_max`, so callers on the default —
+  which is to say every caller, since the field shipped in 4.2.0 — see no change in behaviour.
+
+  Both ends of the ladder are now the same rule, `ω = π / scale`: a half turn over `r_min` at
+  the fast end (Nyquist), and a half turn over `2·r_max` at the slow end (no wrap).
+
+  Treated as a PATCH rather than a MAJOR bump because `headroom` existed for exactly one
+  release, is not known to be set anywhere, and its removal cannot change the behaviour of a
+  caller that did not set it. Note that pydantic's default `extra="ignore"` means a leftover
+  `RoPENDConfig(headroom=...)` is *dropped silently* rather than raising — a caller who had set
+  it to something other than 4.0 will now get different frequencies without being told.
+
+### Fixed
+
+- **`RoPENDConfig` now requires `dim_head >= 4 * coords`** (at least two bands per axis).
+  `dim_head == 2 * coords` divides cleanly and passed validation, but leaves a one-band ladder
+  with nowhere to descend: `torch.linspace(hi, lo, 1)` returns `[hi]`, so the lone band lands on
+  the fast end, `r_max` is discarded entirely, and the encoding becomes periodic with period
+  `2 · r_min` across the whole domain — offsets of `0`, `2·r_min`, `4·r_min` … all produce an
+  identical attention logit. It failed silently; it is now rejected at construction.
+
 ## [4.2.0] — 2026-07-14
 
 ### Added

{stackformers-4.2.0 → stackformers-4.3.0}/CLAUDE.md

@@ -25,6 +25,7 @@ Behavior comes from injected dependencies, not constructor flags.
 | dh     | dim per head |
 | w      | window size |
 | nt     | total tokens in packed sequence |
+| m      | number of masked tokens (gathered across batch/pack) |
 
 ## File organisation
 

{stackformers-4.2.0 → stackformers-4.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: stackformers
-Version: 4.2.0
+Version: 4.3.0
 Summary: Typed, composable, SOLID transformer library for PyTorch
 Project-URL: Homepage, https://github.com/Red-Eyed/stackformers
 Project-URL: Repository, https://github.com/Red-Eyed/stackformers

{stackformers-4.2.0 → stackformers-4.3.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "stackformers"
-version = "4.2.0"
+version = "4.3.0"
 description = "Typed, composable, SOLID transformer library for PyTorch"
 readme = "README.md"
 license = { file = "LICENSE" }

{stackformers-4.2.0 → stackformers-4.3.0}/stackformers/__init__.py

@@ -17,6 +17,11 @@ from stackformers.feedforward.factory import build_ff
 from stackformers.feedforward.protocols import FeedForward
 from stackformers.feedforward.swiglu import SwiGLU
 from stackformers.layers import TransformerLayer
+from stackformers.mlm.config import MLMWrapperConfig
+from stackformers.mlm.head import RegressionHead
+from stackformers.mlm.masking import RandomMasking
+from stackformers.mlm.protocols import EncoderLike, MaskingStrategy, ReconstructionHead
+from stackformers.mlm.wrapper import MLMOutput, MLMWrapper
 from stackformers.norm.config import LayerNormConfig, RMSNormConfig
 from stackformers.norm.factory import NormConfig, build_norm
 from stackformers.norm.protocols import Norm
@@ -85,6 +90,9 @@ __all__ = [
     "CrossAttn",
     "FeedForward",
     "Norm",
+    "EncoderLike",
+    "MaskingStrategy",
+    "ReconstructionHead",
     # configs — attention
     "SelfAttentionConfig",
     "CrossAttentionConfig",
@@ -133,4 +141,10 @@ __all__ = [
     "CrossAttenderConfig",
     "CrossAttender",
     "plain_cross_attender_config",
+    # mlm
+    "MLMWrapperConfig",
+    "RandomMasking",
+    "RegressionHead",
+    "MLMWrapper",
+    "MLMOutput",
 ]

stackformers-4.3.0/stackformers/mlm/README.md

@@ -0,0 +1,54 @@
+# mlm
+
+A masked-token-reconstruction auxiliary loss, wrapped around any encoder that satisfies `EncoderLike`. Domain-agnostic — it has no idea what a token represents.
+
+## Usage
+
+```python
+import torch
+from stackformers import (
+    MLMWrapper,
+    MLMWrapperConfig,
+    TransformerEncoder,
+    make_padded_input,
+    plain_encoder_config,
+)
+
+encoder = TransformerEncoder(plain_encoder_config(dim=512, heads=8, num_layers=6))
+mlm_wrapper = MLMWrapper(MLMWrapperConfig(dim=512, mask_ratio=0.15))
+
+x = torch.randn(2, 128, 512)
+mask = torch.ones(2, 128, dtype=torch.bool)
+input = make_padded_input(x, mask)
+
+res = mlm_wrapper(input, encoder)
+
+task_loss = my_task_head(res.out)  # res.out is always the clean, unmasked encoder output
+loss = task_loss + 0.1 * res.mlm_loss  # zero automatically in eval — nothing to branch on
+```
+
+One call site covers both modes: in training, `mlm_wrapper` runs the encoder twice internally — once clean (returned as `res.out`) and once on a separately-masked copy (used only to compute `res.mlm_loss`) — so masking never reaches whatever consumes `res.out`. In eval, it runs the encoder once and `res.mlm_loss` is a constant zero. Pretraining with no main task yet — just drop `task_loss` and train on `res.mlm_loss` alone.
+
+`mlm_wrapper` and `encoder` are separate modules — `encoder` is passed in at call time, never stored, so `mlm_wrapper.parameters()` holds only `mask_token`, `masking_strategy`, and `head`. Build one optimizer over both: `itertools.chain(encoder.parameters(), mlm_wrapper.parameters())`. And since they're separate modules, `.train()`/`.eval()` on one does not propagate to the other — call both, or register both as submodules of a shared parent so one call does.
+
+## Design
+
+**`out` is always clean; only `mlm_loss` ever reflects masking.** `forward(input, encoder)` runs `encoder(input)` unconditionally for `out`. In training it additionally runs `encoder` a second time on a masked copy, purely to produce `mlm_loss` — that corrupted copy is never returned. This means the main pipeline sees byte-identical output whether or not the aux loss is being trained alongside it, and callers never need to special-case which pass produced `out`.
+
+**Training/eval is the one branch `forward()` is allowed, because it isn't tensor control flow.** `self.training` gates whether the second (masked) pass runs — the same mechanism `nn.Dropout` and `nn.BatchNorm` already use. It lets callers invoke `mlm_wrapper(input, encoder)` unconditionally in both modes, instead of an external `if training: ...` at every call site.
+
+**A wrapper over an encoder supplied at call time, not one it owns.** `MLMWrapper` never stores an encoder — `forward` takes one as an argument, so nothing is registered as a submodule and `mlm_wrapper.parameters()` never includes the encoder's weights.
+
+**The mask token is owned state, not an injected collaborator.** It's the wrapper's own learned parameter, analogous to a `nn.Linear`'s weight — not swappable behavior, so it isn't a `Protocol`.
+
+**The reconstruction target is always `input.x.detach()`, unconditionally.** Not a config flag: computing a detached target is the only behavior the masked pass ever has, so there is nothing to branch on. This severs the gradient path from the reconstruction loss back to whatever produced `input.x`, which removes the standard collapse shortcut for self-supervised regression targets (drive every token toward one constant vector, and reconstruction becomes trivial). See `tests/mlm/test_wrapper.py` for the gradient test that verifies this directly.
+
+**No layout dispatch in the wrapper.** `MaskingStrategy` returns a boolean tensor shaped like `input.x`'s leading dims — `(b, n)` for `PaddedInput`, `(nt,)` for `PackedInput`. `torch.where` and boolean advanced indexing both broadcast/gather correctly against either shape, so `MLMWrapper.forward` never matches on the sequence variant; only `RandomMasking` does, internally.
+
+**Random masking needs no packed-sequence boundary awareness.** Each token's masking decision is independent of every other token's, so document identity (`cu_seqlens`) never enters the decision. Boundary-awareness only becomes necessary for a contiguous-span (blockwise) strategy, which could otherwise straddle two packed documents — add that logic when that strategy is added, not before.
+
+## Extending
+
+To add a masking strategy (e.g. blockwise), write an `nn.Module` satisfying `MaskingStrategy` and pass it as `MLMWrapper(..., masking_strategy=...)`.
+
+To add a reconstruction target type (e.g. a discretized/codebook target for BEiT-style cross-entropy, or a contrastive target), write an `nn.Module` satisfying `ReconstructionHead` and pass it as `MLMWrapper(..., head=...)`. Prediction and loss are scored together by one call — they're a coupled choice (a cross-entropy loss over a raw regression output isn't meaningful), not two independently pluggable axes.

stackformers-4.3.0/stackformers/mlm/config.py

@@ -0,0 +1,10 @@
+from __future__ import annotations
+
+from pydantic import BaseModel, Field
+
+
+class MLMWrapperConfig(BaseModel):
+    """Config for MLMWrapper: token dimension and corruption ratio."""
+
+    dim: int = Field(gt=0)
+    mask_ratio: float = Field(default=0.15, gt=0.0, lt=1.0)

stackformers-4.3.0/stackformers/mlm/head.py

@@ -0,0 +1,27 @@
+from __future__ import annotations
+
+import torch.nn as nn
+import torch.nn.functional as F
+from jaxtyping import Float
+from torch import Tensor
+
+
+class RegressionHead(nn.Module):
+    """Reconstruct masked tokens via a linear projection, scored with MSE.
+
+    The projection gives the encoder a dedicated place to specialise for
+    reconstruction, so the shared representation isn't forced to double as the
+    literal token vector at every layer.
+    """
+
+    def __init__(self, dim: int) -> None:
+        super().__init__()
+        self.proj = nn.Linear(dim, dim)
+
+    def forward(
+        self,
+        encoder_output_at_masked: Float[Tensor, "m d"],
+        target_at_masked: Float[Tensor, "m d"],
+    ) -> Float[Tensor, ""]:
+        prediction = self.proj(encoder_output_at_masked)
+        return F.mse_loss(prediction, target_at_masked)

stackformers-4.3.0/stackformers/mlm/masking.py

@@ -0,0 +1,30 @@
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+from jaxtyping import Bool
+from torch import Tensor
+
+from stackformers.sequence import PackedInput, PaddedInput, SequenceInput
+
+
+class RandomMasking(nn.Module):
+    """Mask each valid token independently with probability mask_ratio.
+
+    Every token's masking decision is drawn independently, so packed sequences need
+    no document-boundary awareness here — that only becomes necessary for a
+    contiguous-span (blockwise) strategy, which could otherwise straddle two packed
+    documents.
+    """
+
+    def __init__(self, mask_ratio: float = 0.15) -> None:
+        super().__init__()
+        self.mask_ratio = mask_ratio
+
+    def forward(self, input: SequenceInput) -> Bool[Tensor, "*batch"]:
+        match input:
+            case PaddedInput(mask=mask):
+                candidate = torch.rand(mask.shape, device=mask.device) < self.mask_ratio
+                return candidate & mask
+            case PackedInput(x=x):
+                return torch.rand(x.shape[0], device=x.device) < self.mask_ratio

stackformers-4.3.0/stackformers/mlm/protocols.py

@@ -0,0 +1,47 @@
+from __future__ import annotations
+
+from typing import Protocol, runtime_checkable
+
+from jaxtyping import Bool, Float
+from torch import Tensor
+
+from stackformers.sequence import SequenceInput
+
+
+@runtime_checkable
+class EncoderLike(Protocol):
+    """Any encoder mapping a sequence to per-token embeddings of the same layout.
+
+    Implementations: Encoder, or any TransformerEncoderBase subclass.
+    """
+
+    def __call__(self, input: SequenceInput) -> Tensor: ...
+
+
+@runtime_checkable
+class MaskingStrategy(Protocol):
+    """Select which token positions to corrupt for reconstruction.
+
+    Returns True at positions to mask, shaped like the leading (non-feature) dims of
+    input.x — (b, n) for PaddedInput, (nt,) for PackedInput. Implementations dispatch
+    on the SequenceInput variant internally, so callers never need to know the layout.
+    Implementation: RandomMasking.
+    """
+
+    def __call__(self, input: SequenceInput) -> Bool[Tensor, "*batch"]: ...
+
+
+@runtime_checkable
+class ReconstructionHead(Protocol):
+    """Predict masked tokens from encoder output and score against the clean target.
+
+    Both arguments are already gathered down to just the masked positions (m = number
+    of masked tokens across the batch or pack). Returns a scalar loss.
+    Implementation: RegressionHead.
+    """
+
+    def __call__(
+        self,
+        encoder_output_at_masked: Float[Tensor, "m d"],
+        target_at_masked: Float[Tensor, "m d"],
+    ) -> Float[Tensor, ""]: ...

stackformers-4.3.0/stackformers/mlm/wrapper.py

@@ -0,0 +1,76 @@
+from __future__ import annotations
+
+from typing import NamedTuple
+
+import torch
+import torch.nn as nn
+from jaxtyping import Float
+from torch import Tensor
+
+from stackformers.mlm.config import MLMWrapperConfig
+from stackformers.mlm.head import RegressionHead
+from stackformers.mlm.masking import RandomMasking
+from stackformers.mlm.protocols import EncoderLike, MaskingStrategy, ReconstructionHead
+from stackformers.sequence import SequenceInput
+
+
+class MLMOutput(NamedTuple):
+    out: Float[Tensor, "*batch d"]
+    mlm_loss: Float[Tensor, ""]
+
+
+class MLMWrapper(nn.Module):
+    """Masked-token-reconstruction auxiliary loss over an encoder supplied at call time.
+
+    Does not store or own an encoder — none is passed to __init__, and forward() takes
+    one as an argument instead of holding a reference across calls. This module's own
+    parameter tree (mask_token, masking_strategy, head) is exactly what it owns; the
+    encoder stays wherever the caller already keeps it, with no duplicate registration
+    under this wrapper.
+
+    `out` is always the encoder's clean, unmasked output — masking is invisible to
+    whatever consumes it, so the main pipeline behaves identically whether or not this
+    aux loss is being trained alongside it. Only `mlm_loss` ever reflects that masking
+    happened. In training mode that costs a second encoder forward pass on a
+    separately-masked copy of input; the two passes share weights but never interfere,
+    since the corrupted copy built for the loss is never returned as `out`. In eval mode
+    only the one clean pass runs and `mlm_loss` reports a constant zero, so callers can
+    invoke this unconditionally in both modes without an if-training branch of their
+    own — the same role self.training already plays in nn.Dropout or nn.BatchNorm.
+
+    The reconstruction target is always the encoder's own clean input embedding,
+    detached — so this loss trains the encoder and head, never whatever produced
+    input.x. That removes the representation-collapse shortcut a trainable tokenizer
+    would otherwise have available (drive every token toward a constant vector to make
+    reconstruction trivial).
+    """
+
+    def __init__(
+        self,
+        config: MLMWrapperConfig,
+        masking_strategy: MaskingStrategy | None = None,
+        head: ReconstructionHead | None = None,
+    ) -> None:
+        super().__init__()
+        self.mask_token = nn.Parameter(torch.empty(config.dim))
+        nn.init.trunc_normal_(self.mask_token, std=0.02, a=-0.04, b=0.04)
+        self.masking_strategy = (
+            masking_strategy if masking_strategy is not None else RandomMasking(config.mask_ratio)
+        )
+        self.head = head if head is not None else RegressionHead(config.dim)
+
+    def _masked_loss(self, input: SequenceInput, encoder: EncoderLike) -> Tensor:
+        should_mask = self.masking_strategy(input)
+        target = input.x.detach()
+        # should_mask's shape always matches input.x's leading dims, so this select and
+        # the boolean indexing below both work for PaddedInput and PackedInput alike.
+        corrupted_x = torch.where(should_mask.unsqueeze(-1), self.mask_token, input.x)
+        masked_output = encoder(input._replace(x=corrupted_x))
+        return self.head(masked_output[should_mask], target[should_mask])
+
+    def forward(self, input: SequenceInput, encoder: EncoderLike) -> MLMOutput:
+        clean_output = encoder(input)
+        if self.training:
+            mlm_loss = self._masked_loss(input, encoder)
+            return MLMOutput(out=clean_output, mlm_loss=mlm_loss)
+        return MLMOutput(out=clean_output, mlm_loss=clean_output.new_zeros(()))

{stackformers-4.2.0 → stackformers-4.3.0}/stackformers/positional/README.md

@@ -8,12 +8,14 @@ Positional encodings applied to Q and K tensors inside attention — not to the
 
 RoPE's shortest wavelength is fixed at 2π coordinate units, and `base` only stretches the long end of the band range. So `base` lands the ladder correctly only when tokens sit *one unit apart* — true for text, true for a patch grid, meaningless for scattered continuous coordinates. Positions normalised to `[0, 1]` leave every band longer than the whole domain, so the encoding barely rotates and carries almost no signal; and `base=10000` on a 14×14 patch grid leaves 14 of 16 bands frozen.
 
-`RotaryEmbeddingND` (`RoPENDConfig`) drops `base` entirely and builds the ladder from the two numbers that actually mean something, in any number of dimensions:
+`RotaryEmbeddingND` (`RoPENDConfig`) drops `base` entirely and builds the ladder from the two numbers that actually mean something, in any number of dimensions. Both ends are the same rule — **a half turn over the scale in question, `ω = π / scale`**:
 
-- `r_min` — the finest separation the model must resolve. Sets the shortest wavelength to `2·r_min`, the Nyquist limit. Measure it as a low percentile of the nearest-neighbour distance.
-- `r_max` — the domain diameter. Sets the longest wavelength (times `headroom`, so the slowest band stays monotone rather than wrapping). Measure it as a high percentile of the pairwise distance distribution.
+- `r_min` — the finest separation the model must resolve. Turned through π on the fastest band, so the shortest wavelength is `2·r_min`: the Nyquist limit, past which distinct offsets alias onto the same rotation. Measure it as a low percentile of the nearest-neighbour distance.
+- `r_max` — the domain diameter. Attention sees *signed* offsets, spanning `[−r_max, +r_max]` — a width of `2·r_max`, not `r_max` — so the slowest band is given a half turn over that full width. It cannot wrap, which leaves it monotone in the offset: a coarse ramp the model reads as displacement. Measure it as a high percentile of the pairwise distance distribution.
 
-The resulting ladder depends only on the *dynamic range* `r_max / r_min`, so metres, pixels and millimetres all give the identical encoding — the property `base` never had. `RoPE1DConfig` and `RoPE2DConfig` keep `base` and are unchanged; use them for text and grids.
+There is no third knob, and no constant to tune: two measured distances fix the ladder. The resulting band range depends only on the *dynamic range* `r_max / r_min`, so metres, pixels and millimetres all give the identical encoding — the property `base` never had. `RoPE1DConfig` and `RoPE2DConfig` keep `base` and are unchanged; use them for text and grids.
+
+`dim_head` must be at least `4 · coords`, so that every axis gets at least two bands. At one band per axis the ladder has nowhere to descend: it collapses onto the fast end, `r_max` is never reached, and the encoding silently becomes periodic with period `2·r_min` across the whole domain. `RoPENDConfig` rejects that at construction.
 
 **Centre your coordinates.** Not for invariance — RoPE is translation-invariant exactly, since the query and key rotations cancel into `ω·(pᵢ − pⱼ)`. It is for float32: the angle is `ω·p`, and coordinates far from the origin push it into the hundreds of radians and spend the mantissa before the cosine is taken. A 1e5 offset costs three orders of magnitude of accuracy. This cannot be done inside the module — cross-attention must subtract the *same* constant from query and key positions, and the module sees them separately — so subtract the centroid at the input boundary.
 

{stackformers-4.2.0 → stackformers-4.3.0}/stackformers/positional/config.py

@@ -53,9 +53,10 @@ class RoPENDConfig(BaseModel):
     Takes no ``base``. RoPE's ``base`` ladder pins its fastest band at a wavelength of 2π
     *coordinate units* regardless of the base, so it only lands correctly when tokens sit one
     unit apart — true for text and patch grids, meaningless for continuous coordinates. The
-    band range is set by the data instead: r_min fixes the fast end, r_max the slow end. The
-    ladder then depends only on the ratio r_max / r_min, so the units the coordinates happen
-    to be expressed in stop mattering.
+    band range is set by the data instead: r_min fixes the fast end, r_max the slow end, and
+    each is a half turn over the scale it names (ω = π / scale). Nothing else is free. The
+    ladder's shape then depends only on the ratio r_max / r_min, so the units the coordinates
+    happen to be expressed in stop mattering.
     """
 
     kind: Literal["rope_nd"] = "rope_nd"
@@ -66,29 +67,22 @@ class RoPENDConfig(BaseModel):
     r_min: float = Field(
         gt=0.0,
         description=(
-            "Finest separation between two nodes that the model must tell apart. Sets the"
-            " shortest wavelength to 2*r_min — the Nyquist limit, below which distinct offsets"
-            " alias onto the same rotation. Measure it as a low percentile of the"
-            " nearest-neighbour distance, not the minimum, which is noise."
+            "Finest separation between two nodes that the model must tell apart. Gets a half"
+            " turn on the fastest band, so the shortest wavelength is 2*r_min — the Nyquist"
+            " limit, below which distinct offsets alias onto the same rotation. Measure it as a"
+            " low percentile of the nearest-neighbour distance, not the minimum, which is noise."
         ),
     )
     r_max: float = Field(
         gt=0.0,
         description=(
-            "Diameter of the domain — the largest offset that must stay distinguishable. Sets"
-            " the longest wavelength. Measure it as a high percentile of the pairwise distance"
+            "Diameter of the domain — the largest offset that must stay distinguishable."
+            " Attention sees signed offsets spanning [-r_max, +r_max], so the slowest band is"
+            " given a half turn over that full width of 2*r_max: it never wraps, which leaves it"
+            " monotone in the offset. Measure it as a high percentile of the pairwise distance"
             " distribution, not the maximum, which is an outlier."
         ),
     )
-    headroom: float = Field(
-        default=4.0,
-        gt=1.0,
-        description=(
-            "How far the longest wavelength reaches past r_max. Keeps the slowest band monotone"
-            " across the whole domain instead of wrapping, so it can act as a coarse absolute"
-            " coordinate. Llama's defaults sit at roughly 13x; 2-4x is ample here."
-        ),
-    )
 
     @model_validator(mode="after")
     def _check(self) -> RoPENDConfig:
@@ -99,6 +93,14 @@ class RoPENDConfig(BaseModel):
                 f" dim_head is split into {self.coords} per-axis blocks, and each block is"
                 " rotated in pairs of channels."
             )
+        if self.bands_per_axis < 2:
+            raise ValueError(
+                f"dim_head ({self.dim_head}) leaves {self.bands_per_axis} band per axis; at"
+                f" least 2 bands are needed, so dim_head must be at least {4 * self.coords}."
+                " A lone band lands on the fast end of the ladder and r_max is never reached,"
+                " leaving the encoding periodic with period 2 * r_min across the whole domain:"
+                " every offset an integer number of periods apart becomes indistinguishable."
+            )
         if self.r_max <= self.r_min:
             raise ValueError(f"r_max ({self.r_max}) must exceed r_min ({self.r_min}).")
         return self

{stackformers-4.2.0 → stackformers-4.3.0}/stackformers/positional/rope_nd.py

@@ -18,18 +18,28 @@ def _frequency_ladder(config: RoPENDConfig) -> Tensor:
     at ω=1, i.e. a wavelength of exactly 2π coordinate units, whatever the base — so ``base``
     can only stretch the slow end, and the ladder lands correctly only when tokens happen to
     sit one unit apart. That holds for text and for patch grids; it is meaningless for
-    continuous coordinates, where the useful band range is fixed by the data instead:
+    continuous coordinates, where the band range is fixed by the data instead.
 
-        ω_hi = π / r_min                  → shortest wavelength = 2·r_min  (Nyquist on the
-                                            finest separation the model must resolve)
-        ω_lo = 2π / (headroom · r_max)    → longest wavelength spans the whole domain
+    Both ends are a half turn over the scale they name — ω = π / scale:
 
-    So the ladder is a function of the *dynamic range* r_max / r_min alone, and is invariant
-    to the units the coordinates are expressed in — metres, pixels, or normalised to [0, 1]
-    all give the same encoding. That is precisely the property ``base`` does not have.
+        ω_hi = π / r_min       half-period r_min. The finest separation the model must resolve
+                               is turned through π, the furthest two points can be driven apart
+                               on a circle. Any faster is past Nyquist, and distinct offsets
+                               start aliasing onto the same rotation.
+
+        ω_lo = π / (2·r_max)   half-period 2·r_max. Attention sees *signed* offsets, spanning
+                               [−r_max, +r_max] — a width of 2·r_max, not r_max — so the slowest
+                               band turns through at most π across the entire domain. It cannot
+                               wrap, which leaves it monotone in the offset: a coarse ramp the
+                               model can read directly as displacement.
+
+    Nothing here is tuned, and there is no knob left to set: two measured distances fix the
+    ladder. Its *shape* depends only on the dynamic range r_max / r_min, so the encoding is
+    invariant to the units the coordinates are expressed in — metres, pixels, or normalised to
+    [0, 1] all give the same result. That is precisely the property ``base`` does not have.
     """
     omega_hi = math.pi / config.r_min
-    omega_lo = 2.0 * math.pi / (config.headroom * config.r_max)
+    omega_lo = math.pi / (2.0 * config.r_max)
     steps = torch.linspace(math.log(omega_hi), math.log(omega_lo), config.bands_per_axis)
     return torch.exp(steps)  # fastest band first, matching the 1-D/2-D convention
 

stackformers-4.3.0/tests/mlm/test_head.py

@@ -0,0 +1,37 @@
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+from stackformers.mlm.head import RegressionHead
+
+M, D = 12, 16
+
+
+def test_regression_head_output_is_scalar(device_dtype: tuple[torch.device, torch.dtype]) -> None:
+    device, dtype = device_dtype
+    head = RegressionHead(dim=D).to(device=device, dtype=dtype)
+    prediction_in = torch.randn(M, D, device=device, dtype=dtype)
+    target = torch.randn(M, D, device=device, dtype=dtype)
+    loss = head(prediction_in, target)
+    assert loss.shape == ()
+
+
+def test_regression_head_zero_when_projection_matches_target(device: torch.device) -> None:
+    head = RegressionHead(dim=D).to(device)
+    nn.init.zeros_(head.proj.weight)
+    nn.init.zeros_(head.proj.bias)
+    target = torch.zeros(M, D, device=device)
+    prediction_in = torch.randn(M, D, device=device)  # proj(x) == 0 regardless, weight/bias are 0
+    loss = head(prediction_in, target)
+    assert torch.allclose(loss, torch.zeros_like(loss))
+
+
+def test_regression_head_gradients_flow(device: torch.device) -> None:
+    head = RegressionHead(dim=D).to(device)
+    prediction_in = torch.randn(M, D, device=device, requires_grad=True)
+    target = torch.randn(M, D, device=device)
+    loss = head(prediction_in, target)
+    loss.backward()
+    assert prediction_in.grad is not None
+    assert head.proj.weight.grad is not None

stackformers-4.3.0/tests/mlm/test_masking.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+import torch
+
+from stackformers.mlm.masking import RandomMasking
+from stackformers.sequence import make_packed_input, make_padded_input
+
+B, N, D = 4, 64, 16
+NT = 20
+
+
+def test_random_masking_padded_shape(device_dtype: tuple[torch.device, torch.dtype]) -> None:
+    device, dtype = device_dtype
+    x = torch.randn(B, N, D, device=device, dtype=dtype)
+    mask = torch.ones(B, N, dtype=torch.bool, device=device)
+    input = make_padded_input(x, mask)
+    should_mask = RandomMasking(mask_ratio=0.5)(input)
+    assert should_mask.shape == mask.shape
+    assert should_mask.dtype == torch.bool
+
+
+def test_random_masking_packed_shape(device_dtype: tuple[torch.device, torch.dtype]) -> None:
+    device, dtype = device_dtype
+    x = torch.randn(NT, D, device=device, dtype=dtype)
+    cu = torch.tensor([0, 12, 20], dtype=torch.int32, device=device)
+    input = make_packed_input(x, cu, max_seqlen=12)
+    should_mask = RandomMasking(mask_ratio=0.5)(input)
+    assert should_mask.shape == (NT,)
+    assert should_mask.dtype == torch.bool
+
+
+def test_random_masking_never_masks_padding(device_dtype: tuple[torch.device, torch.dtype]) -> None:
+    device, dtype = device_dtype
+    x = torch.randn(B, N, D, device=device, dtype=dtype)
+    mask = torch.ones(B, N, dtype=torch.bool, device=device)
+    mask[:, N // 2 :] = False  # second half is padding
+    input = make_padded_input(x, mask)
+    should_mask = RandomMasking(mask_ratio=0.9)(input)
+    assert not (should_mask & ~mask).any()
+
+
+def test_random_masking_respects_ratio_roughly(
+    device_dtype: tuple[torch.device, torch.dtype],
+) -> None:
+    """Statistical check over a large sample — not exact, but far enough from any other ratio."""
+    device, dtype = device_dtype
+    torch.manual_seed(0)
+    x = torch.randn(1, 10_000, D, device=device, dtype=dtype)
+    mask = torch.ones(1, 10_000, dtype=torch.bool, device=device)
+    input = make_padded_input(x, mask)
+    should_mask = RandomMasking(mask_ratio=0.3)(input)
+    fraction = should_mask.float().mean().item()
+    assert 0.25 < fraction < 0.35

stackformers-4.3.0/tests/mlm/test_wrapper.py

@@ -0,0 +1,201 @@
+from __future__ import annotations
+
+import pytest
+import torch
+import torch.nn as nn
+
+from stackformers.attention.config import SelfAttentionConfig
+from stackformers.attention.self_attn import SelfAttention
+from stackformers.encoder import Encoder
+from stackformers.feedforward.config import SwiGLUConfig
+from stackformers.feedforward.swiglu import SwiGLU
+from stackformers.layers import TransformerLayer
+from stackformers.mlm.config import MLMWrapperConfig
+from stackformers.mlm.wrapper import MLMOutput, MLMWrapper
+from stackformers.norm.config import RMSNormConfig
+from stackformers.norm.factory import build_norm
+from stackformers.positional.none import NoPosEncoding
+from stackformers.sequence import (
+    PackedInput,
+    PaddedInput,
+    SequenceInput,
+    make_packed_input,
+    make_padded_input,
+)
+
+B, N, D, H, DH = 2, 16, 16, 4, 4
+NUM_LAYERS = 2
+NT = 10  # two packed seqs: 6 + 4
+
+
+class AllMasking(nn.Module):
+    """Test double: marks every valid position for masking."""
+
+    def forward(self, input: SequenceInput) -> torch.Tensor:
+        match input:
+            case PaddedInput(mask=mask):
+                return mask.clone()
+            case PackedInput(x=x):
+                return torch.ones(x.shape[0], dtype=torch.bool, device=x.device)
+
+
+def _build_encoder(device: torch.device, dtype: torch.dtype) -> Encoder:
+    attn_cfg = SelfAttentionConfig(dim=D, heads=H, dim_head=DH)
+    ff_cfg = SwiGLUConfig(dim=D)
+    norm_cfg = RMSNormConfig(dim=D)
+    layers = [
+        TransformerLayer(
+            self_attn=SelfAttention(attn_cfg, NoPosEncoding()),
+            ff=SwiGLU(ff_cfg),
+            norm_attn=build_norm(norm_cfg),
+            norm_ff=build_norm(norm_cfg),
+        )
+        for _ in range(NUM_LAYERS)
+    ]
+    return Encoder(layers=layers, final_norm=build_norm(norm_cfg)).to(device=device, dtype=dtype)
+
+
+@pytest.fixture
+def config() -> MLMWrapperConfig:
+    return MLMWrapperConfig(dim=D, mask_ratio=0.5)
+
+
+@pytest.fixture
+def x_pad(device_dtype: tuple[torch.device, torch.dtype]) -> PaddedInput:
+    device, dtype = device_dtype
+    x = torch.randn(B, N, D, device=device, dtype=dtype)
+    mask = torch.ones(B, N, dtype=torch.bool, device=device)
+    return make_padded_input(x, mask)
+
+
+@pytest.fixture
+def x_packed(device_dtype: tuple[torch.device, torch.dtype]) -> PackedInput:
+    device, dtype = device_dtype
+    if not device.type == "cuda" or dtype not in (torch.float16, torch.bfloat16):
+        pytest.skip("packed attention requires CUDA with float16 or bfloat16")
+    x = torch.randn(NT, D, device=device, dtype=dtype)
+    cu = torch.tensor([0, 6, 10], dtype=torch.int32, device=device)
+    return make_packed_input(x, cu, max_seqlen=6)
+
+
+def test_wrapper_padded_output_shapes(
+    config: MLMWrapperConfig,
+    x_pad: PaddedInput,
+    device_dtype: tuple[torch.device, torch.dtype],
+) -> None:
+    device, dtype = device_dtype
+    encoder = _build_encoder(device, dtype)
+    wrapper = MLMWrapper(config).to(device=device, dtype=dtype)
+    res = wrapper(x_pad, encoder)
+    assert isinstance(res, MLMOutput)
+    assert res.mlm_loss.shape == ()
+    assert res.out.shape == (B, N, D)
+
+
+def test_wrapper_packed_output_shapes(
+    config: MLMWrapperConfig,
+    x_packed: PackedInput,
+    device_dtype: tuple[torch.device, torch.dtype],
+) -> None:
+    device, dtype = device_dtype
+    encoder = _build_encoder(device, dtype)
+    wrapper = MLMWrapper(config).to(device=device, dtype=dtype)
+    res = wrapper(x_packed, encoder)
+    assert res.mlm_loss.shape == ()
+    assert res.out.shape == (NT, D)
+
+
+def test_wrapper_mask_token_receives_gradient(device: torch.device) -> None:
+    config = MLMWrapperConfig(dim=D, mask_ratio=0.5)
+    encoder = _build_encoder(device, torch.float32)
+    wrapper = MLMWrapper(config).to(device)
+    x = torch.randn(B, N, D, device=device)
+    mask = torch.ones(B, N, dtype=torch.bool, device=device)
+    res = wrapper(make_padded_input(x, mask), encoder)
+    res.mlm_loss.backward()
+    assert wrapper.mask_token.grad is not None
+
+
+def test_wrapper_does_not_own_encoder(device: torch.device) -> None:
+    """MLMWrapper takes the encoder at call time, not at construction — its own
+    parameter tree must never include the encoder's weights, since the encoder is
+    owned and registered wherever the caller already keeps it.
+    """
+    config = MLMWrapperConfig(dim=D, mask_ratio=0.5)
+    encoder = _build_encoder(device, torch.float32)
+    wrapper = MLMWrapper(config).to(device)
+    encoder_param_ids = {id(p) for p in encoder.parameters()}
+    wrapper_param_ids = {id(p) for p in wrapper.parameters()}
+    assert encoder_param_ids.isdisjoint(wrapper_param_ids)
+
+
+def test_wrapper_detaches_target_from_upstream_input(device: torch.device) -> None:
+    """Collapse guard (design doc §5): a trainable tokenizer upstream of x must get no
+    gradient from this loss, or the optimizer can collapse every token to one constant
+    vector to make reconstruction trivial.
+
+    Force every position masked, so torch.where's backward zeroes x's contribution
+    through the corrupted-x branch entirely (its condition is True everywhere, so the
+    "x" branch of the select is never taken). The only remaining path from loss back to
+    x would be through the target — and that's detached — so x's accumulated gradient
+    must be exactly zero.
+    """
+    config = MLMWrapperConfig(dim=D, mask_ratio=0.5)
+    encoder = _build_encoder(device, torch.float32)
+    wrapper = MLMWrapper(config, masking_strategy=AllMasking()).to(device)
+    x = torch.randn(B, N, D, device=device, requires_grad=True)
+    mask = torch.ones(B, N, dtype=torch.bool, device=device)
+    res = wrapper(make_padded_input(x, mask), encoder)
+    res.mlm_loss.backward()
+    assert x.grad is not None
+    assert torch.all(x.grad == 0)
+
+
+def test_wrapper_accepts_custom_masking_strategy(device: torch.device) -> None:
+    config = MLMWrapperConfig(dim=D, mask_ratio=0.5)
+    encoder = _build_encoder(device, torch.float32)
+    wrapper = MLMWrapper(config, masking_strategy=AllMasking()).to(device)
+    x = torch.randn(B, N, D, device=device)
+    mask = torch.ones(B, N, dtype=torch.bool, device=device)
+    res = wrapper(make_padded_input(x, mask), encoder)
+    assert res.mlm_loss.shape == ()
+
+
+def test_wrapper_training_out_is_always_clean(device: torch.device) -> None:
+    """out must never carry masking, even in training — only mlm_loss does. This keeps
+    the main pipeline byte-identical whether or not the MLM aux loss trains alongside
+    it, so callers never need to special-case which pass produced `out`. AllMasking
+    maximizes the contrast: if masking leaked into `out`, every token would be replaced
+    by mask_token, making it maximally different from a clean encoder(input) call.
+    """
+    config = MLMWrapperConfig(dim=D, mask_ratio=0.5)
+    encoder = _build_encoder(device, torch.float32)
+    wrapper = MLMWrapper(config, masking_strategy=AllMasking()).to(device)
+    x = torch.randn(B, N, D, device=device)
+    mask = torch.ones(B, N, dtype=torch.bool, device=device)
+    input = make_padded_input(x, mask)
+
+    res = wrapper(input, encoder)
+    expected = encoder(input)
+
+    assert torch.equal(res.out, expected)
+
+
+def test_wrapper_eval_mode_passes_through_encoder_unmasked(device: torch.device) -> None:
+    """In eval mode, forward() must skip masking entirely: out matches a clean encoder
+    call exactly, and mlm_loss reports zero — so callers can invoke the wrapper
+    unconditionally in both modes without an if-training branch of their own.
+    """
+    config = MLMWrapperConfig(dim=D, mask_ratio=0.5)
+    encoder = _build_encoder(device, torch.float32)
+    wrapper = MLMWrapper(config).to(device)
+    wrapper.eval()
+    x = torch.randn(B, N, D, device=device)
+    mask = torch.ones(B, N, dtype=torch.bool, device=device)
+    input = make_padded_input(x, mask)
+
+    res = wrapper(input, encoder)
+    expected = encoder(input)
+
+    assert torch.equal(res.out, expected)
+    assert torch.equal(res.mlm_loss, torch.zeros_like(res.mlm_loss))

{stackformers-4.2.0 → stackformers-4.3.0}/tests/positional/test_rope_nd.py

@@ -11,7 +11,7 @@ from stackformers.positional.rope_nd import RotaryEmbeddingND
 from tests.conftest import atol
 
 B, N, H = 2, 16, 4
-R_MIN, R_MAX, HEADROOM = 0.5, 100.0, 4.0
+R_MIN, R_MAX = 0.5, 100.0
 
 DIMS = [(1, 64), (2, 64), (3, 96)]  # (coords, dim_head) — dim_head must divide by 2 * coords
 
@@ -19,9 +19,7 @@ DIMS = [(1, 64), (2, 64), (3, 96)]  # (coords, dim_head) — dim_head must divid
 @pytest.fixture(params=DIMS, ids=lambda p: f"c{p[0]}")
 def config(request: pytest.FixtureRequest) -> RoPENDConfig:
     coords, dim_head = request.param
-    return RoPENDConfig(
-        dim_head=dim_head, coords=coords, r_min=R_MIN, r_max=R_MAX, headroom=HEADROOM
-    )
+    return RoPENDConfig(dim_head=dim_head, coords=coords, r_min=R_MIN, r_max=R_MAX)
 
 
 @pytest.fixture
@@ -65,15 +63,32 @@ def test_output_shape(
 
 
 def test_ladder_spans_nyquist_to_domain(config: RoPENDConfig) -> None:
-    """The whole point of dropping `base`: the band range is pinned to the data, not a constant."""
+    """The whole point of dropping `base`: the band range is pinned to the data, not a constant.
+
+    Both ends are a half turn over the scale they name — r_min at the fast end, and the full
+    signed offset width 2*r_max at the slow end — so both are stated here as half-periods.
+    """
     inv_freq: torch.Tensor = RotaryEmbeddingND(config).inv_freq  # type: ignore[assignment]
-    shortest = 2 * math.pi / float(inv_freq.max())
-    longest = 2 * math.pi / float(inv_freq.min())
-    assert shortest == pytest.approx(2 * R_MIN, rel=1e-4)  # Nyquist on the finest separation
-    assert longest == pytest.approx(HEADROOM * R_MAX, rel=1e-4)  # spans the domain, with headroom
+    fastest_half_period = math.pi / float(inv_freq.max())
+    slowest_half_period = math.pi / float(inv_freq.min())
+    assert fastest_half_period == pytest.approx(R_MIN, rel=1e-4)  # Nyquist on the finest gap
+    assert slowest_half_period == pytest.approx(2 * R_MAX, rel=1e-4)  # the whole signed range
     assert inv_freq.shape == (config.bands_per_axis,)
 
 
+def test_slowest_band_never_wraps(config: RoPENDConfig) -> None:
+    """The slow end must stay monotone over every offset attention can see.
+
+    Attention sees signed offsets over [-r_max, +r_max], so the slowest band sweeps
+    omega_lo * 2 * r_max. Let that exceed pi and the band turns back on itself: two different
+    offsets in the domain collide on the same rotation, in the one band whose entire job is to
+    tell the coarse end apart.
+    """
+    inv_freq: torch.Tensor = RotaryEmbeddingND(config).inv_freq  # type: ignore[assignment]
+    sweep = float(inv_freq.min()) * 2 * R_MAX
+    assert sweep == pytest.approx(math.pi, rel=1e-4)
+
+
 def test_units_do_not_matter(config: RoPENDConfig, device: torch.device) -> None:
     """Scale the coordinates and r_min/r_max together — the encoding must not move.
 
@@ -209,6 +224,17 @@ def test_rejects_indivisible_head_dim() -> None:
         RoPENDConfig(dim_head=64, coords=3, r_min=R_MIN, r_max=R_MAX)
 
 
+def test_rejects_single_band_per_axis() -> None:
+    """dim_head == 2 * coords divides cleanly, but leaves a ladder with nowhere to descend.
+
+    The lone band lands on omega_hi and r_max is never reached, so the encoding is periodic with
+    period 2 * r_min across the entire domain — offsets of 0, 2*r_min, 4*r_min ... all give the
+    same rotation. It fails silently, which is why it is rejected at construction.
+    """
+    with pytest.raises(ValidationError, match="at least 2 bands"):
+        RoPENDConfig(dim_head=6, coords=3, r_min=R_MIN, r_max=R_MAX)
+
+
 def test_rejects_inverted_range() -> None:
     with pytest.raises(ValidationError, match="must exceed r_min"):
         RoPENDConfig(dim_head=64, coords=2, r_min=100.0, r_max=1.0)

{stackformers-4.2.0 → stackformers-4.3.0}/uv.lock

@@ -811,7 +811,7 @@ wheels = [
 
 [[package]]
 name = "stackformers"
-version = "4.2.0"
+version = "4.3.0"
 source = { editable = "." }
 dependencies = [
     { name = "beartype" },