diff --git a/docs/plans/2026-03-01-phase4-hiera-backbone-plan.md b/docs/plans/2026-03-01-phase4-hiera-backbone-plan.md new file mode 100644 index 0000000..67c9b42 --- /dev/null +++ b/docs/plans/2026-03-01-phase4-hiera-backbone-plan.md @@ -0,0 +1,120 @@ +# Phase 4: Hiera Backbone MLX Port + +## Context + +Phase 4 of the CorridorKey MLX port. Phases 1-3 complete (reference harness, decoder/refiner, converter). Now need the Hiera backbone — the most complex component. Once done, all model pieces exist for end-to-end inference. + +## Architecture Summary + +Hiera is a hierarchical vision transformer from timm. Key components: + +- **PatchEmbed**: Conv2d(4->112, 7x7, stride=4) + reshape to [B, N, C] +- **pos_embed**: Learned (1, N, 112), bicubic interpolated from training res +- **Unroll**: Permutes tokens for mask-unit windowed attention +- **24 HieraBlocks** in 4 stages: + - Stage 0 (blocks 0-1): dim=112, heads=2, window=64, mask_unit_attn=True + - Stage 1 (blocks 2-4): dim=224, heads=4, window=16, mask_unit_attn=True + - Stage 2 (blocks 5-20): dim=448, heads=8, window=4, mask_unit_attn=False (global) + - Stage 3 (blocks 21-23): dim=896, heads=16, window=1, mask_unit_attn=False (global) +- **Reroll**: Undoes permutation -> spatial [B, H, W, C] +- **Stage transitions** (blocks 2, 5, 21): proj Linear + max-pool reduces tokens 4x +- **Features**: Collected at stage_ends [1, 4, 20, 23], rerolled to NHWC + +### MaskUnitAttention +- qkv = Linear(dim_in, 3*dim_out), reshaped per window +- q_stride > 1 at transitions: max-pool over q_stride dim in queries +- Scaled dot-product attention +- proj = Linear(dim_out, dim_out) + +### MLP +- fc1(dim->4*dim) -> GELU -> fc2(4*dim->dim) + +## Checkpoint Details (from safetensors) + +- 297 encoder keys total +- `encoder.model.patch_embed.proj.weight`: (112, 7, 7, 4) — already transposed by converter +- `encoder.model.pos_embed`: (1, 262144, 112) — raw, needs bicubic interpolation +- Transition blocks (2, 5, 21) have extra `proj.weight/bias` +- All block weights are Linear (no transpose needed) +- No LayerScale weights (init_values=None for base_plus) + +## Implementation Plan + +### Sub-phase 4a: Core backbone module + +**File**: `src/corridorkey_mlx/model/hiera.py` + +Implement (all in [B, N, C] sequence format, NHWC only at boundaries): + +1. **`HieraPatchEmbed`** — Conv2d(4->112, 7x7, stride=4) + flatten to [B, N, C] +2. **`unroll(x, spatial_size, schedule)`** — faithful port of timm's Unroll.forward +3. **`reroll(x, block_idx, schedule_map)`** — faithful port of timm's Reroll.forward (no-mask path) +4. **`undo_windowing(x, shape, mu_shape)`** — helper for reroll +5. **`MaskUnitAttention`** — windowed multi-head attention with optional q max-pool +6. **`HieraMLP`** — fc1 -> GELU -> fc2 +7. **`HieraBlock`** — norm1 -> [proj+maxpool if transition] -> attn + residual -> norm2 -> mlp + residual +8. **`HieraBackbone`** — full assembly: patch_embed + pos_embed + unroll + 24 blocks + reroll at stage_ends + +### Sub-phase 4b: Converter updates + +**File**: `src/corridorkey_mlx/convert/converter.py` + +- Encoder keys are already handled (patch_embed.proj in CONV_WEIGHT_KEYS, all else passthrough) +- No changes needed — verified all 297 keys pass through correctly + +### Sub-phase 4c: Weight loading + pos_embed interpolation + +**In**: `src/corridorkey_mlx/model/hiera.py` (method on HieraBackbone) + +- Load safetensors weights into MLX model +- Bicubic interpolation for pos_embed: (1, 262144, 112) -> (1, N', 112) + - Reshape to (1, H_train, W_train, 112) -> bilinear/bicubic resize -> flatten + +### Sub-phase 4d: Tests + +**Files**: +- `tests/test_hiera_stage_shapes.py` — shape contract tests (no checkpoint needed) +- `tests/test_hiera_stage_parity.py` — numerical parity vs golden fixtures + +Shape tests: +- PatchEmbed output shape +- Each stage feature map shape (128x128x112, 64x64x224, 32x32x448, 16x16x896) +- Correct number of features (4) + +Parity tests (require checkpoint + fixtures): +- Load golden `encoder_feature_{0-3}` from `reference/fixtures/golden.npz` +- Load checkpoint weights into MLX backbone +- Compare each stage output (NCHW fixtures -> NHWC comparison) +- Report max_abs and mean_abs error per stage + +## Key Files + +| File | Action | +|------|--------| +| `src/corridorkey_mlx/model/hiera.py` | create (replace placeholder `backbone.py`) | +| `src/corridorkey_mlx/model/backbone.py` | keep as thin re-export | +| `tests/test_hiera_stage_shapes.py` | create | +| `tests/test_hiera_stage_parity.py` | create | + +## Key Decisions + +1. **Faithful Unroll/Reroll port** — required for exact parity +2. **File naming**: `hiera.py` — more descriptive; `backbone.py` stays as re-export +3. **No masking support** — inference only, skip masked token paths +4. **pos_embed interpolation at load time** — converter outputs raw checkpoint pos_embed +5. **All backbone ops in [B, N, C]** — NHWC only at boundaries (input image, output features) +6. **MLX key prefix**: `encoder.model.` stripped when loading into HieraBackbone + +## Verification + +```bash +uv run pytest tests/test_hiera_stage_shapes.py -v +uv run pytest tests/test_hiera_stage_parity.py -v -s +uv run ruff check src/corridorkey_mlx/model/hiera.py +uv run mypy src/corridorkey_mlx/model/hiera.py +``` + +## Unresolved Questions + +- E2e parity tolerance? Expecting ~1e-4 max_abs, Metal float32 may drift more through 24 blocks +- `mx.fast.scaled_dot_product_attention` availability/API — fallback to manual if needed diff --git a/src/corridorkey_mlx/model/backbone.py b/src/corridorkey_mlx/model/backbone.py index b1e4fdc..b65c721 100644 --- a/src/corridorkey_mlx/model/backbone.py +++ b/src/corridorkey_mlx/model/backbone.py @@ -1,5 +1,8 @@ -"""Hiera backbone — MLX port (not yet implemented). +"""Hiera backbone — thin re-export. -Original: timm Hiera with features_only=True. -Emits 4 multiscale feature maps. +Import from :mod:`corridorkey_mlx.model.hiera` for the full implementation. """ + +from corridorkey_mlx.model.hiera import HieraBackbone + +__all__ = ["HieraBackbone"] diff --git a/src/corridorkey_mlx/model/hiera.py b/src/corridorkey_mlx/model/hiera.py new file mode 100644 index 0000000..2e0ec4b --- /dev/null +++ b/src/corridorkey_mlx/model/hiera.py @@ -0,0 +1,416 @@ +"""Hiera backbone — MLX port. + +Hierarchical vision transformer (timm hiera_base_plus_224) for feature extraction. +Emits 4 multiscale feature maps in NHWC format. + +Ported from: timm/models/hiera.py (Meta Platforms, Apache-2.0 / CC-BY-NC-4.0) +Reference: https://arxiv.org/abs/2306.00989 +""" + +from __future__ import annotations + +from functools import reduce + +import mlx.core as mx +import mlx.nn as nn + +# ── hiera_base_plus_224 constants ────────────────────────────────────── +EMBED_DIM = 112 +NUM_HEADS = 2 +STAGES = (2, 3, 16, 3) +Q_POOL = 3 # number of stages with q-pooling +Q_STRIDE = (2, 2) +MASK_UNIT_SIZE = (8, 8) +MASK_UNIT_ATTN = (True, True, False, False) +PATCH_KERNEL = (7, 7) +PATCH_STRIDE = (4, 4) +PATCH_PADDING = (3, 3) +IN_CHANS = 4 # RGB + alpha hint +MLP_RATIO = 4.0 +DIM_MUL = 2.0 +HEAD_MUL = 2.0 + + +# ── Helpers ──────────────────────────────────────────────────────────── + + +def _prod(seq: tuple[int, ...] | list[int]) -> int: + return reduce(lambda a, b: a * b, seq, 1) + + +def undo_windowing( + x: mx.array, + shape: list[int], + mu_shape: list[int], +) -> mx.array: + """Undo mask-unit windowing: [B, #MUs, MUy, MUx, C] -> [B, H, W, C]. + + Faithful port of timm ``undo_windowing`` (2d only). + """ + ndim = len(shape) # spatial dims (2 for images) + batch_size = x.shape[0] + channels = x.shape[-1] + + num_mus = [s // mu for s, mu in zip(shape, mu_shape, strict=True)] + # [B, #MUy*#MUx, MUy, MUx, C] -> [B, #MUy, #MUx, MUy, MUx, C] + x = x.reshape([batch_size] + num_mus + mu_shape + [channels]) + + # Interleave: [B, #MUy, MUy, #MUx, MUx, C] + perm = ( + [0] + + sum( + [list(p) for p in zip(range(1, 1 + ndim), range(1 + ndim, 1 + 2 * ndim), strict=True)], + [], + ) + + [len(x.shape) - 1] + ) + x = mx.transpose(x, axes=perm) + return x.reshape([batch_size] + shape + [channels]) + + +def unroll(x: mx.array, spatial_size: list[int], schedule: list[tuple[int, int]]) -> mx.array: + """Reorder tokens so patches are contiguous for windowed ops. + + Faithful port of timm ``Unroll.forward`` (2d only, inference path). + + Input: [B, N, C] (flattened patch embeddings) + Output: [B', N', C] where B' = B * prod(all strides), N' = prod(cur_size) + """ + batch_size = x.shape[0] + channels = x.shape[-1] + cur_size = list(spatial_size) + + # [B, N, C] -> [B, H, W, C] + x = x.reshape([batch_size] + cur_size + [channels]) + + for strides in schedule: + cur_size = [i // s for i, s in zip(cur_size, strides, strict=True)] + # [B, H//Sy, Sy, W//Sx, Sx, C] + pairs = sum([[i, s] for i, s in zip(cur_size, strides, strict=True)], []) + new_shape = [batch_size] + pairs + [channels] + x = x.reshape(new_shape) + + # [B, Sy, Sx, H//Sy, W//Sx, C] + ndims = len(new_shape) + perm = [0] + list(range(2, ndims - 1, 2)) + list(range(1, ndims - 1, 2)) + [ndims - 1] + x = mx.transpose(x, axes=perm) + + # Flatten strides into batch: B' = B * Sy * Sx + stride_count = _prod(strides) + x = x.reshape([batch_size * stride_count] + cur_size + [channels]) + batch_size *= stride_count + + # Flatten spatial back using original size → [B, N_orig, C] + # This collapses the inflated batch back to original B + return x.reshape(-1, _prod(spatial_size), channels) + + +def reroll( + x: mx.array, + block_idx: int, + schedule_map: dict[int, tuple[list[tuple[int, int]], list[int]]], +) -> mx.array: + """Undo unroll to recover spatial layout. + + Faithful port of timm ``Reroll.forward`` (2d only, no-mask inference path). + + Input: [B', N, C] + Output: [B, H, W, C] (NHWC) + """ + remaining_schedule, size = schedule_map[block_idx] + batch_size = x.shape[0] + num_tokens = x.shape[1] + channels = x.shape[-1] + + ndim = len(size) # 2 for images + cur_mu_shape = [1] * ndim + + for strides in remaining_schedule: + # [B, *strides, N//(Sy*Sx), *cur_mu_shape, C] + stride_prod = _prod(strides) + inner_n = num_tokens // stride_prod + x = x.reshape([batch_size] + list(strides) + [inner_n] + cur_mu_shape + [channels]) + + # Permute: [B, N//(Sy*Sx), Sy, MUy, Sx, MUx, C] + total_dims = len(x.shape) + perm = ( + [0, 1 + ndim] + + sum( + [ + list(p) + for p in zip( + range(1, 1 + ndim), range(1 + ndim + 1, total_dims - 1), strict=True + ) + ], + [], + ) + + [total_dims - 1] + ) + x = mx.transpose(x, axes=perm) + + # Update mu_shape and reshape + for i in range(ndim): + cur_mu_shape[i] *= strides[i] + x = x.reshape([batch_size, -1] + cur_mu_shape + [channels]) + num_tokens = x.shape[1] + + # [B, #MUs, MUy, MUx, C] + x = x.reshape([batch_size, num_tokens] + cur_mu_shape + [channels]) + + # No mask -> return [B, H, W, C] + return undo_windowing(x, size, cur_mu_shape) + + +# ── Modules ──────────────────────────────────────────────────────────── + + +class HieraPatchEmbed(nn.Module): + """Patch embedding: Conv2d(4->112, 7x7, stride=4) + flatten to [B, N, C].""" + + def __init__(self) -> None: + super().__init__() + self.proj = nn.Conv2d( + IN_CHANS, + EMBED_DIM, + kernel_size=PATCH_KERNEL[0], + stride=PATCH_STRIDE[0], + padding=PATCH_PADDING[0], + ) + + def __call__(self, x: mx.array) -> mx.array: + """Input: (B, H, W, 4) NHWC. Output: (B, N, C) where N = (H/4)*(W/4).""" + x = self.proj(x) # (B, H/4, W/4, 112) + batch_size = x.shape[0] + channels = x.shape[-1] + return x.reshape(batch_size, -1, channels) # (B, N, 112) + + +class HieraMLP(nn.Module): + """MLP: fc1(dim -> 4*dim) -> GELU -> fc2(4*dim -> dim).""" + + def __init__(self, dim: int, mlp_ratio: float = MLP_RATIO) -> None: + super().__init__() + hidden = int(dim * mlp_ratio) + self.fc1 = nn.Linear(dim, hidden) + self.fc2 = nn.Linear(hidden, dim) + + def __call__(self, x: mx.array) -> mx.array: + return self.fc2(nn.gelu(self.fc1(x))) + + +class MaskUnitAttention(nn.Module): + """Windowed multi-head attention with optional q max-pool. + + Operates on unrolled [B, N, C] tokens. When use_mask_unit_attn=True, + attention is computed within windows of size window_size. + """ + + def __init__( + self, + dim: int, + dim_out: int, + heads: int, + q_stride: int = 1, + window_size: int = 0, + use_mask_unit_attn: bool = False, + ) -> None: + super().__init__() + self.dim_out = dim_out + self.heads = heads + self.q_stride = q_stride + self.head_dim = dim_out // heads + self.scale = self.head_dim**-0.5 + self.window_size = window_size + self.use_mask_unit_attn = use_mask_unit_attn + + self.qkv = nn.Linear(dim, 3 * dim_out) + self.proj = nn.Linear(dim_out, dim_out) + + def __call__(self, x: mx.array) -> mx.array: + """Input: [B, N, C]. Output: [B, N', dim_out] (N' = N/q_stride if q_stride>1).""" + batch_size, num_tokens, _ = x.shape + num_windows = ( + (num_tokens // (self.q_stride * self.window_size)) if self.use_mask_unit_attn else 1 + ) + + # QKV projection + reshape to [B, N/num_windows, num_windows, 3, heads, head_dim] + qkv = self.qkv(x) + qkv = qkv.reshape(batch_size, -1, num_windows, 3, self.heads, self.head_dim) + # Permute to [3, B, heads, num_windows, tokens_per_window, head_dim] + qkv = mx.transpose(qkv, axes=(3, 0, 4, 2, 1, 5)) + q, k, v = qkv[0], qkv[1], qkv[2] + + if self.q_stride > 1: + # Max-pool over q_stride tokens in the query + # [B, heads, num_windows, q_stride, tokens/q_stride, head_dim] + q = q.reshape(batch_size, self.heads, num_windows, self.q_stride, -1, self.head_dim) + q = mx.max(q, axis=3) + + # Scaled dot-product attention (manual implementation) + attn = (q * self.scale) @ mx.transpose(k, axes=(0, 1, 2, 4, 3)) + attn = mx.softmax(attn, axis=-1) + x = attn @ v + + # [B, heads, num_windows, tokens, head_dim] -> [B, num_windows, tokens, heads, head_dim] + # -> [B, N', dim_out] + x = mx.transpose(x, axes=(0, 2, 3, 1, 4)) + x = x.reshape(batch_size, -1, self.dim_out) + + return self.proj(x) + + +class HieraBlock(nn.Module): + """Single Hiera transformer block. + + At transition blocks (dim != dim_out): proj Linear + q max-pool reduces tokens 4x. + No DropPath or LayerScale at inference (drop_path=0, init_values=None). + """ + + def __init__( + self, + dim: int, + dim_out: int, + heads: int, + q_stride: int = 1, + window_size: int = 0, + use_mask_unit_attn: bool = False, + ) -> None: + super().__init__() + self.dim = dim + self.dim_out = dim_out + self.do_expand = dim != dim_out + self.q_stride = q_stride + + self.norm1 = nn.LayerNorm(dim) + if self.do_expand: + self.proj = nn.Linear(dim, dim_out) + + self.attn = MaskUnitAttention( + dim, dim_out, heads, q_stride, window_size, use_mask_unit_attn + ) + + self.norm2 = nn.LayerNorm(dim_out) + self.mlp = HieraMLP(dim_out) + + def __call__(self, x: mx.array) -> mx.array: + # Attention + optional Q pooling + x_norm = self.norm1(x) + if self.do_expand: + x = self.proj(x_norm) + # Max-pool over q_stride tokens to reduce spatial resolution + batch_size = x.shape[0] + x = x.reshape(batch_size, self.q_stride, -1, x.shape[-1]) + x = mx.max(x, axis=1) + x = x + self.attn(x_norm) + + # MLP + x = x + self.mlp(self.norm2(x)) + return x + + +class HieraBackbone(nn.Module): + """Full Hiera backbone: patch_embed + pos_embed + unroll + 24 blocks + reroll. + + Outputs 4 multiscale feature maps in NHWC at stage_ends [1, 4, 20, 23]. + Hardcoded for hiera_base_plus_224 config with 4-channel input. + """ + + def __init__(self, img_size: int = 512) -> None: + super().__init__() + self.img_size = img_size + + # Spatial size after patching + self.tokens_spatial_shape = [img_size // PATCH_STRIDE[0], img_size // PATCH_STRIDE[1]] + + # Stage ends and q_pool blocks + self.stage_ends = [sum(STAGES[:i]) - 1 for i in range(1, len(STAGES) + 1)] + # [1, 4, 20, 23] + q_pool_blocks = [self.stage_ends[i] + 1 for i in range(Q_POOL)] + # [2, 5, 21] + + # Unroll schedule + unroll_schedule = [Q_STRIDE] * len(self.stage_ends[:-1]) + # [(2,2), (2,2), (2,2)] + + # Precompute reroll schedule map + self._reroll_schedule: dict[int, tuple[list[tuple[int, int]], list[int]]] = {} + size = list(self.tokens_spatial_shape) + cur_schedule = list(unroll_schedule) + for i in range(self.stage_ends[-1] + 1): + self._reroll_schedule[i] = (list(cur_schedule), list(size)) + if i in self.stage_ends[:Q_POOL]: + if len(cur_schedule) > 0: + size = [n // s for n, s in zip(size, cur_schedule[0], strict=True)] + cur_schedule = cur_schedule[1:] + + # Store unroll params + self._unroll_spatial = list(self.tokens_spatial_shape) + self._unroll_schedule = unroll_schedule + + # Patch embedding + self.patch_embed = HieraPatchEmbed() + + # Positional embedding (loaded from checkpoint, interpolated at load time) + self.pos_embed = mx.zeros((1, _prod(self.tokens_spatial_shape), EMBED_DIM)) + + # Build all 24 blocks + self.blocks: list[HieraBlock] = [] + embed_dim = EMBED_DIM + num_heads = NUM_HEADS + flat_mu_size = _prod(MASK_UNIT_SIZE) + flat_q_stride = _prod(Q_STRIDE) + cur_stage = 0 + + for i in range(sum(STAGES)): + dim_out = embed_dim + use_mu_attn = MASK_UNIT_ATTN[cur_stage] + + if i - 1 in self.stage_ends: + dim_out = int(embed_dim * DIM_MUL) + num_heads = int(num_heads * HEAD_MUL) + cur_stage += 1 + if i in q_pool_blocks: + flat_mu_size //= flat_q_stride + + block = HieraBlock( + dim=embed_dim, + dim_out=dim_out, + heads=num_heads, + q_stride=(flat_q_stride if i in q_pool_blocks else 1), + window_size=flat_mu_size, + use_mask_unit_attn=use_mu_attn, + ) + self.blocks.append(block) + embed_dim = dim_out + + def __call__(self, x: mx.array) -> list[mx.array]: + """Forward pass. + + Args: + x: Input image (B, H, W, 4) in NHWC — ImageNet-normalized RGB + alpha hint. + + Returns: + 4 feature maps in NHWC: + [0]: (B, H/4, W/4, 112) + [1]: (B, H/8, W/8, 224) + [2]: (B, H/16, W/16, 448) + [3]: (B, H/32, W/32, 896) + """ + # Patch embed -> [B, N, C] + x = self.patch_embed(x) + + # Add positional embedding + x = x + self.pos_embed + + # Unroll for windowed attention + x = unroll(x, self._unroll_spatial, self._unroll_schedule) + + # Run blocks, collecting features at stage_ends + features: list[mx.array] = [] + for i, blk in enumerate(self.blocks): + x = blk(x) + if i in self.stage_ends: + feat = reroll(x, i, self._reroll_schedule) + features.append(feat) + + return features