Release LESSViT checkpoint and directly loadable encoder

README.md

@@ -0,0 +1,168 @@
+---
+library_name: transformers
+language:
+- en
+tags:
+- lessvit
+- hyperspectral
+- remote-sensing
+- geospatial
+- vision-transformer
+- custom_code
+arxiv: 2605.18541
+---
+
+# LESSViT-b2-d8-r1
+
+Pretrained **LESSViT** encoder for hyperspectral representation learning, exported
+from the `LESSVIT_b2_d8_r1` checkpoint. The root of this repository contains an
+encoder-only Hugging Face model, with standalone implementation files and a
+working `AutoModel` interface. No project checkout is required.
+
+- [Paper](https://arxiv.org/abs/2605.18541)
+- [Project website and updated manuscript](https://uiuctml.github.io/LESSViT/)
+- [Source repository](https://github.com/uiuctml/LESSViT)
+
+## Install
+
+Reference environment: Python 3.11.
+
+```bash
+pip install torch==2.7.1 torchvision==0.22.1 transformers==4.53.3 \
+  huggingface-hub==0.34.4 timm==1.0.18 numpy==2.2.6 safetensors
+```
+
+Choose matching PyTorch/Torchvision wheels for your hardware. This release uses
+custom model code, so `trust_remote_code=True` is required. You can pin `revision`
+to a repository commit SHA for reproducible loading.
+
+## Load the encoder directly
+
+```python
+import torch
+from transformers import AutoModel
+
+encoder = AutoModel.from_pretrained(
+    "GFM-Bench/LESSViT-b2-d8-r1",
+    trust_remote_code=True,
+).eval()
+
+# Synthetic example: [batch, spectral channels, height, width].
+x = torch.randn(1, 10, 32, 32)
+wavelengths_nm = torch.linspace(450, 2300, 10)[None, :]
+# Fixed full-EnMAP reference bounds for this example. Keep reference bounds
+# consistent across your downstream training and evaluation configurations.
+wavelengths = (wavelengths_nm - 418.24) / (2445.53 - 418.24)
+
+with torch.no_grad():
+    output = encoder(
+        optical=x,
+        optical_channel_wv=wavelengths,
+        spatial_resolution=30,  # meters per pixel
+    )
+
+features = output.last_hidden_state
+print(features.shape)                  # [1, 11, 5, 768]
+global_features = features[:, 0, 0]     # [1, 768], classification
+spatial_features = features[:, 0, 1:]   # [1, 4, 768], dense prediction
+band_patch_features = features[:, 1:, 1:]  # [1, 10, 4, 768]
+```
+
+`load_encoder.py` contains the same runnable example:
+
+```bash
+python load_encoder.py --model-id GFM-Bench/LESSViT-b2-d8-r1 --device cpu
+```
+
+To run offline after downloading a complete snapshot, pass its directory as
+`--model-id`. No network or original repository imports are needed for a local
+snapshot.
+
+### Inputs and outputs
+
+- `optical`: float tensor `[B, C, H, W]`; height and width should be multiples of 16.
+- `optical_channel_wv`: normalized wavelengths `[1, C]` or `[B, C]`, matching input
+  channel order. Use fixed reference bounds for the task, not new min/max bounds
+  for every spectral subset. The numerical bounds above are an EnMAP example,
+  not a universal normalization rule for every sensor.
+- `spatial_resolution`: spatial resolution in meters per pixel.
+- `last_hidden_state`: `[B, C+1, N+1, 768]`, where
+  `N = (H / 16) * (W / 16)`. Index 0 contains aggregation tokens.
+
+For real data, use dataset/sensor-specific intensity normalization and matching
+wavelength metadata. The source repository includes dataset statistics,
+transforms, SRF resampling, and collators. Random values in the example only
+verify that loading and inference work.
+
+## Attach a downstream head
+
+```python
+# Use `encoder` and example tensors from above.
+encoder.train()
+head = torch.nn.Linear(encoder.config.embed_dim, 14)
+features = encoder(
+    optical=x, optical_channel_wv=wavelengths, spatial_resolution=30
+).last_hidden_state[:, 0, 0]
+logits = head(features)
+loss = torch.nn.functional.cross_entropy(logits, torch.zeros(1, dtype=torch.long))
+loss.backward()  # gradients through both encoder and head
+```
+
+For multi-label classification use binary cross entropy with logits. For dense
+prediction, reshape `spatial_features` into the patch grid and attach a spatial
+decoder. Task-specific heads are not included and must be trained. MAE masking
+and random channel dropout are disabled in this encoder interface, including
+when calling `.train()`; ordinary configured dropout remains active in training.
+
+## Architecture
+
+| Setting | Value |
+|---|---|
+| Encoder depth | 12 |
+| Embedding dimension | 768 |
+| Attention heads | 12 |
+| Patch size | 16 |
+| Spectral dimension per head | 2 |
+| Attention rank | 1 |
+| Positional encoding | Spatial–spectral RoPE |
+| LayerScale initialization | 1.0 |
+
+`d8` in the checkpoint name refers to the **pretraining decoder**, not the
+encoder. The encoder-only release has approximately 83.2M parameters.
+
+## Files and provenance
+
+- Root `model.safetensors`: extracted encoder, with legacy attention keys
+  remapped to the current implementation. No encoder tensors were omitted.
+- Root `config.json`: explicit encoder architecture and Hugging Face auto-map.
+- `pretraining/model.safetensors`: the original full checkpoint, copied
+  byte-for-byte, including the MAE decoder. It uses the original checkpoint key
+  layout and is **not** the root `AutoModel` format.
+- `pretraining/architecture.json`: architectural fields from the accompanying
+  training config. Machine-specific paths and training service settings are
+  omitted. This is provenance metadata, not an AutoModel config.
+- `validation.json`: export verification results.
+- `SHA256SUMS`: file hashes.
+
+Encoder implementation is packaged from LESSViT source commit
+`20fb37c10541f1a29eb2a53999afdc0734499bc3`, with a small Hugging Face wrapper to
+preserve architecture settings across config serialization and disable MAE
+masking for encoder use. Loading checks compare every exported encoder tensor
+exactly to the supplied checkpoint after key remapping. Local AutoModel loading,
+forward-output equivalence, and encoder backward propagation were tested on CPU.
+These checks do not independently reproduce the paper's benchmark scores or
+establish the checkpoint's complete training-data provenance.
+
+## Citation
+
+```bibtex
+@misc{si2026lessvitrobusthyperspectralrepresentation,
+  title={LESSViT: Robust Hyperspectral Representation Learning under Spectral Configuration Shift},
+  author={Haozhe Si and Yuxuan Wan and Yuqing Wang and Minh N. Do and Han Zhao},
+  year={2026},
+  eprint={2605.18541},
+  archivePrefix={arXiv},
+  primaryClass={cs.CV},
+  url={https://arxiv.org/abs/2605.18541}
+}
+```

SHA256SUMS

@@ -0,0 +1,14 @@
+66bd1865a04156e70e66415709e2ff39fe7dcf71f686002499a6da650274c28c  README.md
+d6fcf9ed303752ba5ea7792d3fa00e16ac35cd561b6df23e5a19f4d219be6aac  attention_ops.py
+5aee6c840ac260cde0131ba16ac4c637da27cbe63f94aa0cd5d72f6171436e5d  config.json
+0e049d15e1a8229824e4b87f0d3312f55e58c0a2985bd36a69fa5f7790d024de  hyperspectral_patch_embed.py
+780ee1bf2b7119a3cbc3a4f84f5e185c3f0a8c455f9c23deb401f15556333efe  load_encoder.py
+e1e7af883a62f663010397bc46fcae04801ed2465ed830e8bb488a044f05f809  low_rank_attention.py
+935719728265cf0765b32afe634803ab73071025f104c1c708b98142488b8311  model.safetensors
+1b29142386ee2d45deb01325fdf7cc63e91fd1503ca1017ac0d47db9341e3111  modeling_lessvit.py
+4c994cbd89ca4dc5f31ca711b8c7c27b0ec6b3fea7755162e9279fff927c47b7  pos_chan_embed.py
+fe35811d65813247e513749d8fab57b082d553d0e371b78ddc39c60023817f19  pretraining/architecture.json
+8e5dce5229b24255d31dc619c005ad6355784eff7ee9e07036aba6f36a1c606a  pretraining/model.safetensors
+9d58c64770e91d2f2dfae27194b600f680c95649791ead9510fc0bd83fedd5c2  requirements.txt
+5d822f2a1439cea4bba10696277a6ec56cea63e337fd1e05cdec5d2f0ce19ab8  spatial_spectral_low_rank_vit.py
+68301b5051270d261572715d0cbb4e003d4b901dbb9e202387cae29a93afe6d6  validation.json

attention_ops.py

@@ -0,0 +1,432 @@
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from timm.layers import Mlp, DropPath, use_fused_attn
+
+from .low_rank_attention import AttentionBranch, LowDimPool, LowRankAttention
+from .pos_chan_embed import rope_rotate_half
+
+__all__ = [
+    'SpatialSpectralAttention',
+    'LESSAttention',
+    'FullSpatialSpectralAttention',
+    'AdditiveSpatialSpectralAttention',
+    'build_attention',
+    'LayerScale',
+    'LowRankBlock',
+]
+
+
+class SpatialSpectralAttention(nn.Module):
+    """
+    Common interface for the block-level spatial-spectral attention operator.
+
+    forward(x, spatial_mask=None, rope=None) -> Tensor
+        x: [B, C, HW, D] (C and HW here are whatever the caller's grid axis sizes
+           currently are, e.g. C+1/HW+1 when CLS tokens are folded in).
+        spatial_mask: optional perception-field mask, [HW, HW] or None.
+        rope: optional (sin_hw, cos_hw, sin_c, cos_c) tuple from
+           RopePositionChannelEmbedding, or None.
+        returns: [B, C, HW, D], same shape as x.
+    """
+    pass
+
+
+class LESSAttention(SpatialSpectralAttention):
+    """
+    Arm A: the existing LESS (Kronecker) attention. Thin composition of the
+    unmodified LowDimPool + LowRankAttention classes behind the unified interface;
+    no math differs from the original LowRankBlock.forward.
+    """
+    def __init__(
+        self,
+        dim: int,
+        channel_dim: int,
+        spatial_dim: int,
+        num_heads: int = 8,
+        qkv_bias: bool = False,
+        qk_norm: bool = False,
+        attn_drop: float = 0.,
+        proj_drop: float = 0.,
+        norm_layer: nn.Module = nn.LayerNorm,
+        rank: int = 1,
+        skip_pool: bool = False,
+    ) -> None:
+        super().__init__()
+        self.low_dim_pool = LowDimPool(
+            dim=dim,
+            channel_dim=channel_dim,
+            spatial_dim=spatial_dim,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            qk_norm=qk_norm,
+            attn_drop=attn_drop,
+            proj_drop=proj_drop,
+            norm_layer=norm_layer,
+            skip_pool=skip_pool,
+        )
+        self.channel_norm = norm_layer(channel_dim)
+        self.spatial_norm = norm_layer(spatial_dim)
+        self.attn = LowRankAttention(
+            dim=dim,
+            channel_dim=channel_dim,
+            spatial_dim=spatial_dim,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            qk_norm=qk_norm,
+            attn_drop=attn_drop,
+            proj_drop=proj_drop,
+            norm_layer=norm_layer,
+            rank=rank,
+        )
+
+    def forward(self, x: torch.Tensor, spatial_mask: torch.Tensor = None, rope: tuple = None) -> torch.Tensor:
+        x_c, x_s, _ = self.low_dim_pool(x)
+        return self.attn(self.channel_norm(x_c), self.spatial_norm(x_s), spatial_mask, rope)
+
+
+class AdditiveSpatialSpectralAttention(SpatialSpectralAttention):
+    """
+    Arm C: *identical* channel/spatial branch attentions as arm A -- same
+    LowDimPool, same AttentionBranch class, same channel_dim/spatial_dim (so
+    c_head_dim/s_head_dim keep arm A's low-rank widths, e.g. 4 and 16 for ViT-S,
+    NOT widened to D_h). This isolates the fusion mechanism specifically, rather
+    than also widening the branch attentions (an earlier version of this class
+    set channel_dim=spatial_dim=dim, which inflated LowDimPool's projections and
+    the branch attentions themselves and produced a ~3x arm A param count -- this
+    version's param delta over arm A is just P_c/P_s, i.e. genuinely "slight").
+
+    Fusion: each branch's per-head output is projected *up* to the full head_dim
+    D_h only at the fusion step, then broadcast-added:
+
+        Y[c, n] = P_c(Y_C[c]) + P_s(Y_S[n]),   P_c: R^c_head_dim -> R^D_h,
+                                                 P_s: R^s_head_dim -> R^D_h
+
+    P_c/P_s are shared across heads (a single small matrix broadcast over the
+    heads dimension, mirroring how the Kronecker fusion also has no per-head-
+    distinct fusion parameters). `fusion_init_scale` rescales P_c/P_s's
+    initialization so that arm C's post-fusion activation std at initialization
+    matches arm A's (see ablation_sanity_checks.py check #4) -- the two fusion
+    mechanisms are algebraically different (product vs. sum) so matching
+    variance requires an explicit scale, not just "sensible" init.
+    """
+    def __init__(
+        self,
+        dim: int,
+        channel_dim: int,
+        spatial_dim: int,
+        num_heads: int = 8,
+        qkv_bias: bool = False,
+        qk_norm: bool = False,
+        attn_drop: float = 0.,
+        proj_drop: float = 0.,
+        norm_layer: nn.Module = nn.LayerNorm,
+        rank: int = 1,
+        skip_pool: bool = False,
+        fusion_init_scale: float = 1.0,
+    ) -> None:
+        super().__init__()
+        assert channel_dim % num_heads == 0, 'channel_dim should be divisible by num_heads'
+        assert spatial_dim % num_heads == 0, 'spatial_dim should be divisible by num_heads'
+        assert dim % num_heads == 0, 'dim should be divisible by num_heads'
+
+        self.num_heads = num_heads
+        self.dim = dim
+        self.head_dim = dim // num_heads
+        self.c_head_dim = channel_dim // num_heads
+        self.s_head_dim = spatial_dim // num_heads
+        self.rank = rank
+
+        self.low_dim_pool = LowDimPool(
+            dim=dim,
+            channel_dim=channel_dim,
+            spatial_dim=spatial_dim,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            qk_norm=qk_norm,
+            attn_drop=attn_drop,
+            proj_drop=proj_drop,
+            norm_layer=norm_layer,
+            skip_pool=skip_pool,
+        )
+        self.channel_norm = norm_layer(channel_dim)
+        self.spatial_norm = norm_layer(spatial_dim)
+
+        self.channel_branch = AttentionBranch(channel_dim, num_heads, self.c_head_dim, qkv_bias, qk_norm, attn_drop, norm_layer, rank)
+        self.spatial_branch = AttentionBranch(spatial_dim, num_heads, self.s_head_dim, qkv_bias, qk_norm, attn_drop, norm_layer, rank)
+
+        self.P_c = nn.Linear(self.c_head_dim, self.head_dim)
+        self.P_s = nn.Linear(self.s_head_dim, self.head_dim)
+        with torch.no_grad():
+            # Scale the whole affine map (weight AND bias) so that P(x) -> scale*P(x)
+            # exactly, uniformly rescaling arm C's post-fusion output to match arm A's
+            # activation std at init (see ablation_sanity_checks.check_fusion_init_scale_calibration).
+            self.P_c.weight.mul_(fusion_init_scale)
+            self.P_c.bias.mul_(fusion_init_scale)
+            self.P_s.weight.mul_(fusion_init_scale)
+            self.P_s.bias.mul_(fusion_init_scale)
+
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x_grid: torch.Tensor, spatial_mask: torch.Tensor = None, rope: tuple = None) -> torch.Tensor:
+        x_c, x_s, _ = self.low_dim_pool(x_grid)
+        C = x_c.shape[1]
+        HW = x_s.shape[1]
+
+        sin_hw, cos_hw, sin_c, cos_c = rope if rope is not None else (None, None, None, None)
+
+        xc = self.channel_branch(self.channel_norm(x_c), sin=sin_c, cos=cos_c)  # B, rank, num_heads, C, c_head_dim
+        xs = self.spatial_branch(self.spatial_norm(x_s), spatial_mask, sin=sin_hw, cos=cos_hw)  # B, rank, num_heads, HW, s_head_dim
+
+        xc = xc.sum(dim=1) / self.rank  # B, num_heads, C, c_head_dim
+        xs = xs.sum(dim=1) / self.rank  # B, num_heads, HW, s_head_dim
+
+        Yc = self.P_c(xc)  # B, num_heads, C, head_dim
+        Ys = self.P_s(xs)  # B, num_heads, HW, head_dim
+
+        Y = Yc.unsqueeze(3) + Ys.unsqueeze(2)  # B, num_heads, C, HW, head_dim
+        B = Y.shape[0]
+        Y = Y.permute(0, 2, 3, 1, 4).reshape(B, C, HW, -1)  # B, C, HW, D
+        Y = self.proj(Y)
+        Y = self.proj_drop(Y)
+        return Y
+
+
+def _apply_axis_rope(x: torch.Tensor, sin: Optional[torch.Tensor], cos: Optional[torch.Tensor], axis: int) -> torch.Tensor:
+    """
+    Rotate a sub-block of the head dimension of `x` ([B, num_heads, C, HW, W])
+    along one grid axis (axis=2 for the channel axis, axis=3 for the position
+    axis), skipping index 0 along that axis (the channel-CLS row when axis=2, the
+    spatial-CLS column when axis=3). sin/cos: [B, L, W] where L == x.shape[axis]-1.
+
+    Reuses the exact rotate-half formula from pos_chan_embed.rope_apply; only the
+    broadcasting is generalized here to the extra grid axis that arm A's per-branch
+    tensors never had.
+    """
+    if sin is None:
+        return x
+    L = x.shape[axis] - 1
+    x_cls, x_rest = x.split([1, L], dim=axis)
+    view_shape = [x.shape[0]] + [1] * (x.dim() - 2) + [x.shape[-1]]
+    view_shape[axis] = L
+    sin = sin.reshape(view_shape)
+    cos = cos.reshape(view_shape)
+    x_rot = x_rest * cos + rope_rotate_half(x_rest) * sin
+    return torch.cat([x_cls, x_rot], dim=axis)
+
+
+class FullSpatialSpectralAttention(SpatialSpectralAttention):
+    """
+    Arm B: standard full softmax attention over the flattened (channel, position)
+    token grid (paper Appendix E, Algorithm 2). No low-rank branch factorization,
+    no Kronecker fusion -- a single Q/K/V, softmax over all C*HW tokens.
+
+    Joint RoPE: partitions the head dimension D_h into a spatial (u,v) sub-block
+    of width `spatial_split_dim` and a spectral (lambda) sub-block of width
+    `spectral_split_dim` (spatial_split_dim + spectral_split_dim == D_h). Each
+    sub-block is rotated independently -- spatial sub-block using (sin_hw, cos_hw)
+    broadcast over the channel axis, spectral sub-block using (sin_c, cos_c)
+    broadcast over the position axis -- reusing RopePositionChannelEmbedding's
+    existing frequency construction unmodified (see model config: rope is built
+    with rope_spatial_split_dim/rope_channel_split_dim instead of the per-branch
+    spatial_dim/channel_dim that arms A/C use). Rotation is applied while the grid
+    is still 2D-indexed (before the C/HW flatten), so each axis's CLS-skip is
+    independent: real tokens (c>=1, n>=1) get both rotations; the channel-CLS row
+    (c=0) gets spatial rotation only; the spatial-CLS column (n=0) gets spectral
+    rotation only; the global CLS (0,0) gets neither.
+    """
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int = 8,
+        qkv_bias: bool = False,
+        qk_norm: bool = False,
+        attn_drop: float = 0.,
+        proj_drop: float = 0.,
+        norm_layer: nn.Module = nn.LayerNorm,
+    ) -> None:
+        super().__init__()
+        assert dim % num_heads == 0, 'dim should be divisible by num_heads'
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        assert self.head_dim % 2 == 0, 'head_dim must be even for RoPE rotate-half'
+        self.scale = self.head_dim ** -0.5
+        self.fused_attn = use_fused_attn()
+
+        self.spatial_split_dim = self.head_dim // 2
+        self.spectral_split_dim = self.head_dim - self.spatial_split_dim
+        assert self.spatial_split_dim % 2 == 0 and self.spectral_split_dim % 2 == 0, \
+            'both rope sub-blocks must be even for rotate-half'
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.q_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
+        self.k_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def apply_joint_rope(self, q: torch.Tensor, k: torch.Tensor, rope: tuple):
+        sin_hw, cos_hw, sin_c, cos_c = rope
+
+        def rotate(t: torch.Tensor) -> torch.Tensor:
+            t_spa, t_spec = t.split([self.spatial_split_dim, self.spectral_split_dim], dim=-1)
+            t_spa = _apply_axis_rope(t_spa, sin_hw, cos_hw, axis=3)   # rotate along the HW (position) axis
+            t_spec = _apply_axis_rope(t_spec, sin_c, cos_c, axis=2)  # rotate along the C (channel) axis
+            return torch.cat([t_spa, t_spec], dim=-1)
+
+        return rotate(q), rotate(k)
+
+    def forward(self, x: torch.Tensor, spatial_mask: torch.Tensor = None, rope: tuple = None) -> torch.Tensor:
+        if spatial_mask is not None:
+            raise NotImplementedError(
+                'perception_field_mask is not supported for attn_type="full" in this ablation '
+                '(not exercised by the shared base config -- fail loudly rather than silently ignore it)'
+            )
+        B, C, HW, D = x.shape
+        qkv = self.qkv(x).reshape(B, C, HW, 3, self.num_heads, self.head_dim).permute(3, 0, 4, 1, 2, 5)
+        q, k, v = qkv.unbind(0)  # B, num_heads, C, HW, head_dim
+
+        if rope is not None:
+            q, k = self.apply_joint_rope(q, k, rope)
+
+        q, k = self.q_norm(q), self.k_norm(k)
+
+        q = q.reshape(B, self.num_heads, C * HW, self.head_dim)
+        k = k.reshape(B, self.num_heads, C * HW, self.head_dim)
+        v = v.reshape(B, self.num_heads, C * HW, self.head_dim)
+
+        if self.fused_attn:
+            out = F.scaled_dot_product_attention(
+                q, k, v,
+                dropout_p=self.attn_drop.p if self.training else 0.,
+            )
+        else:
+            attn = (q @ k.transpose(-2, -1)) * self.scale
+            attn = attn.softmax(dim=-1)
+            attn = self.attn_drop(attn)
+            out = attn @ v
+
+        out = out.reshape(B, self.num_heads, C, HW, self.head_dim)
+        out = out.permute(0, 2, 3, 1, 4).reshape(B, C, HW, D)
+        out = self.proj(out)
+        out = self.proj_drop(out)
+        return out
+
+
+def build_attention(
+    attn_type: str,
+    *,
+    dim: int,
+    channel_dim: int,
+    spatial_dim: int,
+    num_heads: int,
+    rank: int = 1,
+    qkv_bias: bool = False,
+    qk_norm: bool = False,
+    attn_drop: float = 0.,
+    proj_drop: float = 0.,
+    norm_layer: nn.Module = nn.LayerNorm,
+    skip_pool: bool = False,
+    fusion_init_scale: float = 1.0,
+) -> SpatialSpectralAttention:
+    if attn_type == 'less':
+        return LESSAttention(
+            dim=dim, channel_dim=channel_dim, spatial_dim=spatial_dim, num_heads=num_heads,
+            qkv_bias=qkv_bias, qk_norm=qk_norm, attn_drop=attn_drop, proj_drop=proj_drop,
+            norm_layer=norm_layer, rank=rank, skip_pool=skip_pool,
+        )
+    elif attn_type == 'additive':
+        return AdditiveSpatialSpectralAttention(
+            dim=dim, channel_dim=channel_dim, spatial_dim=spatial_dim, num_heads=num_heads,
+            qkv_bias=qkv_bias, qk_norm=qk_norm, attn_drop=attn_drop, proj_drop=proj_drop,
+            norm_layer=norm_layer, rank=rank, skip_pool=skip_pool, fusion_init_scale=fusion_init_scale,
+        )
+    elif attn_type == 'full':
+        return FullSpatialSpectralAttention(
+            dim=dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_norm=qk_norm,
+            attn_drop=attn_drop, proj_drop=proj_drop, norm_layer=norm_layer,
+        )
+    else:
+        raise ValueError(f'Unknown attn_type: {attn_type!r}, expected one of "less", "full", "additive"')
+
+
+class LayerScale(nn.Module):
+    def __init__(
+            self,
+            dim: int,
+            init_values: float = 1e-5,
+            inplace: bool = False,
+    ) -> None:
+        super().__init__()
+        self.inplace = inplace
+        self.gamma = nn.Parameter(init_values * torch.ones(dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x.mul_(self.gamma) if self.inplace else x * self.gamma
+
+
+class LowRankBlock(nn.Module):
+    def __init__(
+            self,
+            dim: int,
+            num_heads: int,
+            channel_dim: int,
+            spatial_dim: int,
+            mlp_ratio: float = 4.,
+            qkv_bias: bool = False,
+            qk_norm: bool = False,
+            proj_drop: float = 0.,
+            attn_drop: float = 0.,
+            drop_path: float = 0.,
+            init_values: Optional[float] = None,
+            act_layer: nn.Module = nn.GELU,
+            norm_layer: nn.Module = nn.LayerNorm,
+            mlp_layer: nn.Module = Mlp,
+            skip_pool: bool = False,
+            rank: int = 1,
+            use_rope_embed: bool = False,
+            attn_type: str = 'less',
+            fusion_init_scale: float = 1.0,
+    ) -> None:
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+
+        self.attn = build_attention(
+            attn_type,
+            dim=dim,
+            channel_dim=channel_dim,
+            spatial_dim=spatial_dim,
+            num_heads=num_heads,
+            rank=rank,
+            qkv_bias=qkv_bias,
+            qk_norm=qk_norm,
+            attn_drop=attn_drop,
+            proj_drop=proj_drop,
+            norm_layer=norm_layer,
+            skip_pool=skip_pool,
+            fusion_init_scale=fusion_init_scale,
+        )
+        self.ls1 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
+        self.drop_path1 = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+
+        self.norm2 = norm_layer(dim)
+        self.mlp = mlp_layer(
+            in_features=dim,
+            hidden_features=int(dim * mlp_ratio),
+            act_layer=act_layer,
+            drop=proj_drop,
+        )
+        self.ls2 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
+        self.drop_path2 = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+
+        self.use_rope_embed = use_rope_embed
+
+    def forward(self, x: torch.Tensor, spatial_mask: torch.Tensor = None, pos_chan_embedding: torch.Tensor = None) -> torch.Tensor:
+        rope = pos_chan_embedding if self.use_rope_embed else None
+
+        x = x + self.drop_path1(self.ls1(self.attn(self.norm1(x), spatial_mask, rope)))
+        x = x + self.drop_path2(self.ls2(self.mlp(self.norm2(x))))
+        return x

config.json

@@ -0,0 +1,52 @@
+{
+  "architectures": [
+    "LESSViTModel"
+  ],
+  "attention_radius": 640,
+  "attn_drop": 0.1,
+  "attn_type": "less",
+  "channel_dim": 24,
+  "channel_dropout": null,
+  "channel_embed_dims_per_head": 2,
+  "channel_mask_ratio": 0.0,
+  "decoder_channel_dim": 64,
+  "decoder_depth": 8,
+  "decoder_embed_dim": 512,
+  "decoder_num_heads": 16,
+  "decoder_out_chans": 1,
+  "decoder_rope_channel_split_dim": 256,
+  "decoder_rope_spatial_split_dim": 256,
+  "decoder_spatial_dim": 128,
+  "depth": 12,
+  "drop_path_rate": 0.1,
+  "drop_path_uniform": false,
+  "embed_dim": 768,
+  "fusion_init_scale": 1.0,
+  "init_values": 1.0,
+  "mask_ratio": 0.0,
+  "mlp_ratio": 4.0,
+  "model_type": "lessvit",
+  "norm_pix_loss": true,
+  "num_heads": 12,
+  "num_tokens": 1,
+  "patch_size": 16,
+  "pos_chan_embed_residual": true,
+  "pretrain": false,
+  "proj_drop": 0.1,
+  "qk_norm": false,
+  "qkv_bias": true,
+  "rank": 1,
+  "return_dict": true,
+  "rope_channel_split_dim": 384,
+  "rope_embed_base": 100.0,
+  "rope_spatial_split_dim": 384,
+  "spatial_dim": 384,
+  "torch_dtype": "float32",
+  "transformers_version": "4.53.3",
+  "use_perception_field_mask": true,
+  "use_rope_embed": true,
+  "auto_map": {
+    "AutoConfig": "modeling_lessvit.LESSViTConfig",
+    "AutoModel": "modeling_lessvit.LESSViTModel"
+  }
+}

hyperspectral_patch_embed.py

@@ -0,0 +1,57 @@
+import torch
+import torch.nn as nn
+import numpy as np
+
+class HyperspectralPatchEmbed(nn.Module):
+    """Image to Patch Embedding."""
+
+    def __init__(
+        self,
+        patch_size: int = 16,
+        embed_dim: int = 768,
+    ):
+        super().__init__()
+        self.patch_size = patch_size
+        
+        # self.proj = nn.Conv3d(
+        #     1,
+        #     embed_dim,
+        #     kernel_size=(1, patch_size, patch_size),
+        #     stride=(1, patch_size, patch_size),
+        # )
+        
+        self.proj = nn.Sequential(
+            nn.Conv3d(
+                1,
+                embed_dim*2, #2
+                kernel_size=(1, patch_size, patch_size),
+                stride=(1, patch_size, patch_size),
+            ),
+            nn.ReLU(),
+            nn.Conv3d(
+                embed_dim*2,
+                embed_dim,
+                kernel_size=(1, 1, 1),
+                stride=(1, 1, 1),
+            ),
+        )
+
+        self.initialize_weights()
+
+    def initialize_weights(self):
+        try:
+            w = self.proj.weight.data
+            torch.nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
+        except:
+            w1 = self.proj[0].weight.data
+            w2 = self.proj[2].weight.data
+            torch.nn.init.xavier_uniform_(w1.view([w1.shape[0], -1]))
+            torch.nn.init.xavier_uniform_(w2.view([w2.shape[0], -1]))
+
+    def forward(self, x):
+        # x: B, Cin, H, W
+        assert x.shape[2] % self.patch_size == 0, "image height should be divisible by patch size"
+        assert x.shape[3] % self.patch_size == 0, "image width should be divisible by patch size"
+        x = self.proj(x.unsqueeze(1)) # B, D, C, H, W
+        B, D, C, H, W = x.shape
+        return x.permute(0, 2, 3, 4, 1).reshape(B, C, H*W, D) # B, C, H, W, D

load_encoder.py

@@ -0,0 +1,41 @@
+"""Load the released encoder and extract features from a synthetic HSI batch."""
+import argparse
+import torch
+from transformers import AutoModel
+
+
+def load_encoder(model_id="GFM-Bench/LESSViT-b2-d8-r1", device="cpu", revision=None):
+    model = AutoModel.from_pretrained(
+        model_id, trust_remote_code=True, revision=revision,
+    )
+    return model.to(device).eval()
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument('--model-id', default='GFM-Bench/LESSViT-b2-d8-r1',
+                        help='Hugging Face model ID or local snapshot directory')
+    parser.add_argument('--revision', default=None, help='Optional commit SHA')
+    parser.add_argument('--device', default='cpu')
+    args = parser.parse_args()
+    torch.set_num_threads(4)
+    encoder = load_encoder(args.model_id, args.device, args.revision)
+    # Synthetic values demonstrate tensor shapes, not preprocessing for real imagery.
+    image = torch.randn(1, 10, 32, 32, device=args.device)
+    wavelengths_nm = torch.linspace(450, 2300, 10, device=args.device)[None]
+    # Example: fixed full-EnMAP wavelength bounds. Use the same reference bounds
+    # for train/eval in your downstream pipeline; do not renormalize each subset.
+    wavelengths = (wavelengths_nm - 418.24) / (2445.53 - 418.24)
+    with torch.no_grad():
+        output = encoder(optical=image, optical_channel_wv=wavelengths,
+                         spatial_resolution=30)
+    grid = output.last_hidden_state
+    assert torch.isfinite(grid).all()
+    print('Token grid:', tuple(grid.shape))
+    print('Global features:', tuple(grid[:, 0, 0].shape))
+    print('Spatial features:', tuple(grid[:, 0, 1:].shape))
+    print('Per-band patch features:', tuple(grid[:, 1:, 1:].shape))
+
+
+if __name__ == '__main__':
+    main()

low_rank_attention.py

@@ -0,0 +1,340 @@
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.checkpoint
+
+from timm.layers import use_fused_attn
+from itertools import product
+from typing import Optional
+
+from .pos_chan_embed import rope_apply
+import math
+
+__all__ = ['AttentionPool', 'AttentionBranch', 'LowRankAttention', 'LowDimPool', 'get_perception_field_mask']
+
+def get_perception_field_mask(num_patches, patch_size, spacial_resolution, attention_radius, cls_token=False):
+    """
+    Create a distance mask for the image.
+    spacial_resolution: the resolution of the image in meters/pixel
+    attention_radius: the radius of the attention in meters
+    """
+    points = list(product(range(int(math.sqrt(num_patches))), range(int(math.sqrt(num_patches)))))
+    idxs = []
+    points_array = torch.tensor(points).to(torch.float32)
+    distances = torch.cdist(points_array, points_array, p=2) * patch_size * spacial_resolution
+    idxs = (distances <= attention_radius).flatten().tolist()
+    mask = torch.tensor(idxs).to(torch.float32).reshape(num_patches, num_patches)
+    if cls_token:   
+        # add a new row and column of ones to the mask for cls_token
+        new_row = torch.ones(num_patches, 1)
+        new_col = torch.ones(1, num_patches + 1)
+        mask = torch.cat([new_row, mask], dim=1)
+        mask = torch.cat([new_col, mask], dim=0)
+    # convert to boolean
+    mask = mask > 0 # HW HW
+    return mask
+
+class AvgPool(nn.Module):
+    def __init__(self, dim: int, dim_out: int):
+        super().__init__()
+        self.linear = nn.Linear(dim, dim_out, bias=False)
+
+    def forward(self, x: torch.Tensor, mask: torch.Tensor = None) -> torch.Tensor:
+        """
+        Args:
+            x (torch.Tensor): input tensor of shape B C N D
+            mask (torch.Tensor): mask tensor of shape B N
+        """
+        B, C, N, D = x.shape
+        if mask is not None:
+            mask = mask.unsqueeze(1).unsqueeze(-1).expand(B, C, N, 1)
+            x = (x * mask).sum(dim=2) / mask.sum(dim=2)
+        else:
+            x = x.mean(dim=2)
+        x = self.linear(x)
+        return x # B C D
+
+class ProjectionPool(nn.Module):
+    def __init__(self, dim: int, dim_out: int, norm_layer: nn.Module = nn.LayerNorm, **kwargs):
+        super().__init__()
+        self.linear = nn.Linear(dim, dim_out, bias=True)
+        self.norm = norm_layer(dim_out)
+
+    def forward(self, x: torch.Tensor, mask: torch.Tensor = None) -> torch.Tensor:
+        # x is B C N D
+        x = x[:, :, 0] # B C D
+        x = self.linear(x)
+        x = self.norm(x)
+        return x # B C D
+
+class AttentionPool(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int = 1,
+        qkv_bias: bool = False,
+        qk_norm: bool = False,
+        attn_drop: float = 0.,
+        norm_layer: nn.Module = nn.LayerNorm,
+        proj_drop: float = 0.,
+    ) -> None:
+        """
+        AttentionPool module for low-rank attention in vision transformers.
+
+        This module implements a pooling mechanism using multi-head attention.
+        It processes input tensors by applying attention across the feature dimension,
+        resulting in a lower-dimensional output. This is particularly useful for
+        reducing the computational complexity in vision transformer architectures.
+
+        Args:
+            dim (int): Input dimension.
+            dim_out (int): Output dimension.
+            num_heads (int): Number of attention heads.
+            qkv_bias (bool): If True, add bias to query, key, value projections.
+            qk_norm (bool): If True, apply normalization to query and key.
+            attn_drop (float): Dropout rate for attention weights.
+            proj_drop (float): Dropout rate for output.
+            norm_layer (nn.Module): Normalization layer to use.
+
+        This method will use the cls token to pool the feature.
+        """
+        super().__init__()
+        assert dim % num_heads == 0, 'dim should be divisible by num_heads'
+        self.head_dim = dim // num_heads
+        self.num_heads = num_heads
+        self.dim = dim
+        self.dim_out = dim
+        self.scale = (dim // num_heads) ** -0.5  # Scaling factor for attention scores
+        self.fused_attn = use_fused_attn()
+        
+        self.q = nn.Linear(dim, dim, bias=qkv_bias)
+        self.kv = nn.Linear(dim, dim*2, bias=qkv_bias)
+        self.q_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
+        self.k_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
+
+        self.attn_drop = nn.Dropout(attn_drop)
+        
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B, C, N, D = x.shape
+            
+        x = x.reshape(B * C, N, D)  # B*C, N+1, D
+        x_cls = x[:, :1, :] # B*C, 1, D
+
+        q = self.q(x_cls).reshape(B * C, 1, self.num_heads, self.head_dim).permute(0, 2, 1, 3) # B*C, num_heads, 1, head_dim
+        kv = self.kv(x).reshape(B * C, N, 2, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4) # 2, B*C, num_heads, N, head_dim
+        k, v = kv.unbind(0) # B*C, num_heads, N, head_dim
+        q, k = self.q_norm(q), self.k_norm(k)
+        
+        if self.fused_attn:
+            x = F.scaled_dot_product_attention(
+                q, k, v,
+                dropout_p=self.attn_drop.p if self.training else 0.,
+            ) # B*C, num_heads, 1, D
+        else:
+            attn = (q @ k.transpose(-2, -1)) * self.scale # B*C, num_heads, N, 1
+            attn = attn.softmax(dim=-1)
+            attn = self.attn_drop(attn)
+            x = (attn @ v).transpose(1, 2) # B*C, num_heads, 1, D
+            
+        x = x.reshape(B, C, self.dim) # B C dim, take only cls token
+        
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+class LowDimPool(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        channel_dim: int,
+        spatial_dim: int,
+        num_heads: int = 8,
+        qkv_bias: bool = False,
+        qk_norm: bool = False,
+        attn_drop: float = 0.,
+        proj_drop: float = 0.,
+        norm_layer: nn.Module = nn.LayerNorm,
+        skip_pool: bool = False,
+    ) -> None:
+        super().__init__()
+        self.num_heads = num_heads
+        
+        self.channel_norm = norm_layer(channel_dim)
+        self.spatial_norm = norm_layer(spatial_dim)
+        # self.norm = norm_layer(dim)
+        
+        if skip_pool:
+            self.channel_pool = lambda x: x[:, :, 0]
+            self.spatial_pool = lambda x: x[:, :, 0]
+        else:
+            self.channel_pool = AttentionPool(dim=channel_dim, num_heads=num_heads, norm_layer=norm_layer, qkv_bias=qkv_bias, qk_norm=qk_norm, attn_drop=attn_drop, proj_drop=proj_drop) # B C N D -> B C dim
+            self.spatial_pool = AttentionPool(dim=spatial_dim, num_heads=num_heads, norm_layer=norm_layer, qkv_bias=qkv_bias, qk_norm=qk_norm, attn_drop=attn_drop, proj_drop=proj_drop) # B HW N D -> B HW dim
+        
+        self.channel_linear = nn.Linear(dim, channel_dim) # B C dim -> B C channel_dim
+        self.spatial_linear = nn.Linear(dim, spatial_dim) # B HW dim -> B HW spatial_dim
+        
+    def forward(self, x: torch.Tensor, pos_chan_embedding: torch.Tensor = None) -> torch.Tensor:
+        # x is B C HW D
+        B, C, HW, D = x.shape
+        x_c = self.channel_linear(x) # B, C, HW, channel_dim
+        x_s = self.spatial_linear(x) # B, C, HW, spatial_dim
+        
+        x_c = x_c[:, :, 0] + self.channel_pool(self.channel_norm(x_c)) # B, C, dim
+        x_s = x_s[:, 0] + self.spatial_pool(self.spatial_norm(x_s.transpose(1, 2))) # B, HW, dim
+        return x_c, x_s, x
+
+class AttentionBranch(nn.Module):
+    def __init__(
+            self,
+            dim: int,
+            num_heads: int,
+            head_dim: int,
+            qkv_bias: bool = False,
+            qk_norm: bool = False,
+            attn_drop: float = 0.,
+            norm_layer: nn.Module = nn.LayerNorm,
+            rank: int = 1,
+    ) -> None:
+        super().__init__()
+        assert dim % num_heads == 0, 'dim should be divisible by num_heads'
+        self.num_heads = num_heads
+        self.head_dim = head_dim
+        self.scale = head_dim ** -0.5
+        self.fused_attn = use_fused_attn() 
+
+        self.qk = nn.Linear(dim, num_heads * head_dim * rank * 2, bias=qkv_bias)
+        self.v = nn.Linear(dim, num_heads * head_dim * rank, bias=qkv_bias)
+        self.q_norm = norm_layer(head_dim) if qk_norm else nn.Identity()
+        self.k_norm = norm_layer(head_dim) if qk_norm else nn.Identity()
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.rank = rank
+    
+    def apply_rope(self, q: torch.Tensor, k: torch.Tensor, sin: Optional[torch.Tensor], cos: Optional[torch.Tensor]):
+        if sin is None and cos is None:
+            return q, k
+        
+        assert sin is not None and cos is not None, "Both sin and cos must be provided for RoPE."
+        q_dtype = q.dtype
+        k_dtype = k.dtype
+        rope_dtype = sin.dtype
+
+        q = q.to(dtype=rope_dtype)
+        k = k.to(dtype=rope_dtype)
+
+        q = rope_apply(q, sin, cos)
+        k = rope_apply(k, sin, cos)
+
+        q = q.to(dtype=q_dtype)
+        k = k.to(dtype=k_dtype)
+
+        return q, k
+        
+    def forward(self, x: torch.Tensor, mask: torch.Tensor = None, sin: torch.Tensor = None, cos: torch.Tensor = None) -> torch.Tensor:
+        B, N, D = x.shape
+        v = self.v(x).reshape(B, N, self.num_heads, self.rank, self.head_dim).permute(0, 3, 2, 1, 4) # B, rank, num_heads, N, D
+        qk = self.qk(x).reshape(B, N, 2, self.num_heads, self.rank, self.head_dim).permute(2, 0, 4, 3, 1, 5) # 2, B, rank, num_heads, N, head_dim
+        q, k = qk.unbind(0)
+        q, k = self.apply_rope(q, k, sin, cos)
+        q, k = self.q_norm(q), self.k_norm(k) #B, rank, num_heads, N, head_dim
+
+        if self.fused_attn:
+            x = F.scaled_dot_product_attention(
+                q, k, v, attn_mask=mask,
+                dropout_p=self.attn_drop.p if self.training else 0.,
+            ) # B, rank, num_heads, N, D
+        else:
+            L, S = q.shape[-2], k.shape[-2]
+            attn = (q @ k.transpose(-2, -1)) * self.scale
+            if mask is not None:
+                atten_bias = torch.zeros(L, S, device=attn.device)
+                atten_bias = atten_bias.masked_fill(mask == 0, -torch.inf)
+                attn = attn + atten_bias
+            attn = attn.softmax(dim=-1)
+            attn = self.attn_drop(attn)
+            x = attn @ v # B, rank, num_heads, N, D
+
+        return x # B, rank, num_heads, N, D
+
+class LowRankAttention(nn.Module):
+    def __init__(
+            self,
+            dim: int,
+            channel_dim: int,
+            spatial_dim: int,
+            num_heads: int = 8,
+            qkv_bias: bool = False,
+            qk_norm: bool = False,
+            attn_drop: float = 0.,
+            proj_drop: float = 0.,
+            norm_layer: nn.Module = nn.LayerNorm,
+            rank: int = 1,
+    ) -> None:
+        """
+        LowRankAttention module for low-rank attention in vision transformers.
+
+        This module implements a low-rank attention mechanism using multi-head attention.
+        It processes input tensors by applying attention across the feature dimension,
+        resulting in a lower-dimensional output. This is particularly useful for
+        reducing the computational complexity in vision transformer architectures.
+
+        Args:
+            dim (int): Input dimension. 
+            channel_dim (int): Channel dimension.
+            spatial_dim (int): Spatial dimension.
+            num_heads (int): Number of attention heads.
+            qkv_bias (bool): If True, add bias to query, key, value projections.
+            qk_norm (bool): If True, apply normalization to query and key.
+            attn_drop (float): Dropout rate for attention weights.
+            proj_drop (float): Dropout rate for output.
+            norm_layer (nn.Module): Normalization layer to use.
+            dim_ratio (float): Dimension ratio for low-rank approximation.
+            pool (bool): If True, use pooling.
+        """
+        super().__init__()
+        assert channel_dim % num_heads == 0, 'channel_dim should be divisible by num_heads'
+        assert spatial_dim % num_heads == 0, 'spatial_dim should be divisible by num_heads'
+        
+        self.num_heads = num_heads
+        self.dim = dim
+        self.channel_dim = channel_dim
+        self.spatial_dim = spatial_dim
+        
+        self.c_head_dim = self.channel_dim // num_heads
+        self.s_head_dim = self.spatial_dim // num_heads
+        self.head_dim = self.dim // num_heads
+        
+        self.rank = rank
+        
+        assert self.head_dim == self.c_head_dim * self.s_head_dim, 'head_dim should be equal to c_head_dim * s_head_dim'
+
+        self.channel_branch = AttentionBranch(channel_dim, num_heads, self.c_head_dim, qkv_bias, qk_norm, attn_drop, norm_layer, rank)
+        self.spatial_branch = AttentionBranch(spatial_dim, num_heads, self.s_head_dim, qkv_bias, qk_norm, attn_drop, norm_layer, rank)
+
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x_c: torch.Tensor, x_s: torch.Tensor, spatial_mask: torch.Tensor = None, rope: tuple = None) -> torch.Tensor:
+        B, C, D = x_c.shape
+        HW = x_s.shape[1]
+
+        sin_hw, cos_hw, sin_c, cos_c = rope if rope is not None else (None, None, None, None)
+        
+        xc = self.channel_branch(x_c, sin=sin_c, cos=cos_c)  # B, rank, num_heads, C, c_head_dim
+        xs = self.spatial_branch(x_s, spatial_mask, sin=sin_hw, cos=cos_hw)  # B, rank, num_heads, HW, s_head_dim
+        x = torch.einsum('...ca,...nb->...cnab', xc, xs).flatten(-2) # B, rank, num_heads, C, HW, D
+        x = x.sum(dim=1) / self.rank # B, num_heads, C, HW, D
+        x = x.permute(0, 2, 3, 1, 4).reshape(B, C, HW, -1) # B, C, HW, D
+        x = self.proj(x) # B, C, HW, D
+        x = self.proj_drop(x) # B, C, HW, D
+
+        return x
+
+# LayerScale and LowRankBlock moved to attention_ops.py, which composes this
+# module's AttentionBranch/LowDimPool/LowRankAttention (arm A) alongside the
+# arm B/C attention operators behind a single build_attention(attn_type, ...)
+# factory.

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:935719728265cf0765b32afe634803ab73071025f104c1c708b98142488b8311
+size 332935996

modeling_lessvit.py

@@ -0,0 +1,37 @@
+"""Hugging Face interface for the original LESSViT encoder."""
+# Explicit dependencies let Transformers copy the full local/remote module tree.
+from .pos_chan_embed import RopePositionChannelEmbedding
+from .hyperspectral_patch_embed import HyperspectralPatchEmbed
+from .low_rank_attention import LowRankAttention
+from .attention_ops import LowRankBlock
+from .spatial_spectral_low_rank_vit import (
+    SpatialSpectralLowRankViTConfig,
+    SpatialSpectralLowRankViTEncoder,
+)
+
+
+class LESSViTConfig(SpatialSpectralLowRankViTConfig):
+    model_type = "lessvit"
+
+    def __init__(self, channel_embed_dims_per_head=2, **kwargs):
+        # The original config constructor recomputes these derived dimensions.
+        kwargs.pop("channel_dim", None)
+        kwargs.pop("spatial_dim", None)
+        kwargs.pop("pretrain", None)
+        kwargs.pop("mask_ratio", None)
+        kwargs.pop("channel_mask_ratio", None)
+        kwargs.pop("channel_dropout", None)
+        super().__init__(channel_embed_dims_per_head=channel_embed_dims_per_head,
+                         mask_ratio=0.0, channel_mask_ratio=0.0,
+                         channel_dropout=None, **kwargs)
+        self.channel_embed_dims_per_head = channel_embed_dims_per_head
+        self.pretrain = False
+
+
+class LESSViTModel(SpatialSpectralLowRankViTEncoder):
+    config_class = LESSViTConfig
+    main_input_name = "optical"
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.post_init()

pos_chan_embed.py

@@ -0,0 +1,361 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+# --------------------------------------------------------
+# Position embedding utils
+# --------------------------------------------------------
+
+from typing import Any
+import numpy as np
+
+import math
+import torch
+import torch.nn as nn
+from torch import Tensor, nn
+from typing import Literal
+
+
+class PositionalChannelEmbedding():
+    def __init__(self, embed_dim):
+        self.embed_dim = embed_dim
+        self.spactial_resolution = None
+        self.channel_ids = None
+        self.num_patches = None
+        self.pos_embed = None
+        self.channel_embed = None
+        self.cls_token = torch.zeros(1, 1, 1, embed_dim)
+        
+
+    def interpolate_pos_channel_embed(self, pos_embed: torch.Tensor, channel_embed: torch.Tensor, cls_token: bool = False):
+        """
+        pos_embed: (1, HW, D)
+        channel_embed: (1, C, D)
+        return pos_embed: (1, C+1, HW+1, D) if cls_token else (1, C, HW, D)
+        """
+        n_chan = channel_embed.shape[1] # C
+        n_pos = pos_embed.shape[1] # HW
+        
+        if cls_token:
+            # Add cls token embeddings
+            pos_cls_embed = torch.zeros(1, 1, self.embed_dim, device=pos_embed.device)
+            pos_embed = torch.cat([pos_cls_embed, pos_embed], dim=1)  # (1, HW+1, D)
+            
+            chan_cls_embed = torch.zeros(1, 1, self.embed_dim, device=channel_embed.device)
+            channel_embed = torch.cat([chan_cls_embed, channel_embed], dim=1)  # (1, C+1, D)
+            
+            n_chan += 1
+            n_pos += 1
+            
+        interpolated_pos_embed = pos_embed.unsqueeze(1).repeat(1, n_chan, 1, 1) # 1 C HW D
+        interpolated_channel_embed = channel_embed.unsqueeze(2).repeat(1, 1, n_pos, 1) # 1 C HW D
+        
+        pos_channel_embed = (interpolated_channel_embed + interpolated_pos_embed) / 2 # 1 C HW D
+        
+        return pos_channel_embed
+    
+    def get_pos_embed(self, tokens: torch.Tensor, spatial_resolution: float, cls_token: bool = True):
+        _, HW, _ = tokens.shape
+        grid_size = int(np.sqrt(HW))
+        assert grid_size * grid_size == HW, "HW must be a square"
+        
+        if self.spactial_resolution == spatial_resolution and self.num_patches == HW:
+            pos_embed = self.pos_embed
+        else:
+            pos_embed = get_2d_sincos_pos_embed(self.embed_dim, grid_size, spatial_resolution, cls_token=False) # (1, HW, D)
+            self.pos_embed = pos_embed
+            self.spactial_resolution = spatial_resolution
+            self.num_patches = HW
+            
+        if cls_token:
+            pos_embed = torch.cat([torch.zeros_like(pos_embed[:, :1, :]), pos_embed], dim=1)
+            
+        return pos_embed
+    
+    def get_channel_embed(self, tokens: torch.Tensor, channel_ids: torch.Tensor, cls_token: bool = True):
+        """
+        channel_ids: (1, C)
+        return channel_embed: (1, C, D)
+        """
+        _, C, _ = tokens.shape
+        channel_ids = channel_ids.squeeze(0).detach().cpu().numpy()
+        if self.channel_ids is not None and tuple(self.channel_ids) == tuple(channel_ids):
+            channel_embed = self.channel_embed
+        else:
+            channel_embed = get_1d_sincos_channel_embed(self.embed_dim, channel_ids, cls_token=False) # (1, C, D)
+            self.channel_embed = channel_embed
+            self.channel_ids = channel_ids
+        
+        if cls_token:
+            channel_embed = torch.cat([torch.zeros_like(channel_embed[:, :1, :]), channel_embed], dim=1)
+        return channel_embed
+        
+    
+    def __call__(self, tokens: torch.Tensor, spatial_resolution: float, channel_ids: torch.Tensor, cls_token: bool = True):
+        _, C, HW, _ = tokens.shape
+        assert channel_ids.shape == (1, C), "channel_ids must be the same length as the number of channels"
+        grid_size = int(np.sqrt(HW))
+        assert grid_size * grid_size == HW, "HW must be a square"
+        channel_ids = channel_ids.squeeze(0).detach().cpu().numpy()
+        
+        if self.channel_ids is not None and tuple(self.channel_ids) == tuple(channel_ids):
+            channel_embed = self.channel_embed
+        else:
+            channel_embed = get_1d_sincos_channel_embed(self.embed_dim, channel_ids, cls_token=False) # (1, C, D)
+            self.channel_embed = channel_embed
+            self.channel_ids = channel_ids
+            
+        if self.spactial_resolution == spatial_resolution and self.num_patches == HW:
+            pos_embed = self.pos_embed
+        else:
+            pos_embed = get_2d_sincos_pos_embed(self.embed_dim, grid_size, spatial_resolution, cls_token=False) # (1, HW, D)
+            self.pos_embed = pos_embed
+            self.spactial_resolution = spatial_resolution
+            self.num_patches = HW
+            
+        return self.interpolate_pos_channel_embed(pos_embed, channel_embed, cls_token)
+
+# --------------------------------------------------------
+# 2D sine-cosine position embedding
+# References:
+# Transformer: https://github.com/tensorflow/models/blob/master/official/nlp/transformer/model_utils.py
+# MoCo v3: https://github.com/facebookresearch/moco-v3
+# --------------------------------------------------------
+def get_2d_sincos_pos_embed(embed_dim, grid_size, spactial_resolution, cls_token=False):
+    """
+    grid_size: int of the grid height and width
+    return:
+    pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
+    """
+    grid_h = np.arange(grid_size, dtype=np.float32)
+    grid_w = np.arange(grid_size, dtype=np.float32)
+    grid = np.meshgrid(grid_w, grid_h)  # here w goes first
+    grid = np.stack(grid, axis=0) * spactial_resolution
+
+    grid = grid.reshape([2, 1, grid_size, grid_size])
+    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
+    if cls_token:
+        pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0)
+    pos_embed = torch.from_numpy(pos_embed).float().unsqueeze(0)
+    return pos_embed
+
+def get_1d_sincos_channel_embed(embed_dim, channel_idx, cls_token=False):
+    """
+    embed_dim: output dimension for each position
+    channel_idx: a list of channel_idx to be encoded: size (C,)
+    out: (C, D)
+    """
+    channel_embed = get_1d_sincos_pos_embed_from_grid(embed_dim, channel_idx)
+    if cls_token:
+        channel_embed = np.concatenate([np.zeros([1, embed_dim]), channel_embed], axis=0)
+    channel_embed = torch.from_numpy(channel_embed).float().unsqueeze(0)
+    return channel_embed
+
+
+def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
+    assert embed_dim % 2 == 0
+
+    # use half of dimensions to encode grid_h
+    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])  # (H*W, D/2)
+    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])  # (H*W, D/2)
+
+    emb = np.concatenate([emb_h, emb_w], axis=1) # (H*W, D)
+    return emb
+
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
+    """
+    embed_dim: output dimension for each position
+    pos: a list of positions to be encoded: size (M,)
+    out: (M, D)
+    """
+    assert embed_dim % 2 == 0
+    omega = np.arange(embed_dim // 2, dtype=np.float32)
+    omega /= embed_dim / 2.
+    omega = 1. / 10000**omega  # (D/2,)
+
+    pos = pos.reshape(-1)  # (M,)
+    out = np.einsum('m,d->md', pos, omega)  # (M, D/2), outer product
+
+    emb_sin = np.sin(out) # (M, D/2)
+    emb_cos = np.cos(out) # (M, D/2)
+
+    emb = np.concatenate([emb_sin, emb_cos], axis=1)  # (M, D)
+    return emb
+
+
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This software may be used and distributed in accordance with
+# the terms of the DINOv3 License Agreement.
+# --------------------------------------------------------
+# RoPE positional embedding with no mixing of coordinates (axial) and no learnable weights
+# Supports two parametrizations of the rope parameters: either using `base` or `min_period` and `max_period`.
+class RopePositionChannelEmbedding(nn.Module):
+    def __init__(
+        self,
+        # embed_dim: int,
+        spatial_dim: int,
+        channel_dim: int,
+        *,
+        num_heads: int,
+        base: float | None = 100.0,
+        min_period: float | None = None,
+        max_period: float | None = None,
+        normalize_coords: Literal["min", "max", "separate"] = "separate",
+        shift_coords: float | None = None,
+        jitter_coords: float | None = None,
+        rescale_coords: float | None = None,
+        dtype: torch.dtype | None = None,
+        device: torch.device | None = None,
+    ):
+        super().__init__()
+        assert spatial_dim % (4 * num_heads) == 0 and spatial_dim % (4 * num_heads) == 0
+        both_periods = min_period is not None and max_period is not None
+        if (base is None and not both_periods) or (base is not None and both_periods):
+            raise ValueError("Either `base` or `min_period`+`max_period` must be provided.")
+
+        D_head_hw = spatial_dim // num_heads
+        D_head_c = channel_dim // num_heads
+        self.base = base
+        self.min_period = min_period
+        self.max_period = max_period
+        self.D_head_hw = D_head_hw
+        self.D_head_c = D_head_c
+        self.normalize_coords = normalize_coords
+        self.shift_coords = shift_coords
+        self.jitter_coords = jitter_coords
+        self.rescale_coords = rescale_coords
+
+        # Needs persistent=True because we do teacher.load_state_dict(student.state_dict()) to initialize the teacher
+        self.dtype = dtype  # Don't rely on self.periods.dtype
+        self.register_buffer(
+            "periods_hw",
+            torch.empty(D_head_hw // 4, device=device, dtype=dtype),
+            persistent=True,
+        )
+        self.register_buffer(
+            "periods_c",
+            torch.empty(D_head_c // 2, device=device, dtype=dtype),
+            persistent=True,
+        )
+        self._init_weights()
+
+    def forward(self, *, H: int, W: int, C: int, optical_channel_wv) -> list[Tensor, ...]:
+        device = self.periods_hw.device
+        dtype = self.dtype
+        dd = {"device": device, "dtype": dtype}
+
+        # Prepare coords in range [-1, +1]
+        if self.normalize_coords == "max":
+            max_HW = max(H, W)
+            coords_h = torch.arange(0.5, H, **dd) / max_HW  # [H]
+            coords_w = torch.arange(0.5, W, **dd) / max_HW  # [W]
+        elif self.normalize_coords == "min":
+            min_HW = min(H, W)
+            coords_h = torch.arange(0.5, H, **dd) / min_HW  # [H]
+            coords_w = torch.arange(0.5, W, **dd) / min_HW  # [W]
+        elif self.normalize_coords == "separate":
+            coords_h = torch.arange(0.5, H, **dd) / H  # [H]
+            coords_w = torch.arange(0.5, W, **dd) / W  # [W]
+        else:
+            raise ValueError(f"Unknown normalize_coords: {self.normalize_coords}")
+        coords_hw = torch.stack(torch.meshgrid(coords_h, coords_w, indexing="ij"), dim=-1)  # [H, W, 2]
+        coords_hw = coords_hw.flatten(0, 1)  # [HW, 2]
+        coords_hw = 2.0 * coords_hw - 1.0  # Shift range [0, 1] to [-1, +1]
+
+        # coords_c = torch.arange(0.5, C, **dd) / C # [C]
+        # coords_c = 2.0 * coords_c - 1.0
+        coords_c = 2 * optical_channel_wv.squeeze() - 1.0 # [C]
+
+        # Shift coords by adding a uniform value in [-shift, shift]
+        if self.training and self.shift_coords is not None:
+            shift_hw = torch.empty(2, **dd).uniform_(-self.shift_coords, self.shift_coords)
+            coords_hw += shift_hw[None, :]
+            
+            shift_c = torch.empty(1, **dd).uniform_(-self.shift_coords, self.shift_coords)
+            coords_c += shift_c
+
+        # Jitter coords by multiplying the range [-1, 1] by a log-uniform value in [1/jitter, jitter]
+        if self.training and self.jitter_coords is not None:
+            jitter_max = np.log(self.jitter_coords)
+            jitter_min = -jitter_max
+            jitter_hw = torch.empty(2, **dd).uniform_(jitter_min, jitter_max).exp()
+            coords_hw *= jitter_hw[None, :]
+
+            jitter_c = torch.empty(1, **dd).uniform_(jitter_min, jitter_max).exp()
+            coords_c *= jitter_c
+
+        # Rescale coords by multiplying the range [-1, 1] by a log-uniform value in [1/rescale, rescale]
+        if self.training and self.rescale_coords is not None:
+            rescale_max = np.log(self.rescale_coords)
+            rescale_min = -rescale_max
+            rescale = torch.empty(1, **dd).uniform_(rescale_min, rescale_max).exp()
+            coords_hw *= rescale_hw
+            coords_c *= rescale_c
+
+        # Prepare angles and sin/cos
+        angles_hw = 2 * math.pi * coords_hw[:, :, None] / self.periods_hw[None, None, :]  # [HW, 2, D//4]
+        angles_hw = angles_hw.flatten(1, 2)  # [HW, D//2]
+        angles_hw = angles_hw.tile(2)  # [HW, D]
+        cos_hw = torch.cos(angles_hw)  # [HW, D]
+        sin_hw = torch.sin(angles_hw)  # [HW, D]
+
+        angles_c = 2 * math.pi * coords_c[:, None] / self.periods_c[None, :]  # [C, D//2]
+        angles_c = angles_c.tile(2)  # [C, D]
+        cos_c = torch.cos(angles_c)  # [C, D]
+        sin_c = torch.sin(angles_c)  # [C, D]
+
+        # rope = {
+        #     'spatial_rope': (sin_hw, cos_hw), # 2 * [HW, D]
+        #     'spectral_rope': (sin_c, cos_c) # 2 * [C, D]
+        # }
+        # return (sin, cos)  # 2 * [HW, D]
+        return [sin_hw, cos_hw, sin_c, cos_c]
+
+    def _init_weights(self):
+        device = self.periods_hw.device
+        dtype = self.dtype
+        if self.base is not None:
+            periods_hw = self.base ** (
+                2 * torch.arange(self.D_head_hw // 4, device=device, dtype=dtype) / (self.D_head_hw // 2)
+            )  # [D//4]
+        else:
+            base = self.max_period / self.min_period
+            exponents = torch.linspace(0, 1, self.D_head_hw // 4, device=device, dtype=dtype)  # [D//4] range [0, 1]
+            periods_hw = base**exponents  # range [1, max_period / min_period]
+            periods_hw = periods_hw / base  # range [min_period / max_period, 1]
+            periods_hw = periods_hw * self.max_period  # range [min_period, max_period]
+
+        if self.base is not None:
+            periods_c = self.base ** (
+                2 * torch.arange(self.D_head_c // 2, device=device, dtype=dtype) / (self.D_head_c)
+            )  # [D//2]
+        else:
+            base = self.max_period / self.min_period
+            exponents = torch.linspace(0, 1, self.D_head_c // 2, device=device, dtype=dtype)
+            periods_c = base**exponents
+            periods_c = periods_c / base
+            periods_c = periods_c * self.max_period
+
+        self.periods_hw.data = periods_hw
+        self.periods_c.data = periods_c
+
+# RoPE-related functions:
+def rope_rotate_half(x: Tensor) -> Tensor:
+    # x:   [ x0  x1  x2  x3  x4  x5]
+    # out: [-x3 -x4 -x5  x0  x1  x2]
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat([-x2, x1], dim=-1)
+
+
+def rope_apply(x: Tensor, sin: Tensor, cos: Tensor) -> Tensor:
+    # x:   [..., D], eg [x0,     x1,   x2,   x3,   x4,   x5]
+    # sin: [..., D], eg [sin0, sin1, sin2, sin0, sin1, sin2]
+    # cos: [..., D], eg [cos0, cos1, cos2, cos0, cos1, cos2]
+    x_ = x[:, :, :, 1:, :]
+    # print(x_.shape, cos.unsqueeze(1).unsqueeze(1).shape)
+    x_ = (x_ * cos.unsqueeze(1).unsqueeze(1)) + (rope_rotate_half(x_) * sin.unsqueeze(1).unsqueeze(1))
+    x = torch.cat((x[:, :, :, :1, :], x_), dim=3)
+    return x

pretraining/architecture.json

@@ -0,0 +1,50 @@
+{
+  "architectures": [
+    "SpatialSpectralMAEViT"
+  ],
+  "attention_radius": 640,
+  "attn_drop": 0.1,
+  "channel_dim": 24,
+  "channel_dropout": [
+    0.75,
+    0.85
+  ],
+  "channel_mask_ratio": 0.75,
+  "decoder_attention_radius": [
+    640
+  ],
+  "decoder_channel_dim": 32,
+  "decoder_depth": 8,
+  "decoder_embed_dim": 512,
+  "decoder_history_length": [
+    4
+  ],
+  "decoder_num_heads": 16,
+  "decoder_out_chans": 155,
+  "decoder_out_dims": 1,
+  "decoder_spatial_dim": 256,
+  "decoder_temporal_embed_dims_per_head": 4,
+  "depth": 12,
+  "drop_path_rate": 0.1,
+  "drop_path_uniform": false,
+  "embed_dim": 768,
+  "init_values": 1.0,
+  "mask_ratio": 0.75,
+  "mlp_ratio": 4.0,
+  "model_type": "multi_modal_low_rank_vit",
+  "norm_pix_loss": false,
+  "num_heads": 12,
+  "patch_size": 16,
+  "pos_chan_embed_residual": true,
+  "pretrain": true,
+  "proj_drop": 0.1,
+  "qk_norm": false,
+  "qkv_bias": true,
+  "rank": 1,
+  "return_dict": false,
+  "rope_embed_base": 100.0,
+  "spatial_dim": 384,
+  "use_perception_field_mask": true,
+  "use_rope_embed": true,
+  "channel_embed_dims_per_head": 2
+}

pretraining/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8e5dce5229b24255d31dc619c005ad6355784eff7ee9e07036aba6f36a1c606a
+size 430541280

requirements.txt

@@ -0,0 +1,7 @@
+torch==2.7.1
+torchvision==0.22.1
+transformers==4.53.3
+huggingface-hub==0.34.4
+timm==1.0.18
+numpy==2.2.6
+safetensors>=0.4.3,<1

spatial_spectral_low_rank_vit.py

@@ -0,0 +1,704 @@
+import torch
+import random
+import torch.nn as nn
+from transformers import PretrainedConfig, PreTrainedModel
+from transformers.modeling_outputs import BaseModelOutput
+from typing import Optional, Tuple, Dict, Any, List
+from functools import partial
+import numpy as np
+from .hyperspectral_patch_embed import HyperspectralPatchEmbed
+from .pos_chan_embed import PositionalChannelEmbedding, RopePositionChannelEmbedding
+from .low_rank_attention import get_perception_field_mask
+from .attention_ops import LowRankBlock
+import torch.nn.functional as F
+import math
+
+class SpatialSpectralLowRankViTConfig(PretrainedConfig):
+    model_type = "multi_modal_low_rank_vit"
+
+    def __init__(
+        self,
+        patch_size: int = 16,
+        embed_dim: int = 768,
+        channel_embed_dims_per_head: int = 4,
+        depth: int = 12,
+        num_heads: int = 12,
+        mlp_ratio: float = 4.0,
+        qkv_bias: bool = True,
+        qk_norm: bool = False,
+        drop_path_rate: float = 0.0,
+        drop_path_uniform: bool = False,
+        init_values: Optional[float] = None,
+        attn_drop: float = 0.0,
+        proj_drop: float = 0.0,
+        mask_ratio: float = 0.75,
+        channel_mask_ratio: float = 0.5,
+        pos_chan_embed_residual: bool = True,
+        rank: int = 1,
+        # Decoder-specific parameters
+        decoder_embed_dim: int = 512,
+        decoder_depth: int = 8,
+        decoder_channel_embed_dims_per_head: int = 4,
+        decoder_num_heads: int = 16,
+        decoder_out_chans: int = 1,
+        # return dict
+        return_dict: bool = False,
+        norm_pix_loss: bool = True,
+        use_perception_field_mask: bool = False,
+        attention_radius: int = 640,
+        use_rope_embed: bool = False,
+        rope_embed_base: float = 100.0,
+        channel_dropout: Optional[List[float]] = None,
+        classes: Optional[int] = None,
+        attn_type: str = "less",
+        fusion_init_scale: float = 1.0,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        assert attn_type in ("less", "full", "additive"), f"Unknown attn_type: {attn_type!r}"
+        self.attn_type = attn_type
+        self.fusion_init_scale = fusion_init_scale
+        self.patch_size = patch_size
+        self.embed_dim = embed_dim
+        # arm C (additive) reuses arm A's channel_dim/spatial_dim derivation
+        # unchanged -- it keeps the same low-rank branch widths as arm A and only
+        # projects up to head_dim at the fusion step (see
+        # AdditiveSpatialSpectralAttention), so no attn_type-specific override is
+        # needed here (unlike an earlier version of this design).
+        self.channel_dim = channel_embed_dims_per_head * num_heads
+        self.spatial_dim = embed_dim // self.channel_dim * num_heads
+        # Joint-RoPE head-dim partition for attn_type == "full": splits head_dim
+        # (embed_dim // num_heads) into a spatial (u,v) sub-block and a spectral
+        # (lambda) sub-block, default 50/50. Unused by "less"/"additive".
+        head_dim = embed_dim // num_heads
+        rope_spatial_split = head_dim // 2
+        rope_spectral_split = head_dim - rope_spatial_split
+        self.rope_spatial_split_dim = rope_spatial_split * num_heads
+        self.rope_channel_split_dim = rope_spectral_split * num_heads
+        self.depth = depth
+        self.num_heads = num_heads
+        self.mlp_ratio = mlp_ratio
+        self.qkv_bias = qkv_bias
+        self.qk_norm = qk_norm
+        self.drop_path_rate = drop_path_rate
+        self.drop_path_uniform = drop_path_uniform
+        self.init_values = init_values
+        self.attn_drop = attn_drop
+        self.proj_drop = proj_drop
+        self.num_tokens = 1
+        self.return_dict = return_dict
+        self.mask_ratio = mask_ratio
+        self.channel_mask_ratio = channel_mask_ratio
+        self.pretrain = True
+        self.rank = rank
+        # Decoder-specific attributes
+        self.decoder_embed_dim = decoder_embed_dim
+        self.decoder_channel_dim = decoder_channel_embed_dims_per_head * decoder_num_heads
+        self.decoder_spatial_dim = decoder_embed_dim // self.decoder_channel_dim * decoder_num_heads
+        decoder_head_dim = decoder_embed_dim // decoder_num_heads
+        decoder_rope_spatial_split = decoder_head_dim // 2
+        decoder_rope_spectral_split = decoder_head_dim - decoder_rope_spatial_split
+        self.decoder_rope_spatial_split_dim = decoder_rope_spatial_split * decoder_num_heads
+        self.decoder_rope_channel_split_dim = decoder_rope_spectral_split * decoder_num_heads
+        self.decoder_depth = decoder_depth
+        self.decoder_num_heads = decoder_num_heads
+        self.decoder_out_chans = decoder_out_chans
+        
+        # MAE-specific attributes
+        self.norm_pix_loss = norm_pix_loss
+        
+        # Perception field mask
+        self.use_perception_field_mask = use_perception_field_mask
+        self.attention_radius = attention_radius
+        
+        # Positional channel embedding residual
+        self.pos_chan_embed_residual = pos_chan_embed_residual
+
+        # RoPe embedding
+        self.use_rope_embed = use_rope_embed
+        self.rope_embed_base = rope_embed_base
+
+        self.channel_dropout = channel_dropout
+
+    @property
+    def encoder_config(self):
+        return {k: v for k, v in self.__dict__.items() if not k.startswith('decoder_')}
+
+    @property
+    def decoder_config(self):
+        return {
+            'embed_dim': self.decoder_embed_dim,
+            'depth': self.decoder_depth,
+            'num_heads': self.decoder_num_heads,
+            'out_chans': self.decoder_out_chans,
+            'mlp_ratio': self.mlp_ratio,
+            'qkv_bias': self.qkv_bias,
+            'qk_norm': self.qk_norm,
+            'drop_path_rate': self.drop_path_rate,
+            'drop_path_uniform': self.drop_path_uniform,
+            'init_values': self.init_values,
+            'attn_drop': self.attn_drop,
+            'proj_drop': self.proj_drop,
+            'channel_dim': self.channel_dim,
+            'spatial_dim': self.spatial_dim,
+            'patch_size': self.patch_size,
+            'return_dict': self.return_dict,
+        }
+
+class SpatialSpectralLowRankViTEncoder(PreTrainedModel):
+    config_class = SpatialSpectralLowRankViTConfig
+
+    def __init__(self, config: SpatialSpectralLowRankViTConfig):
+        super().__init__(config)
+        self.config = config
+        norm_layer = partial(nn.LayerNorm, eps=1e-6)
+        
+        # Patch embedding layers for optical and radar inputs
+        self.optical_patch_embed = HyperspectralPatchEmbed(config.patch_size, config.embed_dim)
+        self.radar_patch_embed = HyperspectralPatchEmbed(config.patch_size, config.embed_dim)
+        
+        # Learnable tokens and embeddings
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, config.embed_dim))
+        self.spatial_cls_token = nn.Parameter(torch.zeros(1, 1, config.embed_dim))
+        self.channel_cls_token = nn.Parameter(torch.zeros(1, 1, config.embed_dim))
+        if not config.use_rope_embed:
+            self.pos_chan_embed = PositionalChannelEmbedding(config.embed_dim)
+        else:
+            # attn_type == "full" rotates a joint head-dim split (see
+            # FullSpatialSpectralAttention) rather than arms A/C's per-branch widths.
+            rope_spatial_dim = config.rope_spatial_split_dim if config.attn_type == "full" else config.spatial_dim
+            rope_channel_dim = config.rope_channel_split_dim if config.attn_type == "full" else config.channel_dim
+            self.pos_chan_embed = RopePositionChannelEmbedding(
+                # embed_dim=config.embed_dim,
+                spatial_dim=rope_spatial_dim,
+                channel_dim=rope_channel_dim,
+                num_heads=config.num_heads,
+                base=config.rope_embed_base if hasattr(config, 'rope_embed_base') else 100.0
+            )
+
+        if config.drop_path_uniform is True:
+            dpr = [config.drop_path_rate] * config.depth
+        else:
+            dpr = [x.item() for x in torch.linspace(0, config.drop_path_rate, config.depth)]
+
+        # Create transformer blocks
+        self.blocks = nn.ModuleList([
+            LowRankBlock(
+                dim=config.embed_dim,
+                num_heads=config.num_heads,
+                channel_dim=config.channel_dim,
+                spatial_dim=config.spatial_dim,
+                mlp_ratio=config.mlp_ratio,
+                qkv_bias=config.qkv_bias,
+                qk_norm=config.qk_norm,
+                proj_drop=config.proj_drop,
+                attn_drop=config.attn_drop,
+                drop_path=dpr[i],
+                init_values=config.init_values,
+                norm_layer=norm_layer,
+                rank=config.rank,
+                use_rope_embed=config.use_rope_embed,
+                attn_type=config.attn_type,
+                fusion_init_scale=config.fusion_init_scale,
+            )
+            for i in range(config.depth)
+        ])
+        
+        # Final normalization layer
+        self.norm = norm_layer(config.embed_dim)
+        self.head = nn.Identity()
+        
+        # Initialize weights
+        self.initialize_weights()
+
+    def initialize_weights(self):
+        if self.config.use_rope_embed:
+            assert isinstance(self.pos_chan_embed, RopePositionChannelEmbedding)
+            self.pos_chan_embed._init_weights()
+
+        torch.nn.init.normal_(self.cls_token, std=.02)
+        torch.nn.init.normal_(self.spatial_cls_token, std=.02)
+        torch.nn.init.normal_(self.channel_cls_token, std=.02)
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            torch.nn.init.xavier_uniform_(m.weight)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def forward(self, optical=None, radar=None, optical_channel_wv=None, radar_channel_wv=None, spatial_resolution=10, mask_ratio=None, channel_mask_ratio=None):
+        if self.config.channel_dropout is not None and not self.training:
+            assert len(self.config.channel_dropout) in [1, 2], f"channel_dropout should be a float or a list of two floats, but got {self.config.channel_dropout}"
+            for i in range(len(self.config.channel_dropout)):
+                assert 0.0 <= self.config.channel_dropout[i] < 1.0, f"channel_dropout should be between 0 and 1, but got {self.config.channel_dropout}"
+            channel_dropout = sorted(self.config.channel_dropout)
+            num_channels = optical.shape[1] if optical is not None else radar.shape[1]
+            
+            if len(channel_dropout) == 1:
+                C = max(1, int(num_channels * (1 - channel_dropout[0])))
+            else:
+                C_min = max(1, int(num_channels * (1 - channel_dropout[1])))
+                C_max = max(1, int(num_channels * (1 - channel_dropout[0])))
+                C = random.randint(C_min, C_max)
+            idx_to_keep = sorted(random.sample(range(num_channels), C))
+            if optical is not None:
+                optical = optical[:, idx_to_keep, :, :]
+                optical_channel_wv = optical_channel_wv[:, idx_to_keep]
+            if radar is not None:
+                radar = radar[:, idx_to_keep, :, :]
+                radar_channel_wv = radar_channel_wv[:, idx_to_keep]
+        assert optical is not None or radar is not None, "At least one of optical and radar should be provided"
+        assert optical_channel_wv is None or len(optical_channel_wv.shape) == 2, "If optical ids are provided, they should be a 2D tensor"
+        assert radar_channel_wv is None or len(radar_channel_wv.shape) == 2, "If radar ids are provided, they should be a 2D tensor"
+        mask_ratio = self.config.mask_ratio if mask_ratio is None else mask_ratio
+        channel_mask_ratio = self.config.channel_mask_ratio if channel_mask_ratio is None else channel_mask_ratio
+        hidden_states = []
+        
+        dummy_loss = 0 # dummy loss to avoid empty loss error
+        if optical is not None:
+            optical = self.optical_patch_embed(optical)  # B, Co, HW, D
+            assert optical_channel_wv is not None, "Optical ids should be provided"
+            assert optical_channel_wv.shape[1] == optical.shape[1], "Optical ids should have the same number of channels as the optical data"
+        else:
+            optical_channel_wv = None
+            for param in self.optical_patch_embed.parameters():
+                dummy_loss += 0.0 * torch.sum(param)
+        if radar is not None:
+            radar = self.radar_patch_embed(radar)  # B, Cr, HW, D
+            assert radar_channel_wv is not None, "Radar ids should be provided"
+            assert radar_channel_wv.shape[1] == radar.shape[1], "Radar ids should have the same number of channels as the radar data"
+        else:
+            radar_channel_wv = None
+            for param in self.radar_patch_embed.parameters():
+                dummy_loss += 0.0 * torch.sum(param)
+
+        channel_ids = self.maybe_concat(optical_channel_wv, radar_channel_wv) # 1, C = Co + Cr    
+        x = self.maybe_concat(optical, radar) # B, C, HW, D
+        num_patches = x.shape[2]
+        
+        # Create the perception field mask to simulate convolutional attention
+        if self.config.use_perception_field_mask:
+            perception_field_mask = get_perception_field_mask(num_patches, self.config.patch_size, spatial_resolution, attention_radius=self.config.attention_radius, cls_token=False).to(x.device)
+        else:
+            perception_field_mask = None
+        
+        B, C, HW, D = x.shape
+        H = W = int(HW ** 0.5)
+        # Create cls tokens for both Cin and HW dimensions
+        channel_cls_token = self.channel_cls_token.expand(B, 1, HW, D)
+        spatial_cls_token = self.spatial_cls_token.expand(B, C, 1, D)
+        cls_token = self.cls_token.expand(B, 1, 1, D)
+        spatial_cls_token = torch.cat((cls_token, spatial_cls_token), dim=1) # B C+1 HW D
+        
+        # Append cls token to Cin dimension
+        x = torch.cat((channel_cls_token, x), dim=1)
+        
+        # Append cls token to HW dimension
+        x = torch.cat((spatial_cls_token, x), dim=2)
+
+        # Add positional and channel embedding
+        if not self.config.use_rope_embed:
+            pos_chan_embed = self.pos_chan_embed(x[:, 1:, 1:, :], channel_ids=channel_ids, spatial_resolution=spatial_resolution, cls_token=True).to(x.device, dtype=x.dtype)
+            pos_chan_embed = pos_chan_embed.repeat(B, 1, 1, 1)
+        else:
+            pos_chan_embed = self.pos_chan_embed(H=H, W=W, C=C, optical_channel_wv=channel_ids)  # 
+            # below is current walkaround to make RoPE compatibale with masking
+            pos_chan_embed = [embed.repeat(B, 1, 1) for embed in pos_chan_embed]
+        
+        # Apply masks after positional embedding
+        if self.training and self.config.pretrain:
+            # mask the channel
+            x_ = x[:, 1:, :, :] # B C HW+1 D
+            pos_chan_embed_ = pos_chan_embed[:, 1:, :, :] if not self.config.use_rope_embed else pos_chan_embed
+            x_, pos_chan_embed_, channel_mask, channel_ids_restore = self.random_channel_masking(x_, pos_chan_embedding=pos_chan_embed_, channel_mask_ratio=channel_mask_ratio, use_rope_embed=self.config.use_rope_embed) # B N HW+1 D, B C, B C
+            x = torch.cat((x[:, :1, :, :], x_), dim=1) # B N+1 HW+1 D
+            pos_chan_embed = torch.cat((pos_chan_embed[:, :1, :, :], pos_chan_embed_), dim=1) if not self.config.use_rope_embed else pos_chan_embed_ # B N+1 HW+1 D
+            # mask the position
+            x_ = x[:, :, 1:, :] # B N+1 HW D
+            pos_chan_embed_ = pos_chan_embed[:, :, 1:, :] if not self.config.use_rope_embed else pos_chan_embed
+            x_, pos_chan_embed_, pos_mask, pos_ids_restore, perception_field_mask = self.random_pos_masking(x_, pos_chan_embedding=pos_chan_embed_, mask_ratio=mask_ratio, perception_field_mask=perception_field_mask, use_rope_embed=self.config.use_rope_embed) # B N+1 L D, B HW, B HW, L L
+            x = torch.cat((x[:, :, :1, :], x_), dim=2) # B N+1 L+1 D
+            pos_chan_embed = torch.cat((pos_chan_embed[:, :, :1, :], pos_chan_embed_), dim=2) if not self.config.use_rope_embed else pos_chan_embed_ # B N+1 L+1 D
+        else:
+            channel_mask = pos_mask = channel_ids_restore = pos_ids_restore = None
+            
+        x = x + pos_chan_embed if not self.config.use_rope_embed else x # B N+1 L+1 D
+        pos_chan_embed = tuple(pos_chan_embed) if self.config.use_rope_embed else pos_chan_embed
+
+        if perception_field_mask is not None:
+            # append cls token
+            L = perception_field_mask.shape[0]
+            new_row = torch.ones(L, 1, device=perception_field_mask.device)
+            new_col = torch.ones(1, L + 1, device=perception_field_mask.device)
+            perception_field_mask = torch.cat([new_row, perception_field_mask], dim=1)
+            perception_field_mask = torch.cat([new_col, perception_field_mask], dim=0)
+            perception_field_mask = perception_field_mask > 0 # L L bool
+
+        # Apply transformer blocks
+        for i,blk in enumerate(self.blocks):
+            if not self.config.use_rope_embed:
+                pos_chan_embedding = pos_chan_embed*2**(-i) if self.config.pos_chan_embed_residual else pos_chan_embed
+            else:
+                pos_chan_embedding = pos_chan_embed
+            x = blk(x, spatial_mask=perception_field_mask, pos_chan_embedding=pos_chan_embedding)
+
+        # Apply final layer norm
+        x = self.norm(x)  # B N+1 L+1 D
+        cls_token = x[:, 0, 0]
+        patch_tokens = x[:, 1:, 1:] # B N+1 L+1 D -> B N L D
+
+        if self.training and self.config.pretrain:
+            return x + dummy_loss, channel_mask, channel_ids_restore, pos_mask, pos_ids_restore
+        elif self.config.return_dict:
+            return BaseModelOutput(
+                last_hidden_state=x,
+                hidden_states=hidden_states,
+                attentions=None,
+            )
+        else:
+            return x, cls_token, patch_tokens, hidden_states
+
+    def maybe_concat(self, x, y):
+        if x is not None and y is not None:
+            return torch.cat((x, y), dim=1)
+        elif x is not None:
+            return x
+        elif y is not None:
+            return y
+        
+    def random_pos_masking(self, x, pos_chan_embedding=None, mask_ratio=0.75, batch_wise_mask=False, perception_field_mask=None, use_rope_embed=False):
+        """
+        Perform per-sample random masking by per-sample shuffling.
+        Per-sample shuffling is done by argsort random noise.
+        x: [B, C, HW, D], sequence
+        perception_field_mask: [HW, HW] or None, if not None, use distance attention mask is used.
+        """
+        if perception_field_mask is not None:
+            batch_wise_mask = True # if distance mask is provided, we need to use batch-wise masking
+            
+        B, C, HW, D = x.shape  # batch, channels, length, dim
+        if mask_ratio == 0:
+            mask = torch.zeros([B, HW], device=x.device)
+            ids_restore = torch.arange(HW, device=x.device).unsqueeze(0).repeat(B, 1)
+            if pos_chan_embedding is not None:
+                return x, pos_chan_embedding, mask, ids_restore, perception_field_mask
+            else:
+                return x, mask, ids_restore, perception_field_mask
+        
+        x = x.permute(0, 2, 1, 3) # B, HW, C, D
+        len_keep = int(HW * (1 - mask_ratio)) # L
+        
+        if batch_wise_mask:
+            noise = torch.rand(1, HW, device=x.device)  # noise in [0, 1]
+            noise = noise.repeat(B, 1)
+        else:
+            noise = torch.rand(B, HW, device=x.device)  # noise in [0, 1]
+        
+        # sort noise for each sample
+        ids_shuffle = torch.argsort(noise, dim=1)  # ascend: small is keep, large is remove
+        ids_restore = torch.argsort(ids_shuffle, dim=1) # B, HW
+
+        # keep the first subset
+        ids_keep = ids_shuffle[:, :len_keep] # B, L
+        x_masked = torch.gather(x, dim=1, index=ids_keep.unsqueeze(-1).unsqueeze(-1).repeat(1, 1, C, D)) # B, L, C, D
+        
+        if pos_chan_embedding is not None:
+            if use_rope_embed:
+                pos_chan_embedding[0] = torch.gather(pos_chan_embedding[0], dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, pos_chan_embedding[0].shape[-1])) # B, L, D
+                pos_chan_embedding[1] = torch.gather(pos_chan_embedding[1], dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, pos_chan_embedding[1].shape[-1])) # B, L, D
+                pos_chan_embed_masked = pos_chan_embedding
+            else:
+                pos_chan_embedding = pos_chan_embedding.permute(0, 2, 1, 3) # B, HW, C, D
+                pos_chan_embed_masked = torch.gather(pos_chan_embedding, dim=1, index=ids_keep.unsqueeze(-1).unsqueeze(-1).repeat(1, 1, C, D)) # B, L, C, D
+                pos_chan_embed_masked = pos_chan_embed_masked.permute(0, 2, 1, 3) # B, C, L, D
+        
+        if perception_field_mask is not None:
+            # remove the corresponding ids in perception_field_mask
+            ids_keep_ = ids_keep[0].cpu() # L
+            perception_field_mask = perception_field_mask[ids_keep_, :] # L HW
+            perception_field_mask = perception_field_mask[:, ids_keep_] # L L
+
+        # generate the binary mask: 0 is keep, 1 is remove
+        mask = torch.ones([B, HW], device=x.device)
+        mask[:, :len_keep] = 0
+        # unshuffle to get the binary mask
+        mask = torch.gather(mask, dim=1, index=ids_restore) # B, HW
+        
+        x_masked = x_masked.permute(0, 2, 1, 3)  # [B, C, L, D]
+        if pos_chan_embedding is not None:
+            return x_masked, pos_chan_embed_masked, mask, ids_restore, perception_field_mask
+        else:
+            return x_masked, mask, ids_restore, perception_field_mask
+    
+    def random_channel_masking(self, x, pos_chan_embedding=None, channel_mask_ratio=0.5, batch_wise_mask=False, use_rope_embed=False):
+        """
+        Perform per-sample random masking by per-sample shuffling.
+        Per-sample shuffling is done by argsort random noise.
+        x: [B, C, HW, D], sequence
+        pos_chan_embedding: [B, C, HW, D], positional embedding
+        Return:
+            x_masked: [B, N, HW, D]
+            mask: [B, C]
+            ids_restore: [B, C]
+        """
+        B, C, HW, D = x.shape
+        if channel_mask_ratio == 0:
+            mask = torch.zeros([B, C], device=x.device)
+            ids_restore = torch.arange(C, device=x.device).unsqueeze(0).repeat(B, 1)
+            if pos_chan_embedding is not None:
+                return x, pos_chan_embedding, mask, ids_restore
+            else:
+                return x, mask, ids_restore
+        
+        len_keep = max(int(C * (1 - channel_mask_ratio)), 2) # at least 2 channel is kept, N
+        
+        if batch_wise_mask:
+            noise = torch.rand(1, C, device=x.device)  # noise in [0, 1]
+            noise = noise.repeat(B, 1)
+        else:
+            noise = torch.rand(B, C, device=x.device) # noise in [0, 1]
+        
+        # sort noise for each sample
+        ids_shuffle = torch.argsort(noise, dim=1) # B C
+        ids_restore = torch.argsort(ids_shuffle, dim=1) # B C
+
+        # keep the first subset
+        ids_keep = ids_shuffle[:, :len_keep]
+        x_masked = torch.gather(x, dim=1, index=ids_keep.unsqueeze(-1).unsqueeze(-1).repeat(1, 1, HW, D)) # B, N, HW, D
+        if pos_chan_embedding is not None:
+            if use_rope_embed:
+                pos_chan_embedding[2] = torch.gather(pos_chan_embedding[2], dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, pos_chan_embedding[2].shape[-1])) # B, L, D
+                pos_chan_embedding[3] = torch.gather(pos_chan_embedding[3], dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, pos_chan_embedding[3].shape[-1])) # B, L, D
+                pos_chan_embed_masked = pos_chan_embedding
+            else:
+                pos_chan_embed_masked = torch.gather(pos_chan_embedding, dim=1, index=ids_keep.unsqueeze(-1).unsqueeze(-1).repeat(1, 1, HW, D)) # B, N, HW, D
+
+        # generate the binary mask: 0 is keep, 1 is remove
+        mask = torch.ones([B, C], device=x.device) # B C
+        mask[:, :len_keep] = 0
+        # unshuffle to get the binary mask
+        mask = torch.gather(mask, dim=1, index=ids_restore) # B C
+    
+        if pos_chan_embedding is not None:
+            return x_masked, pos_chan_embed_masked, mask, ids_restore
+        else:
+            return x_masked, mask, ids_restore
+
+class SpatialSpectralLowRankViTDecoder(PreTrainedModel):
+    config_class = SpatialSpectralLowRankViTConfig
+
+    def __init__(self, config: SpatialSpectralLowRankViTConfig):
+        super().__init__(config)
+        self.config = config
+        norm_layer = partial(nn.LayerNorm, eps=1e-6)
+        
+        self.decoder_embed = nn.Linear(config.embed_dim, config.decoder_embed_dim, bias=True)
+        # self.decoder_embed_norm = norm_layer(config.decoder_embed_dim)
+        self.mask_token = nn.Parameter(torch.zeros(1, 1, 1, config.decoder_embed_dim))
+        self.spatial_mask_token = nn.Parameter(torch.zeros(1, 1, 1, config.decoder_embed_dim))
+        self.channel_mask_token = nn.Parameter(torch.zeros(1, 1, 1, config.decoder_embed_dim))
+        
+        if config.drop_path_uniform is True:
+            dpr = [config.drop_path_rate] * config.decoder_depth
+        else:
+            dpr = [x.item() for x in torch.linspace(0, config.drop_path_rate, config.decoder_depth)]
+        
+        self.decoder_blocks = nn.ModuleList([
+            LowRankBlock(
+                dim=config.decoder_embed_dim,
+                num_heads=config.decoder_num_heads,
+                channel_dim=config.decoder_channel_dim,
+                spatial_dim=config.decoder_spatial_dim,
+                mlp_ratio=config.mlp_ratio,
+                qkv_bias=config.qkv_bias,
+                qk_norm=config.qk_norm,
+                proj_drop=config.proj_drop,
+                attn_drop=config.attn_drop,
+                drop_path=dpr[i],
+                init_values=config.init_values,
+                norm_layer=norm_layer,
+                skip_pool=False,
+                rank=config.rank,
+                attn_type=config.attn_type,
+                fusion_init_scale=config.fusion_init_scale,
+            )
+            for i in range(config.decoder_depth)
+        ])
+
+        self.decoder_norm = norm_layer(config.decoder_embed_dim)
+        self.decoder_optical_pred = nn.Linear(config.decoder_embed_dim, config.patch_size**2, bias=True)
+        self.decoder_radar_pred = nn.Linear(config.decoder_embed_dim, config.patch_size**2, bias=True)
+
+        if not config.use_rope_embed:
+            self.pos_chan_embed = PositionalChannelEmbedding(config.decoder_embed_dim)
+        else:
+            decoder_rope_spatial_dim = config.decoder_rope_spatial_split_dim if config.attn_type == "full" else config.decoder_spatial_dim
+            decoder_rope_channel_dim = config.decoder_rope_channel_split_dim if config.attn_type == "full" else config.decoder_channel_dim
+            self.pos_chan_embed = RopePositionChannelEmbedding(
+                # embed_dim=config.decoder_embed_dim,
+                spatial_dim=decoder_rope_spatial_dim,
+                channel_dim=decoder_rope_channel_dim,
+                num_heads=config.decoder_num_heads,
+                base=config.rope_embed_base if hasattr(config, 'rope_embed_base') else 100.0
+            )
+        
+        self.initialize_weights()
+
+    def initialize_weights(self):
+        torch.nn.init.normal_(self.mask_token, std=.02)
+        torch.nn.init.normal_(self.spatial_mask_token, std=.02)
+        torch.nn.init.normal_(self.channel_mask_token, std=.02)
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            torch.nn.init.xavier_uniform_(m.weight)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def forward(self, x, pos_ids_restore, channel_ids_restore, optical_channel_wv, radar_channel_wv, spatial_resolution, restore_input_dim=False):
+        hidden_states = []
+
+        # embed tokens
+        x = self.decoder_embed(x)  # B N+1 L+1 D
+
+        if radar_channel_wv is not None:
+            channel_wv = torch.cat((optical_channel_wv, radar_channel_wv), dim=1)  # 1 C = Co + Cr or B C
+            n_radar_channels = radar_channel_wv.shape[1]
+        else:
+            channel_wv = optical_channel_wv
+            n_radar_channels = None
+        n_optical_channels = optical_channel_wv.shape[1]
+        
+        # remove cls token
+        B, N, L, D = x[:, 1:, 1:, :].shape
+        C = channel_ids_restore.shape[1]
+        HW = pos_ids_restore.shape[1]
+
+        # append mask tokens to sequence
+        channel_mask_token = self.channel_mask_token.expand(B, 1, HW - L, -1)
+        mask_tokens = self.mask_token.expand(B, N, HW - L, -1)
+        mask_tokens = torch.cat([channel_mask_token, mask_tokens], dim=1)
+        x_ = torch.cat([x[:, :, 1:], mask_tokens], dim=2)  # B N+1 HW D, remove cls token
+        x_ = torch.gather(x_, dim=2, index=pos_ids_restore.unsqueeze(1).unsqueeze(-1).repeat(1, N + 1, 1, D))  # unshuffle, B N+1 HW D
+        x = torch.cat([x[:, :, :1], x_], dim=2)  # B N+1 HW+1 D, add cls token
+        
+        # append channel mask tokens to sequence
+        spatial_mask_token = self.spatial_mask_token.expand(B, C - N, 1, -1)
+        mask_tokens = self.mask_token.expand(B, C - N, HW, -1)
+        mask_tokens = torch.cat([spatial_mask_token, mask_tokens], dim=2)
+        x_ = torch.cat([x[:, 1:], mask_tokens], dim=1)  # B, C, HW+1, D, remove cls token
+        x_ = torch.gather(x_, dim=1, index=channel_ids_restore.unsqueeze(-1).unsqueeze(-1).repeat(1, 1, HW + 1, D))  # unshuffle, B C HW+1 D
+        x = torch.cat([x[:, :1], x_], dim=1)  # B C+1 HW+1 D, add cls token
+        
+        if n_radar_channels is not None and C == optical_channel_wv.shape[1]: 
+            # missing radar channels from encoder
+            n_radar_channels = radar_channel_wv.shape[1]
+            # extend x to C+n_radar_channels to the end
+            spatial_mask_token = self.spatial_mask_token.expand(B, n_radar_channels, 1, -1)
+            mask_tokens = self.mask_token.expand(B, n_radar_channels, HW, -1)  # B C HW D
+            mask_tokens = torch.cat([spatial_mask_token, mask_tokens], dim=2)
+            x = torch.cat([x, mask_tokens], dim=1)
+        elif n_radar_channels is not None and C == radar_channel_wv.shape[1]:
+            # missing optical channels from encoder
+            n_optical_channels = optical_channel_wv.shape[1]
+            # extend x to C+n_optical_channels to the front behind the cls token
+            spatial_mask_token = self.spatial_mask_token.expand(B, n_optical_channels, 1, -1)
+            mask_tokens = self.mask_token.expand(B, n_optical_channels, HW, -1)
+            mask_tokens = torch.cat([spatial_mask_token, mask_tokens], dim=2)
+            x = torch.cat([x[:, :1, :, :], mask_tokens, x[:, 1:, :, :]], dim=1)
+
+        if not self.config.use_rope_embed:
+            pos_chan_embed = self.pos_chan_embed(x[:, 1:, 1:, :], channel_ids=channel_wv, spatial_resolution=spatial_resolution, cls_token=True).to(x.device, dtype=x.dtype)
+            x = x + pos_chan_embed
+        else:
+            assert math.sqrt(HW).is_integer(), "For RoPE, the input num of tokens must be a perfect square."
+            H = W = int(math.sqrt(HW))
+            pos_chan_embed = self.pos_chan_embed(H=H, W=W, C=C, optical_channel_wv=channel_wv)
+            pos_chan_embed = [embed.repeat(B, 1, 1) for embed in pos_chan_embed] # walk around for RoPE
+            pos_chan_embed = tuple(pos_chan_embed)
+
+        num_patches = x.shape[2] - 1
+        if self.config.use_perception_field_mask:
+            perception_field_mask = get_perception_field_mask(num_patches, self.config.patch_size, spatial_resolution, attention_radius=self.config.attention_radius, cls_token=True).to(x.device)
+        else:
+            perception_field_mask = None
+
+        for i, blk in enumerate(self.decoder_blocks):
+            if not self.config.use_rope_embed:
+                pos_chan_embedding = pos_chan_embed*2**(-i) if self.config.pos_chan_embed_residual else None
+            else:
+                pos_chan_embedding = pos_chan_embed
+            x = blk(x, spatial_mask=perception_field_mask, pos_chan_embedding=pos_chan_embedding)
+
+        x = self.decoder_norm(x)
+        x = x[:, 1:, 1:, :]
+
+        x_optical = x[:, :n_optical_channels, :, :]
+        x_radar = x[:, n_optical_channels:, :, :]
+
+        x_optical = self.decoder_optical_pred(x_optical)
+        x_radar = self.decoder_radar_pred(x_radar)
+
+        x = torch.cat([x_optical, x_radar], dim=1) # B C HW patch_size**2
+
+        if restore_input_dim:
+            x = self.unpatchify(x)
+        
+        if self.config.return_dict:
+            # Return as BaseModelOutput for Hugging Face compatibility
+            return BaseModelOutput(
+                last_hidden_state=x,
+                hidden_states=hidden_states,
+                attentions=None,  # You can modify this if you have attention weights to return
+            )
+        else:
+            return x, hidden_states
+
+    def unpatchify(self, x):
+        """
+        x: (B, C, HW, patch_size**2)
+        imgs: (B, C, H, W)
+        """
+        p = self.config.patch_size
+        h = w = int(x.shape[2]**.5)
+        assert h * w == x.shape[2]
+        
+        x = x.reshape(shape=(x.shape[0], x.shape[1], h, w, p, p))
+        x = torch.einsum('bchwpq->bchpwq', x)
+        imgs = x.reshape(shape=(x.shape[0], x.shape[1], h * p, h * p))
+        return imgs
+
+    def forward_target(self, imgs):
+        """
+        imgs: [B, C, H, W]
+        """
+        target = self.patchify(imgs)
+        if self.config.norm_pix_loss:
+            mean = target.mean(dim=-1, keepdim=True)
+            var = target.var(dim=-1, keepdim=True)
+            target = (target - mean) / (var + 1.e-6)**.5
+        return target # B C HW patch_size**2
+
+    def patchify(self, imgs):
+        """
+        imgs: (B, C, H, W)
+        x: (B, C, HW, patch_size**2)
+        """
+        p = self.config.patch_size
+        assert imgs.shape[2] == imgs.shape[3] and imgs.shape[2] % p == 0
+
+        h = w = imgs.shape[2] // p
+        x = imgs.reshape(shape=(imgs.shape[0], imgs.shape[1], h, p, w, p))
+        x = torch.einsum('bchpwq->bchwpq', x)
+        x = x.reshape(shape=(imgs.shape[0], imgs.shape[1], h * w, p**2))
+        return x
+        

validation.json

@@ -0,0 +1,20 @@
+{
+  "encoder_parameters": 83219328,
+  "verified_tensors": 543,
+  "remapped_tensors": 408,
+  "output_shape": [
+    1,
+    11,
+    5,
+    768
+  ],
+  "auto_model_loading_info": {
+    "missing_keys": [],
+    "unexpected_keys": [],
+    "mismatched_keys": [],
+    "error_msgs": []
+  },
+  "forward_exact_match": true,
+  "encoder_backward": "passed",
+  "original_checkpoint_preserved": true
+}