多模态融合的「流式革命」：CaReFlow 用校正流打通模态鸿沟

# 传统 One-to-One：只采样一个配对
paired_sample = sample_one(language_distribution, visual_point)

# CaReFlow One-to-Many：采样多个目标点，感知全局
multiple_targets = sample_k(language_distribution, visual_point, k=5)
# 每个视觉特征「看到」语言分布的多个角落

# 自适应权重计算
def adaptive_weight(sample_i, sample_j, modality_m, modality_n):
    if same_sample(i, j):
        return1.0# 严格对齐
    elif same_category(i, j):
        return0.5# 中等松弛
    else:
        return0.1# 高度松弛

# 损失函数
L_flow = sum(
    adaptive_weight(i, j) * ||(X_j^n - X_i^m) - v(X_i^m, t)||^2
    for i, j in all_pairs
)

# Forward Flow: 视觉 → 语言
X_v_to_l = rectified_flow_forward(X_v, t)

# Backward Flow: 语言 → 视觉
X_v_recon = rectified_flow_backward(X_v_to_l, t)

# Cycle Consistency Loss
L_cyclic = ||X_v - X_v_recon||^2

import torch
import torch.nn as nn

class UnimodalEncoder(nn.Module):
    def __init__(self, modality, d_model=768):
        super().__init__()
        # 模态特定提取器
        if modality == 'visual':
            self.extractor = CNNBackbone()  # 或 Video Swin Transformer
        elif modality == 'audio':
            self.extractor = VGGish()  # 或 HuBERT
        elif modality == 'text':
            self.extractor = BertModel.from_pretrained('bert-base')
        
        # 投影到共享空间
        self.projection = nn.Linear(self.extractor.dim, d_model)
    
    def forward(self, x):
        features = self.extractor(x)  # [B, T_m, d_m]
        # 时序聚合（mean pooling 或 [CLS]）
        x_unimodal = features.mean(dim=1)  # [B, d_m]
        x_shared = self.projection(x_unimodal)  # [B, d]
        return x_shared

# ODE: dZ/dt = v(Z, t),  t ∈ [0,1]
# 初始: Z_0 ~ π₀ (源模态分布)
# 终止: Z_1 ~ π₁ (目标模态分布)

# 训练目标：最小化插值直线与速度场的差异
# X_t = t*X_1 + (1-t)*X_0  (线性插值)
# L_flow = E[||(X_1 - X_0) - v(X_t, t)||^2]

class RectifiedFlow(nn.Module):
    def __init__(self, d_model):
        super().__init__()
        # 速度场网络：输入 [Z_t, t]，输出速度 v
        self.velocity_net = nn.Sequential(
            nn.Linear(d_model + 1, 512),
            nn.ReLU(),
            nn.Linear(512, 512),
            nn.ReLU(),
            nn.Linear(512, d_model)
        )
    
    def forward(self, z_t, t):
        # z_t: [B, d], t: [B, 1]
        input_vec = torch.cat([z_t, t], dim=-1)
        return self.velocity_net(input_vec)  # [B, d]
    
    def sample(self, z_0, num_steps=10):
        # 从源分布采样，沿 ODE 演化到目标分布
        dt = 1.0 / num_steps
        z = z_0
        for i in range(num_steps):
            t = torch.ones(z.shape[0], 1) * (i * dt)
            v = self.forward(z, t)
            z = z + v * dt  # Euler 积分
        return z  # z ~ π₁

class OneToManyMapping(nn.Module):
    def __init__(self, d_model, k=5):
        super().__init__()
        self.k = k  # 每个源点映射到 k 个目标点
        self.flow = RectifiedFlow(d_model)
    
    def forward(self, x_source, x_target_pool):
        """
        x_source: [B, d] - 源模态特征
        x_target_pool: [N, d] - 目标模态的全局分布样本（N >> B）
        """
        B = x_source.shape[0]
        
        # 为每个源点采样 k 个目标点（基于相似度或随机）
        # 策略：从目标分布中采样 k 个最近邻
        similarities = torch.matmul(x_source, x_target_pool.T)  # [B, N]
        topk_indices = torch.topk(similarities, self.k, dim=-1).indices  # [B, k]
        
        losses = []
        for i in range(B):
            for j in range(self.k):
                target_idx = topk_indices[i, j]
                x_src = x_source[i:i+1]      # [1, d]
                x_tgt = x_target_pool[target_idx:target_idx+1]  # [1, d]
                
                # 构建插值轨迹
                t = torch.rand(1, 1)  # 随机采样时间步
                x_t = t * x_tgt + (1 - t) * x_src
                
                # 计算 flow loss
                v_pred = self.flow(x_t, t)
                v_gt = x_tgt - x_src
                loss = F.mse_loss(v_pred, v_gt)
                losses.append(loss)
        
        return torch.stack(losses).mean()

class AdaptiveRelaxedAlignment(nn.Module):
    def __init__(self, d_model):
        super().__init__()
        self.flow = RectifiedFlow(d_model)
    
    def compute_weight(self, sample_i, sample_j, labels):
        """
        根据样本关系计算对齐权重
        """
        if sample_i == sample_j:
            return1.0# 同一样本：严格对齐
        elif labels[sample_i] == labels[sample_j]:
            return0.5# 同类别不同样本：中等松弛
        else:
            return0.1# 不同类别：高度松弛
    
    def forward(self, x_source, x_target, sample_ids, labels):
        B = x_source.shape[0]
        total_loss = 0
        total_weight = 0
        
        for i in range(B):
            for j in range(B):
                weight = self.compute_weight(sample_ids[i], sample_ids[j], labels)
                if weight < 0.01:
                    continue
                
                # 构建插值
                t = torch.rand(1, 1)
                x_t = t * x_target[j:j+1] + (1 - t) * x_source[i:i+1]
                
                # Flow loss
                v_pred = self.flow(x_t, t)
                v_gt = x_target[j:j+1] - x_source[i:i+1]
                loss = weight * F.mse_loss(v_pred, v_gt)
                
                total_loss += loss
                total_weight += weight
        
        return total_loss / total_weight if total_weight > 0else total_loss

class CyclicRectifiedFlow(nn.Module):
    def __init__(self, d_model):
        super().__init__()
        # Forward: 源 → 目标
        self.flow_forward = RectifiedFlow(d_model)
        # Backward: 目标 → 源
        self.flow_backward = RectifiedFlow(d_model)
    
    def forward(self, x_source, x_target):
        # Forward flow: 源模态映射到目标模态空间
        t_fwd = torch.rand(x_source.shape[0], 1)
        x_t_fwd = t_fwd * x_target + (1 - t_fwd) * x_source
        v_fwd = self.flow_forward(x_t_fwd, t_fwd)
        x_mapped = x_source + v_fwd  # 近似 Z_1
        
        # Backward flow: 映射后的特征还原为源模态
        t_bwd = torch.rand(x_mapped.shape[0], 1)
        x_t_bwd = t_bwd * x_source + (1 - t_bwd) * x_mapped
        v_bwd = self.flow_backward(x_t_bwd, t_bwd)
        x_recon = x_mapped + v_bwd  # 近似还原
        
        # Cycle consistency loss
        L_cyclic = F.mse_loss(x_recon, x_source)
        
        # Forward flow loss（正常训练）
        v_gt_fwd = x_target - x_source
        L_flow = F.mse_loss(v_fwd, v_gt_fwd)
        
        return L_flow + 0.5 * L_cyclic, x_mapped

# Step 1: 准备多模态数据
# 视频 → CNN/Transformer → [B, d_v]
# 音频 → VGGish/HuBERT → [B, d_a]  
# 文本 → BERT → [B, d_l]

# Step 2: 投影到共享空间
projectors = {
    'visual': nn.Linear(d_v, d_shared),
    'audio': nn.Linear(d_a, d_shared),
    'text': nn.Linear(d_l, d_shared)
}

# Step 3: 初始化 CaReFlow
careflow = CaReFlow(
    d_model=d_shared,
    k=5,           # One-to-Many 采样数
    lambda_cyclic=0.5,
    lambda_adaptive=1.0
)

# Step 4: 训练（关键：构造 sample_id 和 label 用于自适应权重）
for batch in dataloader:
    x_v = projectors['visual'](batch.video)      # [B, d]
    x_a = projectors['audio'](batch.audio)        # [B, d]
    x_l = projectors['text'](batch.text)          # [B, d]
    
    # 视觉 → 语言对齐
    loss_v2l, x_v_aligned = careflow(x_v, x_l, batch.sample_id, batch.label)
    # 音频 → 语言对齐  
    loss_a2l, x_a_aligned = careflow(x_a, x_l, batch.sample_id, batch.label)
    
    # 简单融合
    fused = torch.cat([x_v_aligned, x_a_aligned, x_l], dim=-1)
    pred = mlp(fusion)
    
    # 总损失
    loss = task_loss(pred, batch.target) + loss_v2l + loss_a2l
    loss.backward()