Update encoders.py

encoders.py

@@ -1,559 +1,516 @@
-"""
-改进的多模态编码器 - SOTA级别（修复版）
-集成最新的视觉、音频、视频编码技术
-"""
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from typing import Tuple, Optional
-from components import RMSNorm, SwiGLU
-from transformer import OptimizedTransformerBlock
-import math
-
-class LayerScale(nn.Module):
-    """LayerScale - 改进训练稳定性"""
-    def __init__(self, dim: int, init_values: float = 1e-5):
-        super().__init__()
-        self.gamma = nn.Parameter(init_values * torch.ones(dim))
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return x * self.gamma
-
-class StochasticDepth(nn.Module):
-    """随机深度 - Drop Path"""
-    def __init__(self, drop_prob: float = 0.0):
-        super().__init__()
-        self.drop_prob = drop_prob
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        if not self.training or self.drop_prob == 0.0:
-            return x
-        
-        keep_prob = 1 - self.drop_prob
-        shape = (x.shape[0],) + (1,) * (x.ndim - 1)
-        random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
-        random_tensor.floor_()
-        return x.div(keep_prob) * random_tensor
-
-class ImprovedPatchEmbedding(nn.Module):
-    """改进的图像分块嵌入 - 支持重叠patch和多尺度"""
-    def __init__(
-        self,
-        patch_size: int = 14,
-        in_channels: int = 3,
-        embed_dim: int = 2048,
-        overlap: int = 0
-    ):
-        super().__init__()
-        self.patch_size = patch_size
-        stride = patch_size - overlap
-        self.proj = nn.Conv2d(
-            in_channels,
-            embed_dim,
-            kernel_size=patch_size,
-            stride=stride,
-            padding=overlap // 2
-        )
-        
-        self.norm = RMSNorm(embed_dim)
-
-    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, Tuple[int, int]]:
-        B, C, H, W = x.shape
-        x = self.proj(x)
-        grid_size = (x.shape[2], x.shape[3])
-        x = x.flatten(2).transpose(1, 2)
-        x = self.norm(x)
-        return x, grid_size
-
-class ImprovedVisionBlock(nn.Module):
-    """改进的Vision Transformer Block"""
-    def __init__(
-        self,
-        dim: int,
-        n_heads: int,
-        dropout: float = 0.0,
-        drop_path: float = 0.0,
-        use_adapter: bool = False,
-        adapter_dim: int = 64,
-        use_layer_scale: bool = True,
-        layer_scale_init: float = 1e-5
-    ):
-        super().__init__()
-        self.norm1 = RMSNorm(dim)
-        self.attn = nn.MultiheadAttention(
-            dim, n_heads, dropout=dropout, batch_first=True
-        )
-        
-        self.norm2 = RMSNorm(dim)
-        self.mlp = nn.Sequential(
-            nn.Linear(dim, dim * 4),
-            nn.GELU(),
-            nn.Dropout(dropout),
-            nn.Linear(dim * 4, dim),
-            nn.Dropout(dropout)
-        )
-        
-        self.drop_path = StochasticDepth(drop_path) if drop_path > 0 else nn.Identity()
-        
-        if use_layer_scale:
-            self.ls1 = LayerScale(dim, layer_scale_init)
-            self.ls2 = LayerScale(dim, layer_scale_init)
-        else:
-            self.ls1 = nn.Identity()
-            self.ls2 = nn.Identity()
-        
-        # 修复：使用简单的adapter实现，避免外部依赖
-        if use_adapter:
-            self.adapter = nn.Sequential(
-                nn.Linear(dim, adapter_dim),
-                nn.GELU(),
-                nn.Linear(adapter_dim, dim)
-            )
-        else:
-            self.adapter = None
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        # 注意力
-        normx = self.norm1(x)
-        attn_out, _ = self.attn(normx, normx, normx)
-        x = x + self.drop_path(self.ls1(attn_out))
-        
-        # MLP
-        x = x + self.drop_path(self.ls2(self.mlp(self.norm2(x))))
-        
-        # Adapter
-        if self.adapter is not None:
-            x = x + self.adapter(x)
-        
-        return x
-
-class ImprovedVisionTransformer(nn.Module):
-    """
-    改进的视觉Transformer
-    - LayerScale
-    - Stochastic Depth
-    - 改进的位置编码
-    - 可选的Register tokens
-    """
-    def __init__(
-        self,
-        img_size: int = 224,
-        patch_size: int = 14,
-        in_channels: int = 3,
-        embed_dim: int = 2048,
-        depth: int = 24,
-        n_heads: int = 16,
-        dropout: float = 0.0,
-        drop_path_rate: float = 0.1,
-        use_register_tokens: bool = True,
-        num_register_tokens: int = 4,
-        use_adapter: bool = False,
-        adapter_dim: int = 64,
-        use_layer_scale: bool = True,
-        layer_scale_init: float = 1e-5
-    ):
-        super().__init__()
-        self.patch_size = patch_size
-        self.embed_dim = embed_dim
-        self.use_register_tokens = use_register_tokens
-        self.num_register_tokens = num_register_tokens if use_register_tokens else 0
-        
-        # Patch embedding
-        self.patch_embed = ImprovedPatchEmbedding(
-            patch_size, in_channels, embed_dim, overlap=0
-        )
-        
-        self.pretrain_img_size = img_size
-        n_patches_pretrain = (img_size // patch_size) ** 2
-        
-        # CLS token
-        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
-        
-        # Register tokens (DINOv2启发)
-        if use_register_tokens:
-            self.register_tokens = nn.Parameter(
-                torch.zeros(1, num_register_tokens, embed_dim)
-            )
-        
-        # 修复：位置编码总数 = 1(CLS) + n_patches + register_tokens
-        total_tokens = 1 + n_patches_pretrain + self.num_register_tokens
-        self.pos_embed = nn.Parameter(
-            torch.zeros(1, total_tokens, embed_dim)
-        )
-        self.pos_drop = nn.Dropout(dropout)
-        
-        # Stochastic depth
-        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]
-        
-        # Transformer blocks
-        self.blocks = nn.ModuleList([
-            ImprovedVisionBlock(
-                embed_dim,
-                n_heads,
-                dropout,
-                drop_path=dpr[i],
-                use_adapter=use_adapter,
-                adapter_dim=adapter_dim,
-                use_layer_scale=use_layer_scale,
-                layer_scale_init=layer_scale_init
-            )
-            for i in range(depth)
-        ])
-        
-        self.norm = RMSNorm(embed_dim)
-        self._init_weights()
-
-    def _init_weights(self):
-        nn.init.trunc_normal_(self.cls_token, std=0.02)
-        nn.init.trunc_normal_(self.pos_embed, std=0.02)
-        if self.use_register_tokens:
-            nn.init.trunc_normal_(self.register_tokens, std=0.02)
-        
-        self.apply(self._init_module_weights)
-
-    def _init_module_weights(self, m):
-        if isinstance(m, nn.Linear):
-            nn.init.trunc_normal_(m.weight, std=0.02)
-            if m.bias is not None:
-                nn.init.zeros_(m.bias)
-        elif isinstance(m, nn.Conv2d):
-            nn.init.trunc_normal_(m.weight, std=0.02)
-            if m.bias is not None:
-                nn.init.zeros_(m.bias)
-        elif isinstance(m, RMSNorm):
-            if hasattr(m, 'weight') and m.weight is not None:
-                nn.init.ones_(m.weight)
-
-    def _interpolate_pos_encoding(
-        self,
-        patch_tokens: torch.Tensor,
-        grid_size: Tuple[int, int]
-    ) -> torch.Tensor:
-        """
-        修复：改进的位置编码插值
-        只对patch位置编码进行插值，CLS和register token位置编码保持不变
-        """
-        pretrain_grid_h = self.pretrain_img_size // self.patch_size
-        pretrain_grid_w = pretrain_grid_h
-        
-        # 如果尺寸匹配，直接返回原始位置编码
-        if grid_size[0] == pretrain_grid_h and grid_size[1] == pretrain_grid_w:
-            return self.pos_embed
-        
-        # 分离不同部分的位置编码
-        # pos_embed结构: [CLS(1), register_tokens(n), patches(H*W)]
-        num_extra_tokens = 1 + self.num_register_tokens
-        cls_register_pos = self.pos_embed[:, :num_extra_tokens, :]  # [1, 1+n, dim]
-        patch_pos_embed = self.pos_embed[:, num_extra_tokens:, :]  # [1, H*W, dim]
-        
-        # 2D插值patch位置编码
-        patch_pos_embed = patch_pos_embed.reshape(
-            1, pretrain_grid_h, pretrain_grid_w, -1
-        ).permute(0, 3, 1, 2)
-        
-        patch_pos_embed = F.interpolate(
-            patch_pos_embed,
-            size=grid_size,
-            mode='bicubic',
-            align_corners=False
-        )
-        
-        patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).flatten(1, 2)
-        
-        # 拼接回去
-        return torch.cat([cls_register_pos, patch_pos_embed], dim=1)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        B = x.shape[0]
-        
-        # Patch embedding
-        x, grid_size = self.patch_embed(x)
-        
-        # 添加CLS token
-        cls_tokens = self.cls_token.expand(B, -1, -1)
-        
-        # 修复：正确组装tokens序列
-        if self.use_register_tokens:
-            register_tokens = self.register_tokens.expand(B, -1, -1)
-            # 顺序: [CLS, register_tokens, patches]
-            x = torch.cat([cls_tokens, register_tokens, x], dim=1)
-        else:
-            x = torch.cat([cls_tokens, x], dim=1)
-        
-        # 位置编码（插值以适应不同尺寸）
-        pos_embed = self._interpolate_pos_encoding(x, grid_size)
-        x = self.pos_drop(x + pos_embed)
-        
-        # Transformer blocks
-        for block in self.blocks:
-            x = block(x)
-        
-        x = self.norm(x)
-        
-        # 返回所有tokens（调用者可以选择使用CLS token或全局池化）
-        return x
-
-class ImprovedAudioEncoder(nn.Module):
-    """
-    改进的音频编码器
-    - 时序建模
-    - 频率建模
-    - 双流架构
-    """
-    def __init__(
-        self,
-        n_mels: int = 128,
-        target_length: int = 1024,
-        embed_dim: int = 2048,
-        depth: int = 12,
-        n_heads: int = 16,
-        patch_size: int = 16,
-        dropout: float = 0.1,
-        use_adapter: bool = False,
-        adapter_dim: int = 64,
-        use_dual_stream: bool = True
-    ):
-        super().__init__()
-        self.use_dual_stream = use_dual_stream
-        self.patch_size = patch_size
-        
-        # 主编码器
-        self.patch_embed = nn.Conv2d(
-            1, embed_dim, kernel_size=patch_size, stride=patch_size
-        )
-        
-        # 修复：计算实际的patch数量
-        self.n_patches_h = n_mels // patch_size
-        self.n_patches_w = target_length // patch_size
-        n_patches = self.n_patches_h * self.n_patches_w
-        
-        self.pos_embed = nn.Parameter(torch.zeros(1, n_patches, embed_dim))
-        self.pos_drop = nn.Dropout(dropout)
-        
-        # Transformer blocks
-        self.blocks = nn.ModuleList([
-            OptimizedTransformerBlock(
-                embed_dim, n_heads, None, None, dropout,
-                use_adapter=use_adapter, adapter_dim=adapter_dim
-            )
-            for _ in range(depth)
-        ])
-        
-        # 双流：时间流和频率流
-        if use_dual_stream:
-            # 修复：使用正确的池化维度
-            self.temporal_pool = nn.AdaptiveAvgPool1d(1)
-            self.frequency_pool = nn.AdaptiveAvgPool1d(1)
-            
-            self.temporal_proj = nn.Linear(embed_dim, embed_dim)
-            self.frequency_proj = nn.Linear(embed_dim, embed_dim)
-            
-            self.fusion = nn.Linear(embed_dim * 2, embed_dim)
-        
-        self.norm = RMSNorm(embed_dim)
-        self._init_weights()
-
-    def _init_weights(self):
-        nn.init.trunc_normal_(self.pos_embed, std=0.02)
-        self.apply(self._init_module_weights)
-
-    def _init_module_weights(self, m):
-        if isinstance(m, nn.Linear):
-            nn.init.trunc_normal_(m.weight, std=0.02)
-            if m.bias is not None:
-                nn.init.zeros_(m.bias)
-        elif isinstance(m, nn.Conv2d):
-            nn.init.trunc_normal_(m.weight, std=0.02)
-            if m.bias is not None:
-                nn.init.zeros_(m.bias)
-
-    def forward(self, mel_spec: torch.Tensor) -> torch.Tensor:
-        if mel_spec.ndim == 3:
-            mel_spec = mel_spec.unsqueeze(1)
-        
-        # Patch embedding
-        x = self.patch_embed(mel_spec)  # [B, C, H, W]
-        x = x.flatten(2).transpose(1, 2)  # [B, H*W, C]
-        x = self.pos_drop(x + self.pos_embed)
-        
-        # Transformer encoding
-        for block in self.blocks:
-            x, _, _ = block(x)
-        
-        x = self.norm(x)
-        
-        # 修复：双流处理
-        if self.use_dual_stream:
-            B, N, C = x.shape
-            
-            # 重塑为2D网格
-            x_2d = x.transpose(1, 2).reshape(B, C, self.n_patches_h, self.n_patches_w)
-            
-            # 时间流：沿频率维度池化（保留时间）
-            temporal = x_2d.mean(dim=2)  # [B, C, W]
-            temporal = self.temporal_pool(temporal).squeeze(-1)  # [B, C]
-            temporal = self.temporal_proj(temporal).unsqueeze(1)  # [B, 1, C]
-            
-            # 频率流：沿时间维度池化（保留频率）
-            frequency = x_2d.mean(dim=3)  # [B, C, H]
-            frequency = self.frequency_pool(frequency).squeeze(-1)  # [B, C]
-            frequency = self.frequency_proj(frequency).unsqueeze(1)  # [B, 1, C]
-            
-            # 融合
-            x = self.fusion(torch.cat([temporal, frequency], dim=-1))
-        else:
-            # 简单全局平均池化
-            x = x.mean(dim=1, keepdim=True)
-        
-        return x
-
-class ImprovedVideoEncoder(nn.Module):
-    """
-    改进的视频编码器
-    - 因果时序建模
-    - 时空分离注意力
-    - 可选的3D卷积
-    """
-    def __init__(
-        self,
-        img_size: int = 224,
-        patch_size: int = 14,
-        in_channels: int = 3,
-        embed_dim: int = 2048,
-        spatial_depth: int = 12,
-        temporal_depth: int = 4,
-        n_heads: int = 16,
-        num_frames: int = 16,
-        dropout: float = 0.1,
-        use_adapter: bool = False,
-        adapter_dim: int = 64,
-        use_3d_conv: bool = False
-    ):
-        super().__init__()
-        self.num_frames = num_frames
-        self.use_3d_conv = use_3d_conv
-        self.patch_size = patch_size
-        self.img_size = img_size
-        
-        if use_3d_conv:
-            # 3D卷积处理���空信息
-            self.patch_embed = nn.Conv3d(
-                in_channels,
-                embed_dim,
-                kernel_size=(2, patch_size, patch_size),
-                stride=(2, patch_size, patch_size)
-            )
-            # 修复：计算3D卷积后的尺寸
-            self.n_temporal_patches = num_frames // 2
-            self.n_spatial_patches = (img_size // patch_size) ** 2
-        else:
-            # 2D卷积 + 时序建模
-            self.patch_embed = ImprovedPatchEmbedding(
-                patch_size, in_channels, embed_dim
-            )
-            self.n_spatial_patches = (img_size // patch_size) ** 2
-        
-        # 空间位置编码
-        self.spatial_pos_embed = nn.Parameter(
-            torch.zeros(1, self.n_spatial_patches, embed_dim)
-        )
-        self.spatial_pos_drop = nn.Dropout(dropout)
-        
-        # 空间编码器
-        self.spatial_blocks = nn.ModuleList([
-            OptimizedTransformerBlock(
-                embed_dim, n_heads, None, None, dropout,
-                use_adapter=use_adapter, adapter_dim=adapter_dim
-            )
-            for _ in range(spatial_depth)
-        ])
-        
-        # 时间位置编码
-        if use_3d_conv:
-            self.temporal_pos_embed = nn.Parameter(
-                torch.zeros(1, self.n_temporal_patches, embed_dim)
-            )
-        else:
-            self.temporal_pos_embed = nn.Parameter(
-                torch.zeros(1, num_frames, embed_dim)
-            )
-        self.temporal_pos_drop = nn.Dropout(dropout)
-        
-        # 时序编码器
-        self.temporal_blocks = nn.ModuleList([
-            OptimizedTransformerBlock(
-                embed_dim, n_heads, None, None, dropout,
-                use_adapter=use_adapter, adapter_dim=adapter_dim
-            )
-            for _ in range(temporal_depth)
-        ])
-        
-        self.norm = RMSNorm(embed_dim)
-        self._init_weights()
-
-    def _init_weights(self):
-        nn.init.trunc_normal_(self.spatial_pos_embed, std=0.02)
-        nn.init.trunc_normal_(self.temporal_pos_embed, std=0.02)
-        self.apply(self._init_module_weights)
-
-    def _init_module_weights(self, m):
-        if isinstance(m, nn.Linear):
-            nn.init.trunc_normal_(m.weight, std=0.02)
-            if m.bias is not None:
-                nn.init.zeros_(m.bias)
-        elif isinstance(m, (nn.Conv2d, nn.Conv3d)):
-            nn.init.trunc_normal_(m.weight, std=0.02)
-            if m.bias is not None:
-                nn.init.zeros_(m.bias)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        B, T, C, H, W = x.shape
-        
-        if self.use_3d_conv:
-            # 修复：3D卷积路径
-            x = x.transpose(1, 2)  # [B, C, T, H, W]
-            x = self.patch_embed(x)  # [B, embed_dim, T', H', W']
-            
-            # 重塑: [B, D, T', H'*W'] -> [B, T', H'*W', D]
-            B, D, T_new, H_new, W_new = x.shape
-            x = x.view(B, D, T_new, -1).permute(0, 2, 3, 1)  # [B, T', H'*W', D]
-            
-            # 空间位置编码（每帧独立）
-            x = x + self.spatial_pos_embed.unsqueeze(1)
-            
-            # 逐帧空间编码
-            x_flat = x.reshape(B * T_new, -1, D)
-            for block in self.spatial_blocks:
-                x_flat, _, _ = block(x_flat)
-            
-            # 重塑回时序维度
-            x = x_flat.view(B, T_new, -1, D)
-            
-            # 修复：时序聚合 - 使用平均池化而非取第一个token
-            x = x.mean(dim=2)  # [B, T', D]
-            
-        else:
-            # 2D卷积 + 分离时空建模
-            x_flat = x.view(B * T, C, H, W)
-            x_patched, grid_size = self.patch_embed(x_flat)
-            
-            # 空间位置编码
-            x_patched = self.spatial_pos_drop(x_patched + self.spatial_pos_embed)
-            
-            # 空间编码
-            for block in self.spatial_blocks:
-                x_patched, _, _ = block(x_patched)
-            
-            # 修复：时序聚合 - 全局平均池化而非仅mean(dim=2)
-            _, N, D = x_patched.shape
-            x_spatial = x_patched.view(B, T, N, D)
-            x = x_spatial.mean(dim=2)  # [B, T, D] - 对每帧的所有patch取平均
-        
-        # 时序位置编码
-        x = self.temporal_pos_drop(x + self.temporal_pos_embed)
-        
-        # 时序编码
-        for block in self.temporal_blocks:
-            x, _, _ = block(x)
-        
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Tuple, Optional
+from components import RMSNorm, SwiGLU
+from transformer import OptimizedTransformerBlock
+import math
+
+class LayerScale(nn.Module):
+    def __init__(self, dim: int, init_values: float = 1e-5):
+        super().__init__()
+        self.gamma = nn.Parameter(init_values * torch.ones(dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x * self.gamma
+
+class StochasticDepth(nn.Module):
+    def __init__(self, drop_prob: float = 0.0):
+        super().__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if not self.training or self.drop_prob == 0.0:
+            return x
+        
+        keep_prob = 1 - self.drop_prob
+        shape = (x.shape[0],) + (1,) * (x.ndim - 1)
+        random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
+        random_tensor.floor_()
+        return x.div(keep_prob) * random_tensor
+
+class ImprovedPatchEmbedding(nn.Module):
+    def __init__(
+        self,
+        patch_size: int = 14,
+        in_channels: int = 3,
+        embed_dim: int = 2048,
+        overlap: int = 0
+    ):
+        super().__init__()
+        self.patch_size = patch_size
+        stride = patch_size - overlap
+        self.proj = nn.Conv2d(
+            in_channels,
+            embed_dim,
+            kernel_size=patch_size,
+            stride=stride,
+            padding=overlap // 2
+        )
+        
+        self.norm = RMSNorm(embed_dim)
+
+    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, Tuple[int, int]]:
+        B, C, H, W = x.shape
+        x = self.proj(x)
+        grid_size = (x.shape[2], x.shape[3])
+        x = x.flatten(2).transpose(1, 2)
+        x = self.norm(x)
+        return x, grid_size
+
+class ImprovedVisionBlock(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        n_heads: int,
+        dropout: float = 0.0,
+        drop_path: float = 0.0,
+        use_adapter: bool = False,
+        adapter_dim: int = 64,
+        use_layer_scale: bool = True,
+        layer_scale_init: float = 1e-5
+    ):
+        super().__init__()
+        self.norm1 = RMSNorm(dim)
+        self.attn = nn.MultiheadAttention(
+            dim, n_heads, dropout=dropout, batch_first=True
+        )
+        
+        self.norm2 = RMSNorm(dim)
+        self.mlp = nn.Sequential(
+            nn.Linear(dim, dim * 4),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(dim * 4, dim),
+            nn.Dropout(dropout)
+        )
+        
+        self.drop_path = StochasticDepth(drop_path) if drop_path > 0 else nn.Identity()
+        
+        if use_layer_scale:
+            self.ls1 = LayerScale(dim, layer_scale_init)
+            self.ls2 = LayerScale(dim, layer_scale_init)
+        else:
+            self.ls1 = nn.Identity()
+            self.ls2 = nn.Identity()
+        
+        if use_adapter:
+            self.adapter = nn.Sequential(
+                nn.Linear(dim, adapter_dim),
+                nn.GELU(),
+                nn.Linear(adapter_dim, dim)
+            )
+        else:
+            self.adapter = None
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # 注意力
+        normx = self.norm1(x)
+        attn_out, _ = self.attn(normx, normx, normx)
+        x = x + self.drop_path(self.ls1(attn_out))
+        
+        # MLP
+        x = x + self.drop_path(self.ls2(self.mlp(self.norm2(x))))
+        
+        # Adapter
+        if self.adapter is not None:
+            x = x + self.adapter(x)
+        
+        return x
+
+class ImprovedVisionTransformer(nn.Module):
+    def __init__(
+        self,
+        img_size: int = 224,
+        patch_size: int = 14,
+        in_channels: int = 3,
+        embed_dim: int = 2048,
+        depth: int = 24,
+        n_heads: int = 16,
+        dropout: float = 0.0,
+        drop_path_rate: float = 0.1,
+        use_register_tokens: bool = True,
+        num_register_tokens: int = 4,
+        use_adapter: bool = False,
+        adapter_dim: int = 64,
+        use_layer_scale: bool = True,
+        layer_scale_init: float = 1e-5
+    ):
+        super().__init__()
+        self.patch_size = patch_size
+        self.embed_dim = embed_dim
+        self.use_register_tokens = use_register_tokens
+        self.num_register_tokens = num_register_tokens if use_register_tokens else 0
+        
+        # Patch embedding
+        self.patch_embed = ImprovedPatchEmbedding(
+            patch_size, in_channels, embed_dim, overlap=0
+        )
+        
+        self.pretrain_img_size = img_size
+        n_patches_pretrain = (img_size // patch_size) ** 2
+        
+        # CLS token
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        
+        # Register tokens
+        if use_register_tokens:
+            self.register_tokens = nn.Parameter(
+                torch.zeros(1, num_register_tokens, embed_dim)
+            )
+        
+        total_tokens = 1 + n_patches_pretrain + self.num_register_tokens
+        self.pos_embed = nn.Parameter(
+            torch.zeros(1, total_tokens, embed_dim)
+        )
+        self.pos_drop = nn.Dropout(dropout)
+        
+        # Stochastic depth
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]
+        
+        # Transformer blocks
+        self.blocks = nn.ModuleList([
+            ImprovedVisionBlock(
+                embed_dim,
+                n_heads,
+                dropout,
+                drop_path=dpr[i],
+                use_adapter=use_adapter,
+                adapter_dim=adapter_dim,
+                use_layer_scale=use_layer_scale,
+                layer_scale_init=layer_scale_init
+            )
+            for i in range(depth)
+        ])
+        
+        self.norm = RMSNorm(embed_dim)
+        self._init_weights()
+
+    def _init_weights(self):
+        nn.init.trunc_normal_(self.cls_token, std=0.02)
+        nn.init.trunc_normal_(self.pos_embed, std=0.02)
+        if self.use_register_tokens:
+            nn.init.trunc_normal_(self.register_tokens, std=0.02)
+        
+        self.apply(self._init_module_weights)
+
+    def _init_module_weights(self, m):
+        if isinstance(m, nn.Linear):
+            nn.init.trunc_normal_(m.weight, std=0.02)
+            if m.bias is not None:
+                nn.init.zeros_(m.bias)
+        elif isinstance(m, nn.Conv2d):
+            nn.init.trunc_normal_(m.weight, std=0.02)
+            if m.bias is not None:
+                nn.init.zeros_(m.bias)
+        elif isinstance(m, RMSNorm):
+            if hasattr(m, 'weight') and m.weight is not None:
+                nn.init.ones_(m.weight)
+
+    def _interpolate_pos_encoding(
+        self,
+        patch_tokens: torch.Tensor,
+        grid_size: Tuple[int, int]
+    ) -> torch.Tensor:
+        pretrain_grid_h = self.pretrain_img_size // self.patch_size
+        pretrain_grid_w = pretrain_grid_h
+        
+        # 如果尺寸匹配，直接返回原始位置编码
+        if grid_size[0] == pretrain_grid_h and grid_size[1] == pretrain_grid_w:
+            return self.pos_embed
+        
+        # 分离不同部分的位置编码
+        # pos_embed结构: [CLS(1), register_tokens(n), patches(H*W)]
+        num_extra_tokens = 1 + self.num_register_tokens
+        cls_register_pos = self.pos_embed[:, :num_extra_tokens, :]  # [1, 1+n, dim]
+        patch_pos_embed = self.pos_embed[:, num_extra_tokens:, :]  # [1, H*W, dim]
+        
+        # 2D插值patch位置编码
+        patch_pos_embed = patch_pos_embed.reshape(
+            1, pretrain_grid_h, pretrain_grid_w, -1
+        ).permute(0, 3, 1, 2)
+        
+        patch_pos_embed = F.interpolate(
+            patch_pos_embed,
+            size=grid_size,
+            mode='bicubic',
+            align_corners=False
+        )
+        
+        patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).flatten(1, 2)
+        
+        # 拼接回去
+        return torch.cat([cls_register_pos, patch_pos_embed], dim=1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B = x.shape[0]
+        
+        # Patch embedding
+        x, grid_size = self.patch_embed(x)
+        
+        # 添加CLS token
+        cls_tokens = self.cls_token.expand(B, -1, -1)
+        
+        if self.use_register_tokens:
+            register_tokens = self.register_tokens.expand(B, -1, -1)
+            # 顺序: [CLS, register_tokens, patches]
+            x = torch.cat([cls_tokens, register_tokens, x], dim=1)
+        else:
+            x = torch.cat([cls_tokens, x], dim=1)
+        
+        # 位置编码
+        pos_embed = self._interpolate_pos_encoding(x, grid_size)
+        x = self.pos_drop(x + pos_embed)
+        
+        # Transformer blocks
+        for block in self.blocks:
+            x = block(x)
+        
+        x = self.norm(x)
+        
+        return x
+
+class ImprovedAudioEncoder(nn.Module):
+    def __init__(
+        self,
+        n_mels: int = 128,
+        target_length: int = 1024,
+        embed_dim: int = 2048,
+        depth: int = 12,
+        n_heads: int = 16,
+        patch_size: int = 16,
+        dropout: float = 0.1,
+        use_adapter: bool = False,
+        adapter_dim: int = 64,
+        use_dual_stream: bool = True
+    ):
+        super().__init__()
+        self.use_dual_stream = use_dual_stream
+        self.patch_size = patch_size
+        
+        # 主编码器
+        self.patch_embed = nn.Conv2d(
+            1, embed_dim, kernel_size=patch_size, stride=patch_size
+        )
+        
+        self.n_patches_h = n_mels // patch_size
+        self.n_patches_w = target_length // patch_size
+        n_patches = self.n_patches_h * self.n_patches_w
+        
+        self.pos_embed = nn.Parameter(torch.zeros(1, n_patches, embed_dim))
+        self.pos_drop = nn.Dropout(dropout)
+        
+        # Transformer blocks
+        self.blocks = nn.ModuleList([
+            OptimizedTransformerBlock(
+                embed_dim, n_heads, None, None, dropout,
+                use_adapter=use_adapter, adapter_dim=adapter_dim
+            )
+            for _ in range(depth)
+        ])
+        
+        # 双流：时间流和频率流
+        if use_dual_stream:
+            self.temporal_pool = nn.AdaptiveAvgPool1d(1)
+            self.frequency_pool = nn.AdaptiveAvgPool1d(1)
+            
+            self.temporal_proj = nn.Linear(embed_dim, embed_dim)
+            self.frequency_proj = nn.Linear(embed_dim, embed_dim)
+            
+            self.fusion = nn.Linear(embed_dim * 2, embed_dim)
+        
+        self.norm = RMSNorm(embed_dim)
+        self._init_weights()
+
+    def _init_weights(self):
+        nn.init.trunc_normal_(self.pos_embed, std=0.02)
+        self.apply(self._init_module_weights)
+
+    def _init_module_weights(self, m):
+        if isinstance(m, nn.Linear):
+            nn.init.trunc_normal_(m.weight, std=0.02)
+            if m.bias is not None:
+                nn.init.zeros_(m.bias)
+        elif isinstance(m, nn.Conv2d):
+            nn.init.trunc_normal_(m.weight, std=0.02)
+            if m.bias is not None:
+                nn.init.zeros_(m.bias)
+
+    def forward(self, mel_spec: torch.Tensor) -> torch.Tensor:
+        if mel_spec.ndim == 3:
+            mel_spec = mel_spec.unsqueeze(1)
+        
+        # Patch embedding
+        x = self.patch_embed(mel_spec)  # [B, C, H, W]
+        x = x.flatten(2).transpose(1, 2)  # [B, H*W, C]
+        x = self.pos_drop(x + self.pos_embed)
+        
+        # Transformer encoding
+        for block in self.blocks:
+            x, _, _ = block(x)
+        
+        x = self.norm(x)
+        
+        if self.use_dual_stream:
+            B, N, C = x.shape
+            
+            # 重塑为2D网格
+            x_2d = x.transpose(1, 2).reshape(B, C, self.n_patches_h, self.n_patches_w)
+            
+            # 时间流：沿频率维度池化（保留时间）
+            temporal = x_2d.mean(dim=2)  # [B, C, W]
+            temporal = self.temporal_pool(temporal).squeeze(-1)  # [B, C]
+            temporal = self.temporal_proj(temporal).unsqueeze(1)  # [B, 1, C]
+            
+            # 频率流：沿时间维度池化（保留频率）
+            frequency = x_2d.mean(dim=3)  # [B, C, H]
+            frequency = self.frequency_pool(frequency).squeeze(-1)  # [B, C]
+            frequency = self.frequency_proj(frequency).unsqueeze(1)  # [B, 1, C]
+            
+            # 融合
+            x = self.fusion(torch.cat([temporal, frequency], dim=-1))
+        else:
+            # 简单全局平均池化
+            x = x.mean(dim=1, keepdim=True)
+        
+        return x
+
+class ImprovedVideoEncoder(nn.Module):
+    def __init__(
+        self,
+        img_size: int = 224,
+        patch_size: int = 14,
+        in_channels: int = 3,
+        embed_dim: int = 2048,
+        spatial_depth: int = 12,
+        temporal_depth: int = 4,
+        n_heads: int = 16,
+        num_frames: int = 16,
+        dropout: float = 0.1,
+        use_adapter: bool = False,
+        adapter_dim: int = 64,
+        use_3d_conv: bool = False
+    ):
+        super().__init__()
+        self.num_frames = num_frames
+        self.use_3d_conv = use_3d_conv
+        self.patch_size = patch_size
+        self.img_size = img_size
+        
+        if use_3d_conv:
+            # 3D卷积处理时空信息
+            self.patch_embed = nn.Conv3d(
+                in_channels,
+                embed_dim,
+                kernel_size=(2, patch_size, patch_size),
+                stride=(2, patch_size, patch_size)
+            )
+            self.n_temporal_patches = num_frames // 2
+            self.n_spatial_patches = (img_size // patch_size) ** 2
+        else:
+            # 2D卷积 + 时序建模
+            self.patch_embed = ImprovedPatchEmbedding(
+                patch_size, in_channels, embed_dim
+            )
+            self.n_spatial_patches = (img_size // patch_size) ** 2
+        
+        # 空间位置编码
+        self.spatial_pos_embed = nn.Parameter(
+            torch.zeros(1, self.n_spatial_patches, embed_dim)
+        )
+        self.spatial_pos_drop = nn.Dropout(dropout)
+        
+        # 空间编码器
+        self.spatial_blocks = nn.ModuleList([
+            OptimizedTransformerBlock(
+                embed_dim, n_heads, None, None, dropout,
+                use_adapter=use_adapter, adapter_dim=adapter_dim
+            )
+            for _ in range(spatial_depth)
+        ])
+        
+        # 时间位置编码
+        if use_3d_conv:
+            self.temporal_pos_embed = nn.Parameter(
+                torch.zeros(1, self.n_temporal_patches, embed_dim)
+            )
+        else:
+            self.temporal_pos_embed = nn.Parameter(
+                torch.zeros(1, num_frames, embed_dim)
+            )
+        self.temporal_pos_drop = nn.Dropout(dropout)
+        
+        # 时序编码器
+        self.temporal_blocks = nn.ModuleList([
+            OptimizedTransformerBlock(
+                embed_dim, n_heads, None, None, dropout,
+                use_adapter=use_adapter, adapter_dim=adapter_dim
+            )
+            for _ in range(temporal_depth)
+        ])
+        
+        self.norm = RMSNorm(embed_dim)
+        self._init_weights()
+
+    def _init_weights(self):
+        nn.init.trunc_normal_(self.spatial_pos_embed, std=0.02)
+        nn.init.trunc_normal_(self.temporal_pos_embed, std=0.02)
+        self.apply(self._init_module_weights)
+
+    def _init_module_weights(self, m):
+        if isinstance(m, nn.Linear):
+            nn.init.trunc_normal_(m.weight, std=0.02)
+            if m.bias is not None:
+                nn.init.zeros_(m.bias)
+        elif isinstance(m, (nn.Conv2d, nn.Conv3d)):
+            nn.init.trunc_normal_(m.weight, std=0.02)
+            if m.bias is not None:
+                nn.init.zeros_(m.bias)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B, T, C, H, W = x.shape
+        
+        if self.use_3d_conv:
+            x = x.transpose(1, 2)  # [B, C, T, H, W]
+            x = self.patch_embed(x)  # [B, embed_dim, T', H', W']
+            
+            # 重塑: [B, D, T', H'*W'] -> [B, T', H'*W', D]
+            B, D, T_new, H_new, W_new = x.shape
+            x = x.view(B, D, T_new, -1).permute(0, 2, 3, 1)  # [B, T', H'*W', D]
+            
+            # 空间位置编码（每帧独立）
+            x = x + self.spatial_pos_embed.unsqueeze(1)
+            
+            # 逐帧空间编码
+            x_flat = x.reshape(B * T_new, -1, D)
+            for block in self.spatial_blocks:
+                x_flat, _, _ = block(x_flat)
+            
+            # 重塑回时序维度
+            x = x_flat.view(B, T_new, -1, D)
+            x = x.mean(dim=2)  # [B, T', D]
+            
+        else:
+            # 2D卷积 + 分离时空建模
+            x_flat = x.view(B * T, C, H, W)
+            x_patched, grid_size = self.patch_embed(x_flat)
+            
+            # 空间位置编码
+            x_patched = self.spatial_pos_drop(x_patched + self.spatial_pos_embed)
+            
+            # 空间编码
+            for block in self.spatial_blocks:
+                x_patched, _, _ = block(x_patched)
+            
+            _, N, D = x_patched.shape
+            x_spatial = x_patched.view(B, T, N, D)
+            x = x_spatial.mean(dim=2)  
+        
+        # 时序位置编码
+        x = self.temporal_pos_drop(x + self.temporal_pos_embed)
+        
+        # 时序编码
+        for block in self.temporal_blocks:
+            x, _, _ = block(x)
+        
         return self.norm(x)