"""Architecture ablations for the fixed-32 PASTIS fine-tuning protocol. All existing encoder parameter names are retained for checkpoint loading. The validity experiment changes only handling of collate/internal padding; the seasonal experiment adds a residual temporal-pooling branch. """ import torch from torch import nn import torch.nn.functional as F from torch.utils.checkpoint import checkpoint import config as C from model import ( EVAL_TEMPORAL_CHUNK, SPATIAL_CHUNK, TEMPORAL_CHUNK, STA, UTAE_ASPP, UTAEClassificationDual, ) from temporal_srda import TemporalSRDABlock, TemporalSRDAEncoder def valid_from_index_positions(pos): """The unmodified dataset uses 0,1,...,T-1 and pads positions with zero. Temporal subsampling may drop acquisition zero; its first retained position is then positive. Only an initial zero is a valid acquisition-zero token. This cannot be reused for arbitrary calendar positions without a length mask. """ first_is_zero = pos[:, :1].eq(0) first = torch.zeros_like(pos, dtype=torch.bool) first[:, :1] = first_is_zero valid = pos.gt(0) | first if not valid.any(dim=1).all(): raise ValueError('Each time series must contain an observation') return valid class MaskedTemporalSRDABlock(TemporalSRDABlock): def forward(self, x, valid): n, t, d = x.shape r = self.R x = x * valid.unsqueeze(-1) # Remove partial-group padding *after* LayerNorm. Learned LN biases # would otherwise turn padded zero vectors into nonzero KV features. z = self.pw(self.norm_kv(x)) * valid.unsqueeze(-1) z = z.reshape(n, t // r, d) group_valid = valid.reshape(n, t // r, r).any(dim=-1) zn = self.norm_first(z) z_sa, _ = self.first_sa( zn, zn, zn, key_padding_mask=~group_valid, need_weights=False ) kv = (z + z_sa) * group_valid.unsqueeze(-1) q = self.q_proj(self.norm_q(x)) k, v = self.kv_proj(kv).chunk(2, dim=-1) q = q.reshape(n, t, self.h, self.dh).transpose(1, 2) k = k.reshape(n, -1, self.h, self.dh).transpose(1, 2) v = v.reshape(n, -1, self.h, self.dh).transpose(1, 2) scores = (q @ k.transpose(-2, -1)) * self.scale scores = scores.masked_fill(~group_valid[:, None, None, :], torch.finfo(scores.dtype).min) weights = self.attn_drop(scores.softmax(dim=-1)) attended = (weights @ v).transpose(1, 2).reshape(n, t, d) x = (x + self.proj_drop(self.out_proj(attended))) * valid.unsqueeze(-1) y = F.gelu(self.fc1(self.norm_mlp(x))) * valid.unsqueeze(-1) y = y + self.dw(y.transpose(1, 2)).transpose(1, 2) x = (x + self.fc2(y)) * valid.unsqueeze(-1) return x class MaskedTemporalSRDAEncoder(TemporalSRDAEncoder): def __init__(self, d_model=256, n_heads=8, R_schedule=(4, 4, 4), **kwargs): super().__init__(d_model=d_model, n_heads=n_heads, R_schedule=R_schedule, **kwargs) self.blocks = nn.ModuleList( MaskedTemporalSRDABlock(d_model, n_heads, R=r) for r in R_schedule ) def forward(self, x, pos=None, valid=None): if valid is None: raise ValueError('Validity mask is required for this ablation') if pos is not None: x = x + pos original_t = x.shape[1] pad = (-original_t) % self.R_max if pad: x = F.pad(x, (0, 0, 0, pad)) valid = F.pad(valid, (0, pad), value=False) x = x * valid.unsqueeze(-1) for block in self.blocks: x = block(x, valid) return x[:, :original_t] if pad else x class MaskedUTAE_ASPP(UTAE_ASPP): def __init__(self, n_channels=10, d_model=256, n_heads=8, **kwargs): super().__init__(n_channels=n_channels, d_model=d_model, n_heads=n_heads, **kwargs) self.temporal_encoder = MaskedTemporalSRDAEncoder( d_model=d_model, n_heads=n_heads, R_schedule=C.R_SCHEDULE ) def forward(self, x, pos): b, t, channels, h, w = x.shape valid = valid_from_index_positions(pos) frames = x.reshape(b * t, channels, h, w) flat_valid = valid.flatten() # Padding is absent from the spatial BN batch as well as temporal # attention. This matters during training with batch size two. active = frames[flat_valid] encoded = self._encode_spatial(active) d = encoded.shape[1] x2 = encoded.new_zeros((b * t, d, h, w)) x2 = x2.index_copy(0, flat_valid.nonzero().flatten(), encoded) pixels = x2.reshape(b, t, d, h, w).permute(0, 3, 4, 1, 2).reshape(b * h * w, t, d) pe = self.pos_encoder(pos) expanded_pos = pe[:, None, None].expand(b, h, w, t, d).reshape(b * h * w, t, d) expanded_valid = valid[:, None, None].expand(b, h, w, t).reshape(b * h * w, t) n = pixels.shape[0] if self.training and n > TEMPORAL_CHUNK: output = torch.cat([ checkpoint( self.temporal_encoder, pixels[s:s + TEMPORAL_CHUNK], expanded_pos[s:s + TEMPORAL_CHUNK], expanded_valid[s:s + TEMPORAL_CHUNK], use_reentrant=False ) for s in range(0, n, TEMPORAL_CHUNK) ], dim=0) elif not self.training and n > EVAL_TEMPORAL_CHUNK: output = torch.cat([ self.temporal_encoder( pixels[s:s + EVAL_TEMPORAL_CHUNK], expanded_pos[s:s + EVAL_TEMPORAL_CHUNK], expanded_valid[s:s + EVAL_TEMPORAL_CHUNK], ) for s in range(0, n, EVAL_TEMPORAL_CHUNK) ], dim=0) else: output = self.temporal_encoder(pixels, expanded_pos, expanded_valid) return output.reshape(b, h, w, t, d), x2 class MaskedSTA(STA): def forward(self, feats, valid): scores = self.l(feats) scores = scores.masked_fill(~valid[:, None, None, :, None], torch.finfo(scores.dtype).min) weights = scores.softmax(dim=3) pooled = (feats * weights).sum(dim=3).permute(0, 3, 1, 2).contiguous() pooled = pooled + self.conv2(self.conv1(pooled)) return self.conv3(pooled) class MaskedUTAEClassificationDual(UTAEClassificationDual): def __init__(self, encoder, d_model=256, num_classes=20): super().__init__(encoder, d_model=d_model, num_classes=num_classes) self.sta = MaskedSTA(d_model) def forward(self, x, pos): features, _ = self.utae(x, pos) dense = self.sta(features, valid_from_index_positions(pos)) semantic = self.sem_head(dense) boundary = self.bnd_head(dense) return { 'sem_logits': semantic, 'bnd_logits': boundary, 'refined_logits': self.refine(dense, semantic, boundary), } class SeasonalSTA(STA): def __init__(self, d_model=256, n_queries=4): super().__init__(d_model) self.seasonal_scores = nn.Linear(d_model, n_queries) self.seasonal_projection = nn.Conv2d(n_queries * d_model, d_model, 1) nn.init.zeros_(self.seasonal_projection.weight) nn.init.zeros_(self.seasonal_projection.bias) def forward(self, feats): base = super().forward(feats) # Each query pools a different temporal distribution at every pixel. weights = self.seasonal_scores(feats).softmax(dim=3) pooled = torch.einsum('bhwtd,bhwtq->bhwqd', feats, weights) b, h, w, q, d = pooled.shape pooled = pooled.reshape(b, h, w, q * d).permute(0, 3, 1, 2).contiguous() return base + self.seasonal_projection(pooled) class SeasonalUTAEClassificationDual(UTAEClassificationDual): def __init__(self, encoder, d_model=256, num_classes=20): super().__init__(encoder, d_model=d_model, num_classes=num_classes) self.sta = SeasonalSTA(d_model)