from omegaconf import DictConfig from tqdm import tqdm import torch.nn.functional as F import clip.clip as clip import torch import torch.nn as nn from torch.utils.data import DataLoader from .utils import get_class_ids_per_task, get_class_names, batch, merge_we_router, wise_we, moving_avg, l2_loss, \ virtual_vocab from .cc import conceptual_captions from . import utils import os import random from .dynamic_dataset import DynamicDataset DFA_R1 = 16 DFA_R2 = 8 DFA_ADAPTER_DROPOUT = 0.1 DFA_ADAPTER_SCALAR = 0.1 GENERIC_TEMPLATES = ( "a photo of a {}", "a clean photo of a {}", "a bad photo of the {}", "a professional photo of a {}", "a studio shot of a {}", "a photo of a {} a realistic setting", "an image of a {}", "a photo of a {} in a photo book", "an image of a {} by a professional photographer", "a photo of a {} using in an experiment", "a photo of a {} by a camera", "a photo of a nice {}", "a photo of the nice {}", "a photo of many {}", "a photo of the cool {}", "a photo of the {}", "art of the {}", "a blurry photo of the {}", "a photo of the {} in a photo book", "a pixelated photo of the {}", ) class DFA3Block(nn.Module): """Per-transformer-block DFA-MoE adapter with hierarchical routing. The block follows the paper's functional decoupling: - `e1`: task-agnostic alignment expert (Alignment Pathway) - `e2_list`: task-specific expert group (Plasticity Pathway) - `w_gate_e2` / `w_noise_e2`: inner router over task-specific experts - `router_top`: outer router balancing alignment vs. plasticity outputs Routing is computed from the CLS token and broadcast to the full token sequence. The returned tensor is a residual branch added at the FFN segment. """ def __init__(self, embed_dim: int, e2_top_k: int = 2, num_task_experts: int = 2): super().__init__() self.embed_dim = int(embed_dim) self.e2_top_k = max(1, int(e2_top_k)) self.num_task_experts = max(1, int(num_task_experts)) # Alignment Pathway: task-agnostic expert E_a. self.e1 = nn.Sequential( nn.LayerNorm(self.embed_dim), nn.Linear(self.embed_dim, DFA_R1), nn.ReLU(inplace=True), nn.Dropout(p=DFA_ADAPTER_DROPOUT), nn.Linear(DFA_R1, self.embed_dim), ) # Plasticity Pathway: task-specific expert group E_s. self.e2_list = nn.ModuleList([ nn.Sequential( nn.LayerNorm(self.embed_dim), nn.Linear(self.embed_dim, DFA_R2), nn.ReLU(inplace=True), nn.Dropout(p=DFA_ADAPTER_DROPOUT), nn.Linear(DFA_R2, self.embed_dim), ) for _ in range(self.num_task_experts) ]) with torch.no_grad(): # Start from an identity-like residual path. if isinstance(self.e1[-1], nn.Linear): nn.init.zeros_(self.e1[-1].weight) nn.init.zeros_(self.e1[-1].bias) for mod in self.e2_list: if isinstance(mod[-1], nn.Linear): nn.init.zeros_(mod[-1].weight) nn.init.zeros_(mod[-1].bias) self.router_top = nn.Sequential( nn.LayerNorm(self.embed_dim), nn.Linear(self.embed_dim, 2), ) self.w_gate_e2 = nn.Linear(self.embed_dim, self.num_task_experts, bias=False) self.w_noise_e2 = nn.Linear(self.embed_dim, self.num_task_experts, bias=False) def _noisy_topk_cls(self, cls_x: torch.Tensor, train: bool = True, noise_epsilon: float = 1e-2) -> torch.Tensor: clean = self.w_gate_e2(cls_x) if train: raw = self.w_noise_e2(cls_x) noise_std = F.softplus(raw) + float(noise_epsilon) logits = clean + torch.randn_like(clean) * noise_std else: logits = clean E = logits.size(-1) k = min(int(self.e2_top_k), int(E)) top_logits, top_idx = logits.topk(k, dim=-1) top_w = torch.softmax(top_logits, dim=-1) if top_w.dtype != logits.dtype: top_w = top_w.to(dtype=logits.dtype) gates = torch.zeros_like(logits, dtype=logits.dtype) gates.scatter_(-1, top_idx, top_w) return gates # [N, E] def forward(self, pre_ffn: torch.Tensor) -> torch.Tensor: # Input to the DFA branch at the FFN segment: [L, N, D]. # Stage A trains only the alignment expert, so routing is bypassed. if getattr(self, 'e1_only', False): return DFA_ADAPTER_SCALAR * self.e1(pre_ffn) # Stage B / inference: combine both pathways via hierarchical routing. cls_feat = pre_ffn.permute(1, 0, 2)[:, 0, :] e1_res = DFA_ADAPTER_SCALAR * self.e1(pre_ffn) e2_res_list = [DFA_ADAPTER_SCALAR * m(pre_ffn) for m in self.e2_list] w_top = torch.softmax(self.router_top(cls_feat.float()), dim=-1) # [N,2] w_spec = self._noisy_topk_cls(cls_feat.float(), train=self.training) # [N,E] w_top_b = w_top.unsqueeze(0).unsqueeze(-1) w_spec_b = w_spec.unsqueeze(0).unsqueeze(-1) e2_stack = torch.stack(e2_res_list, dim=2) e2_res = (e2_stack * w_spec_b).sum(dim=2) res_stack = torch.stack([e1_res, e2_res], dim=2) res = (res_stack * w_top_b).sum(dim=2) return res class ResidualAttentionBlockWithDFA(nn.Module): """Inject the DFA-MoE residual branch into a CLIP transformer block. The wrapper preserves the original attention residual and augments the FFN residual with the per-block DFA adapter output. """ def __init__(self, base_block: nn.Module, dfa_block: nn.Module): super().__init__() self.base = base_block self.dfa = dfa_block def forward(self, x: torch.Tensor) -> torch.Tensor: # `x` has shape [L, N, D]. x = x + self.base.attention(self.base.ln_1(x)) # The DFA branch receives the representation before the FFN residual. pre_ffn = x adapter_out = self.dfa(pre_ffn) x = x + self.base.mlp(self.base.ln_2(x)) + adapter_out return x class ClassIncremental(nn.Module): def __init__(self, cfg, device, jit=False): super().__init__() self.prompt_template = cfg.prompt_template self.device = device self.classes_names = None self.model, self.transforms, _ = clip.load(cfg.model_name, device=device, jit=jit) self.ref_model = None # When class order is not precomputed, derive the split directly from data. if getattr(cfg, 'class_order', None) is not None: self.class_ids_per_task = list(get_class_ids_per_task(cfg)) else: self.class_ids_per_task = None self.current_class_names = [] # Cumulative absolute class ids aligned with `current_class_names` / `text_tokens`. self.seen_class_ids = [] self.text_tokens = None self.dynamic_dataset = DynamicDataset(cfg) visual_embed_dim = getattr(getattr(self.model, 'visual', None), 'output_dim', None) if visual_embed_dim is None: try: dummy = torch.zeros( 1, 3, getattr(cfg, 'input_resolution', 224), getattr(cfg, 'input_resolution', 224), device=device, ) with torch.no_grad(): visual_embed_dim = self.model.encode_image(dummy).shape[-1] except Exception: visual_embed_dim = getattr(self.model, 'embed_dim', 1024) self.dfa_blocks = [] self._inject_dfa_into_visual_blocks(cfg) # Mirror the same per-block DFA structure on the text transformer. self.text_dfa_blocks = [] self._inject_dfa_into_text_blocks(cfg) # Historical feature queue shared by the alignment and plasticity losses. self.queue_size = int(getattr(cfg, 'moco_queue_size', 4096)) self.register_buffer('queue_img', torch.randn(visual_embed_dim, self.queue_size)) self.register_buffer('queue_txt', torch.randn(visual_embed_dim, self.queue_size)) self.register_buffer('queue_ptr', torch.zeros(1, dtype=torch.long)) self.queue_img = nn.functional.normalize(self.queue_img, dim=0) self.queue_txt = nn.functional.normalize(self.queue_txt, dim=0) def forward(self, image, _taskid): with torch.no_grad(): logits_per_image, _ = self.model(image, self.text_tokens, 0, is_train=False) probs = logits_per_image.softmax(dim=-1) return probs def adaptation(self, task_id, cfg, train_dataset, _train_classes_names): # Derive task-local class ids from the scenario slice to avoid config drift. task_slice = train_dataset[task_id:task_id + 1] tmp_loader = DataLoader(task_slice, batch_size=256, shuffle=False, num_workers=2) uniq = set() for _inputs, _targets, _tids in tmp_loader: uniq.update(_targets.tolist()) real_ids = sorted(int(x) for x in uniq) # Record the insertion offset used by downstream evaluation. start_idx = len(self.current_class_names) # Materialize the discovered class split for this incremental task. if self.class_ids_per_task is None: self.class_ids_per_task = [] while len(self.class_ids_per_task) < task_id: self.class_ids_per_task.append([]) if len(self.class_ids_per_task) == task_id: self.class_ids_per_task.append(real_ids) else: self.class_ids_per_task[task_id] = real_ids # Append new classes in the same order used for cumulative text prompts. self.current_class_names += get_class_names(self.classes_names, real_ids) self.seen_class_ids += list(real_ids) self.text_tokens = clip.tokenize( [self.prompt_template.format(c) for c in self.current_class_names] ).to(self.device) # Expose the latest task mapping for the evaluation loop. self.last_task_real_ids = list(real_ids) self.last_task_start_index = int(start_idx) if cfg.method != "zeroshot": self.train(task_id, cfg, train_dataset, _train_classes_names) def train(self, task_id, cfg, train_dataset, _train_classes_names): # Current incremental task dataloader. train_loader = DataLoader( train_dataset[task_id:task_id + 1], batch_size=cfg.batch_size, shuffle=True, num_workers=min(8, os.cpu_count() or 8), pin_memory=True, ) # The paper optimizes the two pathways in separate stages. epochs_global = int(getattr(cfg, 'epochs', 1)) epochs_a = int(getattr(cfg, 'epochs_a', epochs_global)) epochs_b = int(getattr(cfg, 'epochs_b', epochs_global)) total_iterations_a = max(1, epochs_a * len(train_loader)) total_iterations_b = max(1, epochs_b * len(train_loader)) # Remap absolute dataset labels to task-local indices. task_class_ids = [int(c) for c in self.class_ids_per_task[task_id]] local_C = len(task_class_ids) max_cid = max(task_class_ids) map_table_cpu = torch.full((max_cid + 1,), -1, dtype=torch.long) for i, cid in enumerate(task_class_ids): map_table_cpu[cid] = i # Current-task class prompts used by the plasticity pathway. classnames = get_class_names(self.classes_names, self.class_ids_per_task[task_id]) texts_task = clip.tokenize([self.prompt_template.format(c) for c in classnames]).to(self.device) # Alignment-pathway negatives are drawn from seen classes and generic prompts. all_seen_names = [] for tid in range(task_id + 1): all_seen_names.extend(get_class_names(self.classes_names, self.class_ids_per_task[tid])) generic_templates = GENERIC_TEMPLATES # Backbone CLIP remains frozen throughout continual adaptation. for p in self.model.parameters(): p.requires_grad = False # -------- Stage A: Alignment Pathway -------- train_text_e1_stage_a = hasattr(self, 'text_dfa_blocks') and bool(self.text_dfa_blocks) # Train only the task-agnostic expert E_a and disable all routing. for m in getattr(self, 'dfa_blocks', []): m.e1_only = True for p in m.e1.parameters(): p.requires_grad = True for e2 in getattr(m, 'e2_list', []): for p in e2.parameters(): p.requires_grad = False for p in m.w_gate_e2.parameters(): p.requires_grad = False for p in m.w_noise_e2.parameters(): p.requires_grad = False for p in m.router_top.parameters(): p.requires_grad = False # Optionally optimize the text-side alignment adapters with the same loss. if train_text_e1_stage_a: for m in self.text_dfa_blocks: m.e1_only = True for p in m.e1.parameters(): p.requires_grad = True for e2 in getattr(m, 'e2_list', []): for p in e2.parameters(): p.requires_grad = False for p in m.w_gate_e2.parameters(): p.requires_grad = False for p in m.w_noise_e2.parameters(): p.requires_grad = False for p in m.router_top.parameters(): p.requires_grad = False lr_a = float(getattr(cfg, 'lr_e1', getattr(cfg, 'lr', 1e-4))) stage_a_base_lrs = [lr_a] # Collect alignment-expert parameters for the Stage-A optimizer. img_params_a = [] for m in getattr(self, 'dfa_blocks', []): img_params_a += list(m.e1.parameters()) text_params_a = [] text_lr_a = float(getattr(cfg, 'text_lr_e1', lr_a)) if train_text_e1_stage_a: for m in self.text_dfa_blocks: text_params_a += list(m.e1.parameters()) param_groups_a = [] if img_params_a: param_groups_a.append({ "params": img_params_a, "lr": lr_a, "weight_decay": float(getattr(cfg, 'weight_decay', 0.0)), }) if text_params_a: param_groups_a.append({ "params": text_params_a, "lr": text_lr_a, "weight_decay": float(getattr(cfg, 'weight_decay', 0.0)), }) stage_a_base_lrs.append(text_lr_a) if not param_groups_a: raise RuntimeError("No parameters collected for Stage A optimization in block mode") opt_a = torch.optim.AdamW(param_groups_a) sched_a = utils.cosine_lr(opt_a, stage_a_base_lrs if len(stage_a_base_lrs) > 1 else stage_a_base_lrs[0], 30, total_iterations_a) tau_con = float(getattr(cfg, 'tau_con', 0.05)) use_moco = bool(getattr(cfg, 'use_moco_queue', True)) # Buffer current-task positive pairs and refresh the queue once per task. mq_img_buf: list = [] mq_txt_buf: list = [] self.model.eval() it = 0 for epoch in range(epochs_a): for inputs, targets_abs, _ in tqdm(train_loader, desc=f"Task {task_id} A (E1)"): sched_a(it) it += 1 inputs = inputs.to(self.device, non_blocking=True) targets_abs = targets_abs.to(self.device, non_blocking=True) # Convert absolute labels to task-local labels. map_table = map_table_cpu.to(targets_abs.device) targets = map_table[targets_abs] if (targets < 0).any(): raise RuntimeError("Label mapping failed in Stage A") # Build multi-template positives for the symmetric InfoNCE objective. templates_cfg = GENERIC_TEMPLATES # Keep the task prompt plus additional generic templates without duplication. all_templates = [] base_t = str(getattr(cfg, 'prompt_template', self.prompt_template)) if base_t not in all_templates: all_templates.append(base_t) for t in templates_cfg: if t not in all_templates: all_templates.append(t) Tmpl = len(all_templates) # Restrict positive anchors to the classes present in the minibatch. present_local = torch.unique(targets).tolist() present_local.sort() present_names = [classnames[i] for i in present_local] texts_present = [tmpl.format(c) for c in present_names for tmpl in all_templates] other_seen = [n for n in dict.fromkeys(all_seen_names) if n not in set(present_names)] texts_neg = [tmpl.format(c) for c in other_seen for tmpl in all_templates] if other_seen else [] with torch.no_grad(): scale = self.model.logit_scale.exp() if train_text_e1_stage_a: txt_present_anchors = self.model.encode_text(clip.tokenize(texts_present).to(self.device)) txt_present_anchors = txt_present_anchors / txt_present_anchors.norm(dim=-1, keepdim=True) if texts_neg: txt_neg = self.model.encode_text(clip.tokenize(texts_neg).to(self.device)) txt_neg = txt_neg / txt_neg.norm(dim=-1, keepdim=True) else: txt_neg = None else: with torch.no_grad(): txt_present_anchors = self.model.encode_text(clip.tokenize(texts_present).to(self.device)) txt_present_anchors = txt_present_anchors / txt_present_anchors.norm(dim=-1, keepdim=True) if texts_neg: txt_neg = self.model.encode_text(clip.tokenize(texts_neg).to(self.device)) txt_neg = txt_neg / txt_neg.norm(dim=-1, keepdim=True) else: txt_neg = None # Re-encode with gradients so the alignment pathway receives updates. img0_grad = self.model.encode_image(inputs) feats = img0_grad feats = feats / feats.norm(dim=-1, keepdim=True) # Map task-local labels to the compact minibatch class set. present_index_map = {int(lid): idx for idx, lid in enumerate(present_local)} # Build per-sample paired text features for the task-level queue refresh. B = feats.size(0) Tmpl = max(1, Tmpl) txt_pos_mean = [] for i in range(B): cls_lid = int(targets[i].item()) pidx = present_index_map[cls_lid] start = pidx * Tmpl end = start + Tmpl txt_block = txt_present_anchors[start:end] txt_pos_mean.append(txt_block.mean(dim=0, keepdim=True)) txt_pos_mean = torch.cat(txt_pos_mean, dim=0) # [B, D] opt_a.zero_grad(set_to_none=True) # Candidate sets combine current positives, seen-class negatives, and valid queue entries. qcnt_a = int(getattr(self, 'queue_count', 0)) if use_moco and self.queue_size > 0 and qcnt_a > 0: k_valid_a = min(qcnt_a, int(self.queue_size)) queue_txt_batch = self.queue_txt[:, :k_valid_a].clone().detach().t().to(feats.device) # [K, D] cand_txt = torch.cat( [ txt_present_anchors, txt_neg if txt_neg is not None else txt_present_anchors.new_zeros((0, txt_present_anchors.size(1))), queue_txt_batch, ], dim=0, ) queue_img_batch = self.queue_img[:, :k_valid_a].clone().detach().t().to(feats.device) # [K, D] cand_img = torch.cat([feats, queue_img_batch], dim=0) # [B+K, D] else: cand_txt = torch.cat( [ txt_present_anchors, txt_neg if txt_neg is not None else txt_present_anchors.new_zeros((0, txt_present_anchors.size(1))), ], dim=0, ) cand_img = feats # Image-to-text InfoNCE with all templates of the matched class as positives. N_txt = cand_txt.size(0) pos_mask_i2t = torch.zeros(B, N_txt, dtype=torch.bool, device=feats.device) for i in range(B): cls_lid = int(targets[i].item()) pidx = present_index_map[cls_lid] start = pidx * Tmpl end = start + Tmpl pos_mask_i2t[i, start:end] = True logits_i2t = (scale * feats @ cand_txt.t()) / max(1e-6, tau_con) # Multi-positive InfoNCE in log-sum-exp form. pos_logits = logits_i2t.masked_fill(~pos_mask_i2t, float('-inf')) numer = torch.logsumexp(pos_logits, dim=1) denom = torch.logsumexp(logits_i2t, dim=1) loss_i2t = -(numer - denom).mean() # Text-to-image InfoNCE: each class/template anchor matches all in-batch images of that class. A = txt_present_anchors.size(0) # n_present * Tmpl N_img = cand_img.size(0) pos_mask_t2i = torch.zeros(A, N_img, dtype=torch.bool, device=feats.device) # Recover the class identity of each text anchor. for a in range(A): cls_local = present_local[a // Tmpl] # Positives are the in-batch image features of the same class. match_idx = (targets == int(cls_local)).nonzero(as_tuple=False).squeeze(1) if match_idx.numel() > 0: pos_mask_t2i[a, match_idx] = True logits_t2i = (scale * txt_present_anchors @ cand_img.t()) / max(1e-6, tau_con) pos_logits_t2i = logits_t2i.masked_fill(~pos_mask_t2i, float('-inf')) numer_t2i = torch.logsumexp(pos_logits_t2i, dim=1) denom_t2i = torch.logsumexp(logits_t2i, dim=1) # Ignore anchors whose class is absent from the current minibatch. valid_rows = pos_mask_t2i.any(dim=1) if valid_rows.any(): loss_t2i = -((numer_t2i[valid_rows] - denom_t2i[valid_rows]).mean()) else: loss_t2i = torch.zeros((), device=feats.device, dtype=feats.dtype) loss_a = (loss_i2t + loss_t2i) * 0.5 loss_a.backward() opt_a.step() # Defer queue writes until the task finishes; no online replay is stored. if use_moco and self.queue_size > 0: try: mq_img_buf.append(feats.detach().cpu()) mq_txt_buf.append(txt_pos_mean.detach().cpu()) except Exception: pass # Restore normal two-pathway routing after Stage A. for m in getattr(self, 'dfa_blocks', []): m.e1_only = False if train_text_e1_stage_a: for m in self.text_dfa_blocks: m.e1_only = False # Queue refresh is task-level, so the actual FIFO update happens after Stage B. # -------- Stage B: Plasticity Pathway -------- # Learning rates follow the paper's separation between experts and routers. lr_e2 = float(getattr(cfg, 'lr_e2', getattr(cfg, 'lr', 1e-3))) lr_e2_text = float(getattr(cfg, 'text_lr_e2', max(lr_e2 * 0.1, 1e-6))) lr_e2_router = float(getattr(cfg, 'lr_e2_router', lr_e2)) lr_top_router = float(getattr(cfg, 'lr_top_router', 5.0e-6)) weight_decay = float(getattr(cfg, 'weight_decay', 0.0)) # Collect trainable parameters for the task-specific experts and both routers. img_adapt_params, txt_adapt_params = [], [] e2_router_params = [] for m in self.dfa_blocks: if hasattr(m, 'e1'): for p in m.e1.parameters(): p.requires_grad = False for e2 in getattr(m, 'e2_list', []): for p in e2.parameters(): p.requires_grad = True img_adapt_params += list(e2.parameters()) for p in m.w_gate_e2.parameters(): p.requires_grad = True for p in m.w_noise_e2.parameters(): p.requires_grad = True e2_router_params += list(m.w_gate_e2.parameters()) + list(m.w_noise_e2.parameters()) if hasattr(self, 'text_dfa_blocks') and self.text_dfa_blocks: for m in self.text_dfa_blocks: if hasattr(m, 'e1'): for p in m.e1.parameters(): p.requires_grad = False for e2 in getattr(m, 'e2_list', []): for p in e2.parameters(): p.requires_grad = True txt_adapt_params += list(e2.parameters()) for p in m.w_gate_e2.parameters(): p.requires_grad = True for p in m.w_noise_e2.parameters(): p.requires_grad = True e2_router_params += list(m.w_gate_e2.parameters()) + list(m.w_noise_e2.parameters()) # The outer router balances Alignment vs. Plasticity contributions. top_router_params = [] if hasattr(self, 'dfa_blocks') and self.dfa_blocks: for m in self.dfa_blocks: for p in m.router_top.parameters(): p.requires_grad = True top_router_params += list(m.router_top.parameters()) if hasattr(self, 'text_dfa_blocks') and self.text_dfa_blocks: for m in self.text_dfa_blocks: for p in m.router_top.parameters(): p.requires_grad = True top_router_params += list(m.router_top.parameters()) param_groups = [] base_lrs = [] if img_adapt_params: param_groups.append({"params": img_adapt_params, "lr": lr_e2, "weight_decay": weight_decay}) base_lrs.append(lr_e2) if txt_adapt_params: param_groups.append({"params": txt_adapt_params, "lr": lr_e2_text, "weight_decay": weight_decay}) base_lrs.append(lr_e2_text) if e2_router_params: param_groups.append({"params": e2_router_params, "lr": lr_e2_router, "weight_decay": weight_decay}) base_lrs.append(lr_e2_router) if top_router_params: param_groups.append({"params": top_router_params, "lr": lr_top_router, "weight_decay": weight_decay}) base_lrs.append(lr_top_router) if not param_groups: raise RuntimeError("No parameters collected for Stage B optimization in block mode") opt_b = torch.optim.AdamW(param_groups) sched_b = utils.cosine_lr(opt_b, base_lrs if len(base_lrs) > 1 else base_lrs[0], 30, total_iterations_b) it = 0 for epoch in range(epochs_b): for inputs, targets_abs, _ in tqdm(train_loader, desc=f"Task {task_id} B (E2+Router)"): sched_b(it) it += 1 inputs = inputs.to(self.device, non_blocking=True) targets_abs = targets_abs.to(self.device, non_blocking=True) map_table = map_table_cpu.to(targets_abs.device) targets = map_table[targets_abs] if (targets < 0).any(): raise RuntimeError("Label mapping failed in Stage B") # If text-side DFA blocks exist, let the plasticity pathway update them as well. if hasattr(self, 'text_dfa_blocks') and self.text_dfa_blocks: txt = self.model.encode_text(texts_task) txt = txt / txt.norm(dim=-1, keepdim=True) with torch.no_grad(): scale = self.model.logit_scale.exp() else: with torch.no_grad(): txt = self.model.encode_text(texts_task) txt = txt / txt.norm(dim=-1, keepdim=True) scale = self.model.logit_scale.exp() opt_b.zero_grad(set_to_none=True) img = self.model.encode_image(inputs) img = img / img.norm(dim=-1, keepdim=True) fused = img fused = fused / fused.norm(dim=-1, keepdim=True) logits = scale * fused @ txt.t() loss_ce = F.cross_entropy(logits[:, :local_C], targets, label_smoothing=float(getattr(cfg, 'ls', 0.0))) loss = loss_ce # Auxiliary contrastive loss for the Plasticity Pathway. lambda_b_con = float(getattr(cfg, 'lambda_b_con', 0.0)) if lambda_b_con > 0.0 and local_C > 1 and hasattr(self, 'queue_img') and hasattr(self, 'queue_txt'): tau_b_con = float(getattr(cfg, 'tau_b_con', getattr(cfg, 'tau_con', 0.05))) qcnt_con = int(getattr(self, 'queue_count', 0)) B_con = fused.size(0) # Reuse the historical queue to recall prior task distributions. if qcnt_con > 0 and self.queue_size > 0: k_valid_con = min(qcnt_con, int(self.queue_size)) queue_txt_neg = self.queue_txt[:, :k_valid_con].t().to(fused.device) # [K, D] queue_txt_neg = queue_txt_neg / (queue_txt_neg.norm(dim=-1, keepdim=True) + 1e-12) cand_txt_b = torch.cat([txt, queue_txt_neg], dim=0) # [local_C + K, D] queue_img_neg = self.queue_img[:, :k_valid_con].t().to(fused.device) # [K, D] queue_img_neg = queue_img_neg / (queue_img_neg.norm(dim=-1, keepdim=True) + 1e-12) cand_img_b = torch.cat([fused, queue_img_neg], dim=0) # [B + K, D] else: cand_txt_b = txt cand_img_b = fused # Image-to-text InfoNCE with the class prompt as the positive target. N_txt_b = cand_txt_b.size(0) logits_i2t_b = (scale * fused @ cand_txt_b.t()) / max(1e-6, tau_b_con) # [B, N_txt] # Each image matches its task-local class prompt. pos_mask_i2t_b = torch.zeros(B_con, N_txt_b, dtype=torch.bool, device=fused.device) for i in range(B_con): pos_idx = int(targets[i].item()) if 0 <= pos_idx < local_C: pos_mask_i2t_b[i, pos_idx] = True pos_logits_i2t_b = logits_i2t_b.masked_fill(~pos_mask_i2t_b, float('-inf')) numer_i2t_b = torch.logsumexp(pos_logits_i2t_b, dim=1) denom_i2t_b = torch.logsumexp(logits_i2t_b, dim=1) loss_i2t_b = -(numer_i2t_b - denom_i2t_b).mean() # Each task-local class prompt matches all in-batch images of that class. N_img_b = cand_img_b.size(0) logits_t2i_b = (scale * txt @ cand_img_b.t()) / max(1e-6, tau_b_con) # [local_C, N_img] pos_mask_t2i_b = torch.zeros(local_C, N_img_b, dtype=torch.bool, device=fused.device) for c in range(local_C): match_idx = (targets == c).nonzero(as_tuple=False).squeeze(1) if match_idx.numel() > 0: pos_mask_t2i_b[c, match_idx] = True pos_logits_t2i_b = logits_t2i_b.masked_fill(~pos_mask_t2i_b, float('-inf')) numer_t2i_b = torch.logsumexp(pos_logits_t2i_b, dim=1) denom_t2i_b = torch.logsumexp(logits_t2i_b, dim=1) valid_rows_b = pos_mask_t2i_b.any(dim=1) if valid_rows_b.any(): loss_t2i_b = -((numer_t2i_b[valid_rows_b] - denom_t2i_b[valid_rows_b]).mean()) else: loss_t2i_b = torch.zeros((), device=fused.device, dtype=fused.dtype) loss_b_con = (loss_i2t_b + loss_t2i_b) * 0.5 loss = loss + lambda_b_con * loss_b_con loss.backward() opt_b.step() # After the task finishes, refresh the FIFO queue with buffered feature pairs. if use_moco and self.queue_size > 0: if 'mq_img_buf' in locals() and len(mq_img_buf) > 0: try: img_cat = torch.cat(mq_img_buf, dim=0) txt_cat = torch.cat(mq_txt_buf, dim=0) # Match the paper's task-level queue update with a bounded enqueue size. if int(task_id) >= 1: task_quota = int(getattr(cfg, 'moco_task_enqueue_quota', 32)) task_quota = max(0, min(task_quota, int(self.queue_size))) if task_quota > 0: if img_cat.size(0) > task_quota: sample_idx = torch.randperm(img_cat.size(0), device=img_cat.device)[:task_quota] img_cat = img_cat[sample_idx] txt_cat = txt_cat[sample_idx] if img_cat.numel() > 0: img_cat = img_cat.to(self.device, non_blocking=True) txt_cat = txt_cat.to(self.device, non_blocking=True) # Chunk the write so FIFO pointer updates remain bounded. chunk = min(1024, self.queue_size) for s in range(0, img_cat.size(0), chunk): self._dequeue_and_enqueue(img_cat[s:s+chunk], txt_cat[s:s+chunk]) except Exception: pass def compute_logits(self, inputs, text_tokens): """Run the unified DFA-MoE inference path against the provided text bank.""" m_prev = self.model.training b_prev = [m.training for m in getattr(self, 'dfa_blocks', [])] t_prev = [m.training for m in getattr(self, 'text_dfa_blocks', [])] if hasattr(self, 'text_dfa_blocks') and self.text_dfa_blocks else [] self.model.eval() for m in getattr(self, 'dfa_blocks', []): m.eval() if hasattr(self, 'text_dfa_blocks') and self.text_dfa_blocks: for m in self.text_dfa_blocks: m.eval() with torch.no_grad(): img = self.model.encode_image(inputs) img = img / img.norm(dim=-1, keepdim=True) # The wrapped CLIP blocks already apply alignment/plasticity routing internally. fused = img fused = fused / fused.norm(dim=-1, keepdim=True) txt = self.model.encode_text(text_tokens) txt = txt / txt.norm(dim=-1, keepdim=True) scale = self.model.logit_scale.exp() logits = scale * fused @ txt.t() if hasattr(self, 'text_dfa_blocks') and self.text_dfa_blocks: for m, s in zip(self.text_dfa_blocks, t_prev): m.train(s) for m, s in zip(getattr(self, 'dfa_blocks', []), b_prev): m.train(s) self.model.train(m_prev) return logits @torch.no_grad() def _dequeue_and_enqueue(self, img_feats: torch.Tensor, txt_feats: torch.Tensor): """Update the historical FIFO feature queue used by both training stages.""" if not hasattr(self, 'queue_count'): try: self.queue_count = 0 except Exception: self.queue_count = 0 batch_size = img_feats.size(0) ptr = int(self.queue_ptr) # Standard FIFO overwrite of the oldest slots. if ptr + batch_size <= self.queue_size: self.queue_img[:, ptr:ptr + batch_size] = img_feats.t() self.queue_txt[:, ptr:ptr + batch_size] = txt_feats.t() ptr = (ptr + batch_size) % self.queue_size else: # Wrap around when the enqueue spans the circular buffer boundary. remain = self.queue_size - ptr self.queue_img[:, ptr:] = img_feats[:remain].t() self.queue_txt[:, ptr:] = txt_feats[:remain].t() overflow = batch_size - remain if overflow > 0: self.queue_img[:, :overflow] = img_feats[remain:].t() self.queue_txt[:, :overflow] = txt_feats[remain:].t() ptr = overflow self.queue_ptr[0] = ptr # Track the valid prefix used by Stage A / Stage B contrastive losses. try: self.queue_count = min(self.queue_size, int(self.queue_count) + batch_size) except Exception: self.queue_count = self.queue_size def _inject_dfa_into_visual_blocks(self, cfg): """Inject per-block DFA-MoE adapters into the visual transformer.""" visual = getattr(self.model, 'visual', None) transformer = getattr(visual, 'transformer', None) resblocks = getattr(transformer, 'resblocks', None) if visual is None or transformer is None or resblocks is None: return # Each visual transformer block receives its own alignment/plasticity module. new_blocks = [] self.dfa_blocks = [] top_k = int(getattr(cfg, 'e2_top_k', 2)) for b in resblocks: # Infer the hidden width from the original transformer block. width = None try: width = int(getattr(b, 'ln_1').weight.shape[0]) except Exception: width = int(getattr(visual, 'width', 768)) dfab = DFA3Block( width, e2_top_k=top_k, num_task_experts=int(getattr(cfg, 'num_task_experts', 2)), ).to(self.device) wrap = ResidualAttentionBlockWithDFA(b, dfab) new_blocks.append(wrap) self.dfa_blocks.append(dfab) # CLIP expects `resblocks(x)`, so preserve that contract with an `nn.Sequential` wrapper. transformer.resblocks = nn.Sequential(*new_blocks) def _inject_dfa_into_text_blocks(self, cfg): """Inject the same DFA-MoE structure into the text transformer.""" text_tf = getattr(self.model, 'transformer', None) resblocks = getattr(text_tf, 'resblocks', None) if text_tf is None or resblocks is None: return # Text-side routing follows the same top-k setting unless explicitly overridden. top_k = int(getattr(cfg, 'text_e2_top_k', getattr(cfg, 'e2_top_k', 2))) # Mirror the per-block DFA adapter design used on the visual side. new_blocks = [] self.text_dfa_blocks = [] blocks = list(resblocks) for b in blocks: # Infer the hidden width from the original transformer block. try: width = int(getattr(b, 'ln_1').weight.shape[0]) except Exception: # Fall back to the transformer-wide width when block metadata is missing. width = int(getattr(text_tf, 'width', 512)) dfab = DFA3Block( width, e2_top_k=top_k, num_task_experts=int(getattr(cfg, 'num_task_experts', 2)), ).to(self.device) wrap = ResidualAttentionBlockWithDFA(b, dfab) new_blocks.append(wrap) self.text_dfa_blocks.append(dfab) text_tf.resblocks = nn.Sequential(*new_blocks) class DomainIncremental(nn.Module): pass class TaskAgnostic(nn.Module): pass def load_model(cfg: DictConfig, device: torch.device) -> nn.Module: r"""Instantiate the continual-learning model for the requested scenario.""" if cfg.scenario == "class": return ClassIncremental(cfg, device) elif cfg.scenario == "domain": return DomainIncremental(cfg, device) elif cfg.scenario == "task-aganostic": return TaskAgnostic(cfg, device) else: raise ValueError(f""" `{cfg.scenarios}` is not a valid scenario, Please choose from ['class', "domain', 'task-agnostic'] """)