from typing import Tuple import torch import torch.nn as nn from einops import rearrange from timee.model.layers import ( ClassEmbedding, CLSTokens, DecoderMLP, InstanceNorm, Patch, PatchEncoder, TransformerEncoderLayer, ) class TimeeModel(nn.Module): """TIMEE: in-context learning (ICL) model for time series classification. Architecture: 1. Patch the input time series (non-overlapping patches). 2. Normalize each series (arcsinh instance norm) and across series. 3. Embed patches linearly; add patch-level and series-level stats embeddings. 4. Inject class embeddings into training-set patches. 5. Run n_horizontal_layers of temporal self-attention (RoPE), followed by n_vertical_layers of cross-series attention. 6. Compress each series to a fixed-size representation (num_cls_tokens CLS tokens) via a final attention layer. 7. Run n_icl_layers of ICL attention over the compressed representations. 8. Decode test-set representations to class logits via an MLP. Default arguments match the released checkpoint. """ def __init__( self, n_max_classes: int = 10, patch_size: int = 16, patch_stride: int = 16, d_model: int = 128, n_horizontal_layers: int = 5, n_vertical_layers: int = 5, n_icl_layers: int = 2, n_heads: int = 4, d_kv: int = 32, d_kv_icl: int = 64, mlp_hidden_dim: int = 512, decoder_mlp_hidden_dim: int = 512, dropout_rate: float = 0.1, num_cls_tokens: int = 4, ) -> None: super().__init__() self.n_max_classes = n_max_classes self.patch_size = patch_size self.patch_stride = patch_stride self.d_model = d_model self.n_horizontal_layers = n_horizontal_layers self.n_vertical_layers = n_vertical_layers self.n_icl_layers = n_icl_layers self.n_heads = n_heads self.d_kv = d_kv self.d_kv_icl = d_kv_icl self.mlp_hidden_dim = mlp_hidden_dim self.decoder_mlp_hidden_dim = decoder_mlp_hidden_dim self.dropout_rate = dropout_rate self.num_cls_tokens = num_cls_tokens self.icl_dim = d_model * num_cls_tokens self.instance_norm_sequence = InstanceNorm(use_arcsinh=True) self.instance_norm_series = InstanceNorm() self.patch = Patch(patch_size=patch_size, patch_stride=patch_stride) self.patch_encoder = PatchEncoder(patch_size=patch_size, d_model=d_model) self.patch_stats_projection = nn.Linear(2, d_model, bias=False) self.series_stats_projection = nn.Linear(2, d_model, bias=False) self.cls_tokens = CLSTokens(num_cls_tokens, d_model) layer_kwargs = dict( d_model=d_model, n_heads=n_heads, d_kv=d_kv, mlp_hidden_dim=mlp_hidden_dim, dropout=dropout_rate, ) self.encoder_horizontal = nn.ModuleList( [ TransformerEncoderLayer(**layer_kwargs, use_rope=True) for _ in range(n_horizontal_layers) ] ) self.encoder_vertical = nn.ModuleList( [ TransformerEncoderLayer(**layer_kwargs, use_rope=False) for _ in range(n_vertical_layers) ] ) self.last_horizontal_encoder = TransformerEncoderLayer(**layer_kwargs, use_rope=True) self.cls_tokens_ln = nn.LayerNorm(d_model) self.class_embedding_for_patches = ClassEmbedding( num_embeddings=n_max_classes, embedding_dim=d_model, ) self.class_embedding_icl = ClassEmbedding( num_embeddings=n_max_classes, embedding_dim=self.icl_dim, ) self.icl_block = nn.ModuleList( [ TransformerEncoderLayer( d_model=self.icl_dim, n_heads=n_heads, d_kv=d_kv_icl, mlp_hidden_dim=mlp_hidden_dim, dropout=dropout_rate, use_rope=False, ) for _ in range(n_icl_layers) ] ) self.decoder_mlp_ln = nn.LayerNorm(self.icl_dim) self.decoder_mlp = DecoderMLP( d_model=self.icl_dim, hidden_dim=decoder_mlp_hidden_dim, output_dim=n_max_classes, dropout=dropout_rate, ) def _compute_stats_embeddings(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: """Compute scale-aware stat embeddings from raw (un-normalized) x (B, N, S). Returns (patch_stats, series_stats): patch_stats (B, N, P, 2): arcsinh(mean, std) per patch, z-scored within each series. series_stats (B, N, 1, 2): arcsinh(mean, std) per series, z-scored within the episode. """ x_BNPQ = self.patch(x) raw_mean = torch.nan_to_num(torch.nanmean(x_BNPQ, dim=-1, keepdim=True), nan=0.0) raw_std = ((x_BNPQ - raw_mean).nan_to_num(0.0) ** 2).mean(dim=-1, keepdim=True).sqrt() patch_stats = torch.arcsinh(torch.cat([raw_mean, raw_std], dim=-1)) p_mean = patch_stats.mean(dim=2, keepdim=True) p_std = patch_stats.std(dim=2, correction=0, keepdim=True).clamp(min=1e-5) patch_stats = (patch_stats - p_mean) / p_std series_mean = torch.nan_to_num(torch.nanmean(x, dim=-1, keepdim=True), nan=0.0) series_std = ((x - series_mean).nan_to_num(0.0) ** 2).mean(dim=-1, keepdim=True).sqrt() series_stats = torch.arcsinh(torch.cat([series_mean, series_std], dim=-1)) ep_mean = series_stats.mean(dim=1, keepdim=True) ep_std = series_stats.std(dim=1, keepdim=True).clamp(min=1e-5) series_stats = ((series_stats - ep_mean) / ep_std).unsqueeze(2) return patch_stats, series_stats def forward( self, x: torch.Tensor, y: torch.Tensor, eval_pos: int, ) -> torch.Tensor: """ Args: x: (B, N, S) float — N series per episode (support first, then queries). y: (B, N) int — class labels; query positions (>= eval_pos) are ignored. eval_pos: number of support series; x[:, :eval_pos] are labeled, x[:, eval_pos:] are the queries to classify. Returns: logits: (B, N − eval_pos, n_max_classes). """ x_BNS = x y_BN = y n_train = eval_pos batch_size, num_series, _ = x_BNS.shape patch_stats_BNPC, series_stats_BN12 = self._compute_stats_embeddings(x_BNS) # Normalize: instance norm per series, then cross-series norm. x_BNS, _ = self.instance_norm_sequence(x_BNS) x_BSN = rearrange(x_BNS, "b n s -> b s n") x_BSN, _ = self.instance_norm_series(x_BSN, eval_pos=n_train) x_BNS = rearrange(x_BSN, "b s n -> b n s") x_BNPQ = self.patch(x_BNS) x_BNPQ = torch.nan_to_num(x_BNPQ, nan=0.0) x_BNPD = self.patch_encoder(x_BNPQ) x_BNPD = x_BNPD + self.patch_stats_projection(patch_stats_BNPC) x_BNPD = x_BNPD + self.series_stats_projection( series_stats_BN12.expand(-1, -1, x_BNPD.shape[2], -1) ) # Inject class embeddings into train patches. class_embed_BN1D = self.class_embedding_for_patches(y_BN[:, :n_train], n_train).unsqueeze(2) x_BNPD = torch.cat([x_BNPD[:, :n_train] + class_embed_BN1D, x_BNPD[:, n_train:]], dim=1) x_BNPD = torch.cat([self.cls_tokens(batch_size, num_series), x_BNPD], dim=2) num_patches_plus_cls = x_BNPD.shape[2] num_patches = num_patches_plus_cls - self.num_cls_tokens # Mask convention: True = attend, False = blocked. # Horizontal mask is fully open — each series attends to all its own patches. horizontal_mask = torch.ones( (batch_size * num_series, num_patches_plus_cls, num_patches_plus_cls), device=x_BNPD.device, dtype=torch.bool, ) # Queries cannot attend to other queries; they can only see support series (and themselves). vertical_mask = torch.ones( (batch_size * num_patches_plus_cls, num_series, num_series), device=x_BNPD.device, dtype=torch.bool, ) vertical_mask[:, :, n_train:] = False vertical_mask |= torch.eye(num_series, device=x_BNPD.device, dtype=torch.bool) position_ids = ( torch.arange(num_patches, device=x_BNPD.device) .unsqueeze(0) .repeat(batch_size * num_series, 1) ) for layer in self.encoder_horizontal: x_BNPD = rearrange(x_BNPD, "b n p d -> (b n) p d") x_BNPD = layer(x=x_BNPD, attention_mask=horizontal_mask, position_ids=position_ids) x_BNPD = rearrange(x_BNPD, "(b n) p d -> b n p d", b=batch_size, n=num_series) for layer in self.encoder_vertical: x_BNPD = rearrange(x_BNPD, "b n p d -> (b p) n d") x_BNPD = layer(x=x_BNPD, attention_mask=vertical_mask, position_ids=None) x_BNPD = rearrange(x_BNPD, "(b p) n d -> b n p d", b=batch_size, p=num_patches_plus_cls) # No position attends to CLS tokens as keys, so CLS tokens aggregate only from patches. compression_mask = torch.ones( (batch_size * num_series, num_patches_plus_cls, num_patches_plus_cls), device=x_BNPD.device, dtype=torch.bool, ) compression_mask[:, :, : self.num_cls_tokens] = False x_BNPD = rearrange(x_BNPD, "b n p d -> (b n) p d") x_BNPD = self.last_horizontal_encoder( x=x_BNPD, attention_mask=compression_mask, position_ids=position_ids, ) x_BNPD = rearrange(x_BNPD, "(b n) p d -> b n p d", b=batch_size, n=num_series) row_repr_BNDc = self.cls_tokens_ln(x_BNPD[:, :, : self.num_cls_tokens]).flatten(-2, -1) # ICL: cross-series attention over compressed representations. # Same masking policy as vertical layers — queries attend to support + self only. icl_mask = torch.ones( (batch_size, num_series, num_series), device=x_BNPD.device, dtype=torch.bool, ) icl_mask[:, :, n_train:] = False icl_mask |= torch.eye(num_series, device=x_BNPD.device, dtype=torch.bool) class_embed_icl = self.class_embedding_icl(y_BN[:, :n_train], n_train) row_repr_BNDc = torch.cat( [row_repr_BNDc[:, :n_train] + class_embed_icl, row_repr_BNDc[:, n_train:]], dim=1 ) for layer in self.icl_block: row_repr_BNDc = layer(x=row_repr_BNDc, attention_mask=icl_mask, position_ids=None) return self.decoder_mlp(self.decoder_mlp_ln(row_repr_BNDc[:, n_train:]))