cida-plugin 1.0.0__py3-none-any.whl

cida_plugin/deliberation.py

@@ -0,0 +1,376 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from .agent import AgentState
+
+
+class MessageFormulator(nn.Module):
+    """
+    Creates a structured message out of the explicit agent components.
+
+    v4: Anonymous mode (Choi et al. ACL 2025).
+    Убирает s_i из сообщения → агент слышит аргумент, не авторитет.
+    Доказано устраняет 100% identity bias в debate.
+
+    Было:  m = φ([s; e; b; u]) — агент знает источник → sycophancy/self-bias
+    Стало: m = φ([e; b; u]) — role prior только в инициализации b₀
+    """
+    def __init__(self, d_hidden: int, num_classes: int, d_message: int, anonymous: bool = True):
+        super().__init__()
+        self.anonymous = anonymous
+        if anonymous:
+            in_dim = d_hidden + num_classes + 1  # [e; b; u] — без s
+        else:
+            in_dim = d_hidden + d_hidden + num_classes + 1  # [s; e; b; u]
+        self.proj = nn.Sequential(
+            nn.Linear(in_dim, d_message),
+            nn.LayerNorm(d_message),
+            nn.SiLU(),
+            nn.Linear(d_message, d_message)
+        )
+
+    def forward(self, s, e, b, u):
+        if self.anonymous:
+            concat = torch.cat([e, b, u], dim=-1)
+        else:
+            concat = torch.cat([s, e, b, u], dim=-1)
+        return self.proj(concat)
+
+
+class AgentEvidenceExtractor(nn.Module):
+    """
+    Computes evidence pointers over tokens, giving each agent a distinct focus.
+    p_i^t = softmax((H W_p)(W_q s_i^t)^T / sqrt(d))
+    e_i^t = sum_k p_{i,k}^t H_k
+    """
+    def __init__(self, d_hidden: int):
+        super().__init__()
+        self.W_q = nn.Linear(d_hidden, d_hidden)
+        self.W_p = nn.Linear(d_hidden, d_hidden)
+        self.scale = d_hidden ** -0.5
+
+    def forward(self, s, H, mask=None):
+        K = self.W_p(H)
+        Q = self.W_q(s)
+        V = H
+
+        attn_mask = None
+        if mask is not None:
+            # mask: (B, L) -> (B, 1, L) для бродкаста на M агентов
+            attn_mask = mask.bool().unsqueeze(1)
+
+        # F.scaled_dot_product_attention автоматически использует FlashAttention на H100
+        # Масштабирование (scale) также применяется автоматически внутри.
+        e = F.scaled_dot_product_attention(
+            Q, K, V,
+            attn_mask=attn_mask,
+            dropout_p=0.0,
+            is_causal=False,
+        )
+
+        # Возвращаем пустой тензор для p, так как FlashAttention не материализует матрицу внимания.
+        # Это экономит O(B*M*L) памяти.
+        p_dummy = torch.empty((Q.size(0), Q.size(1), K.size(1)), device=Q.device)
+        return p_dummy, e
+
+
+class CounterargumentCommunication(nn.Module):
+    """
+    Agents collect messages from others, weighted by an attention mechanism
+    that explicitly values disagreement while punishing simple similarity.
+
+    v4: Disagreement-Routed Sparse Communication (Li et al. EMNLP 2024).
+    Каждый агент слушает только top-K по несогласию — противоположные
+    убеждения дают сильную связь, похожие — разрыв.
+    Это защищает minority voice и сокращает вычисления на 75%.
+    """
+    def __init__(self, d_hidden: int, d_message: int, lambda_d: float = 1.0, lambda_s: float = 1.0):
+        super().__init__()
+        self.W_q = nn.Linear(d_hidden, d_hidden)
+        self.W_k = nn.Linear(d_message, d_hidden)
+        self.W_v = nn.Linear(d_message, d_hidden)
+        self.lambda_d = lambda_d
+        self.lambda_s = lambda_s
+        self.scale = d_hidden ** -0.5
+
+    def forward(self, s, m, b, e, sparse_k: int = 0):
+        """
+        s: (B, M, d), m: (B, M, d_msg), b: (B, M, K), e: (B, M, d)
+        sparse_k: если > 0 — включает sparse communication (top-K по disagreement)
+        """
+        Q = self.W_q(s)
+        K = self.W_k(m)
+        V = self.W_v(m)
+
+        attn = torch.bmm(Q, K.transpose(1, 2)) * self.scale
+
+        b_i = b.unsqueeze(2)
+        b_j = b.unsqueeze(1)
+        D = (b_i - b_j).abs().sum(dim=-1)
+
+        e_norm = F.normalize(e, p=2, dim=-1)
+        S = torch.bmm(e_norm, e_norm.transpose(1, 2))
+
+        scores = attn + self.lambda_d * D - self.lambda_s * S
+
+        M_agents = s.size(1)
+        # Само-маскировка (агент не слушает себя)
+        self_mask = torch.eye(M_agents, device=s.device, dtype=torch.bool).unsqueeze(0)
+        scores = scores.masked_fill(self_mask, float('-inf'))
+
+        # v4: Sparse communication — top-K по disagreement
+        if sparse_k > 0 and sparse_k < M_agents - 1:
+            # Находим top-K наиболее несогласных агентов для каждого
+            b_diff = D.clone()  # (B, M, M)
+            b_diff = b_diff.masked_fill(self_mask, 0.0)
+            _, top_idx = b_diff.topk(min(sparse_k, M_agents - 1), dim=-1)  # (B, M, k)
+            sparse_mask = torch.zeros(scores.shape, device=s.device, dtype=torch.bool)
+            sparse_mask.scatter_(-1, top_idx, True)
+            # Маскируем всех кроме top-K
+            scores = scores.masked_fill(~sparse_mask, float('-inf'))
+
+        a = F.softmax(scores, dim=-1)
+        r = torch.bmm(a, V)
+        return r
+
+
+# ─── v2: Transformer Agent Updater ───────────────────────────────────────────
+
+class AgentUpdater(nn.Module):
+    """
+    [v2] Обновляет состояние агентов через Transformer Decoder Step.
+
+    Замена GRUCell на cross-attention даёт агентам возможность
+    самостоятельно решать, НА ЧТО обращать внимание в сообщениях других,
+    а не слепо конкатенировать всё подряд.
+
+    Шаги:
+        1. Cross-Attention: каждый агент (query) смотрит на сигналы
+           counterargument (key/value) от других агентов.
+        2. Gated Evidence Fusion: ворота решают, сколько evidence добавить.
+        3. FFN + Residual + LayerNorm.
+        4. Dirichlet Belief Update: alpha = softplus(g(s)) + 1.
+    """
+
+    def __init__(self, d_hidden: int, num_classes: int, num_heads: int = 4, multi_label: bool = False):
+        super().__init__()
+        self.multi_label = multi_label
+        # ── Step 1: Cross-attention over counterarguments ─────────────────────
+        self.cross_attn = nn.MultiheadAttention(
+            embed_dim=d_hidden,
+            num_heads=num_heads,
+            batch_first=True,
+            dropout=0.1,
+        )
+        self.norm1 = nn.LayerNorm(d_hidden)
+
+        # ── Step 2: Gated evidence fusion ─────────────────────────────────────
+        self.evidence_gate = nn.Sequential(
+            nn.Linear(d_hidden * 2, d_hidden),
+            nn.Sigmoid(),
+        )
+        self.evidence_proj = nn.Linear(d_hidden, d_hidden)
+        self.norm2 = nn.LayerNorm(d_hidden)
+
+        # ── Step 3: FFN ───────────────────────────────────────────────────────
+        self.ffn = nn.Sequential(
+            nn.Linear(d_hidden, d_hidden * 2),
+            nn.SiLU(),
+            nn.Linear(d_hidden * 2, d_hidden),
+        )
+        self.ffn_norm = nn.LayerNorm(d_hidden) # Replaced norm3 for clarity
+
+        # ── Step 4: Belief generator ──────────────────────────────────────────
+        self.g = nn.Sequential(
+            nn.Linear(d_hidden, d_hidden),
+            nn.SiLU(),
+            nn.Linear(d_hidden, num_classes),
+        )
+        self.K = float(num_classes)
+
+    def forward(self, state: AgentState, r: torch.Tensor, e: torch.Tensor) -> AgentState:
+        """
+        state : AgentState — текущее состояние агентов
+        r     : (B, M, d_hidden) — aggregated counterargument signals
+        e     : (B, M, d_hidden) — evidence vectors из AgentEvidenceExtractor
+        """
+        s = state.s  # (B, M, d_hidden)
+
+        # ── Step 1: Cross-attention ───────────────────────────────────────────
+        # Query: текущие состояния агентов
+        # Key/Value: counterargument signals — что говорят другие агенты
+        # NOTE: если r все нули (начальный шаг), cross-attention = identity
+        attn_out, _ = self.cross_attn(query=s, key=r, value=r)
+        s = self.norm1(s + attn_out)
+
+        # ── Step 2: Gated evidence fusion ─────────────────────────────────────
+        # Ворота: насколько новый evidence должен изменить состояние?
+        gate = self.evidence_gate(torch.cat([s, e], dim=-1))   # (B, M, d)
+        e_proj = self.evidence_proj(e)                          # (B, M, d)
+        s = self.norm2(s + gate * e_proj)
+
+        # ── Step 3: FFN + residual ────────────────────────────────────────────
+        s = self.ffn_norm(s + self.ffn(s))
+
+        # ── Step 4: Belief update ─────────────────────────────────────────────
+        g_out = self.g(s)
+        
+        if self.multi_label:
+            # Multi-label: Independent sigmoids
+            b_next = torch.sigmoid(g_out) # (B, M, K)
+            # v4: Разделение неопределённости
+            u_alea = (4.0 * b_next * (1.0 - b_next)).mean(dim=-1, keepdim=True)  # шум данных
+            u_epi = u_alea * 0.5  # приближение для multi-label
+            u_next = u_alea + u_epi
+            alpha = b_next # Dummy for AgentState
+        else:
+            # Single-label: Dirichlet
+            alpha = F.softplus(g_out) + 1.0      # (B, M, K), всегда > 1
+            alpha_sum = alpha.sum(dim=-1, keepdim=True)
+            b_next = alpha / alpha_sum
+            # v4: Разделение неопределённости (Josang 2002 + Meinert 2024)
+            u_epi = self.K ** 2 / alpha_sum ** 2     # Vacuity: неполнота знаний
+            u_alea = 1.0 - b_next.max(dim=-1, keepdim=True).values  # Шум данных
+            u_next = u_epi + u_alea               # Совместимость с u (total)
+
+        return AgentState(
+            s=s,
+            b=b_next,
+            u=u_next,
+            alpha=alpha,
+            p=state.p,
+            e=e,
+            u_epi=u_epi,
+            u_alea=u_alea,
+        )
+
+
+# ─── v3: Structural Specialization ───────────────────────────────────────────
+
+class PerspectiveProjector(nn.Module):
+    """
+    [v3+] Каждый агент имеет свой собственный "взгляд" на входные данные.
+    Вместо клонирования h_cls, каждый агент i проецирует pooled_output
+    через свою собственную матрицу.
+
+    [v4] Добавлен orthogonality_loss(): штраф за коллинеарность матриц
+    проекций. Заставляет агентов смотреть на РАЗНЫЕ аспекты входа.
+    Без этого штрафа все M проекций коллапсируют к одной — агенты
+    становятся копиями друг друга и дебаты теряют смысл.
+    """
+
+    def __init__(self, num_agents: int, d_input: int, d_hidden: int):
+        super().__init__()
+        self.num_agents = num_agents
+        # Каждый агент имеет свою проекцию
+        self.agent_projections = nn.ModuleList([
+            nn.Linear(d_input, d_hidden) for _ in range(num_agents)
+        ])
+        # Опционально: позиционное смещение для севенциальных данных
+        # Предполагаем макс. длину 512 для инициализации
+        self.position_bias = nn.Parameter(
+            torch.randn(num_agents, 512) * 0.01
+        )
+
+    def forward(self, pooled: torch.Tensor) -> torch.Tensor:
+        """
+        pooled: (B, d_input)
+        Returns: (B, M, d_hidden)
+        """
+        out = [proj(pooled) for proj in self.agent_projections]
+        return torch.stack(out, dim=1)
+
+    def orthogonality_loss(self) -> torch.Tensor:
+        """
+        [v4] Штраф за коллинеарность проекций агентов.
+
+        Gram(W) = W_norm @ W_norm^T, где W_norm — нормализованные строки
+        матриц весов каждого агента. Идеал: Gram = I (единичная матрица).
+        Loss = mean((Gram - I)^2) → 0 при полной ортогональности.
+
+        Математическое обоснование:
+            cos(θ_ij) = <w_i, w_j> / (||w_i|| ||w_j||)
+            При L_orth → 0: cos(θ_ij) → 0 для i≠j ⟹ θ_ij → 90°
+            Агенты гарантированно смотрят на линейно независимые подпространства.
+        """
+        # Собираем матрицы весов: (M, d_hidden, d_input)
+        W = torch.stack([proj.weight for proj in self.agent_projections])
+        M = W.size(0)
+
+        # Flatten: (M, d_hidden * d_input)
+        W_flat = W.view(M, -1)
+
+        # L2-нормализация строк
+        W_norm = F.normalize(W_flat, p=2, dim=-1)
+
+        # Gram matrix: (M, M) — косинусные сходства между агентами
+        gram = W_norm @ W_norm.T
+
+        # Цель: единичная матрица (каждый агент ортогонален остальным)
+        eye = torch.eye(M, device=gram.device)
+
+        return ((gram - eye) ** 2).mean()
+
+
+class AdaptiveDSchedule(nn.Module):
+    """
+    [v3+] Адаптивное расписание разногласий с PI-регулятором.
+
+    Исправление: current_disagreement теперь ИСПОЛЬЗУЕТСЯ для обратной связи.
+    Прежняя версия принимала параметр, но полностью его игнорировала.
+
+    Математическое обоснование:
+        d_target = base(h_cls) + Kp * e_t + Ki * Σe_k
+        где e_t = base - actual_disagreement — ошибка регулирования.
+
+    PI-регулятор гарантирует экспоненциальное стремление фактического
+    разногласия к целевому (доказано через анализ устойчивости замкнутой системы):
+        |λ| = |1 - η(1 - Kp)| < 1  при 0 < η < 2, 0 < Kp < 1.
+
+    Возвращает (d_target, error) — ошибку нужно накапливать в integral_acc
+    в вызывающем коде (core.py), а не внутри буфера модуля,
+    чтобы интегральный член был локальным для каждого forward-прохода.
+    """
+
+    def __init__(self, d_hidden: int, Kp: float = 0.5, Ki: float = 0.1):
+        super().__init__()
+        self.complexity_estimator = nn.Sequential(
+            nn.Linear(d_hidden, 64),
+            nn.SiLU(),
+            nn.Linear(64, 1),
+            nn.Sigmoid(),
+        )
+        self.Kp = Kp   # Пропорциональный коэффициент
+        self.Ki = Ki   # Интегральный коэффициент
+
+    def forward(
+        self,
+        h_cls: torch.Tensor,
+        current_disagreement: torch.Tensor,
+        integral_acc: torch.Tensor = None,
+    ):
+        """
+        h_cls               : (B, d_hidden) — начальное представление примера
+        current_disagreement: (B, 1, 1) или (B, 1) — фактическое разногласие
+        integral_acc        : (B, 1) — накопленная ошибка из предыдущих раундов
+
+        Возвращает:
+            d_target : (B, 1) — скорректированная цель разногласия ∈ [0, 1.5]
+            error    : (B, 1) — ошибка e_t для обновления integral_acc снаружи
+        """
+        complexity = self.complexity_estimator(h_cls)             # (B, 1)
+        base_target = complexity * 0.8 + (1 - complexity) * 0.1  # ∈ [0.1, 0.9]
+
+        # Приводим фактическое разногласие к форме (B, 1)
+        disag = current_disagreement.view(h_cls.size(0), 1)
+
+        # Пропорциональная ошибка (stop_grad — чистый управляющий сигнал)
+        error = (base_target - disag).detach()
+
+        # Интегральный член (накоплен вызывающим кодом за предыдущие раунды)
+        integral = integral_acc if integral_acc is not None else torch.zeros_like(error)
+
+        # PI-цель: base + proportional + integral
+        d_target = base_target + self.Kp * error + self.Ki * integral
+
+        return torch.clamp(d_target, 0.0, 1.5), error

cida_plugin/diagnostics.py

@@ -0,0 +1,124 @@
+import torch
+
+class DebateDiagnostics:
+    """
+    Validation mechanics to ensure explicit multi-agent deliberation avoids simple
+    over-confidence and actually provides semantic divergence over standard training.
+    """
+    
+    @staticmethod
+    def expected_calibration_error(p_final: torch.Tensor, y: torch.Tensor, n_bins: int = 10):
+        """Computes Expected Calibration Error (ECE) for final probability vector."""
+        if y.dim() == 2: # Multi-label
+            # If p_final contains logits (values > 1 or < 0 possible), apply sigmoid
+            if p_final.max() > 1.0 or p_final.min() < 0.0:
+                probs = torch.sigmoid(p_final)
+            else:
+                probs = p_final
+            
+            K = probs.size(1)
+            total_ece = 0.0
+            for k in range(K):
+                total_ece += DebateDiagnostics._binary_ece(probs[:, k], y[:, k], n_bins)
+            return total_ece / K
+            
+        # Single-label (original logic)
+        confidences, predictions = torch.max(p_final, dim=1)
+        accuracies = (predictions == y)
+        return DebateDiagnostics._calculate_ece_from_bins(confidences, accuracies, n_bins, p_final.device)
+
+    @staticmethod
+    def _binary_ece(probs: torch.Tensor, targets: torch.Tensor, n_bins: int):
+        """Binary ECE for a single class."""
+        accuracies = (probs > 0.5) == targets
+        return DebateDiagnostics._calculate_ece_from_bins(probs, accuracies, n_bins, probs.device)
+
+    @staticmethod
+    def _calculate_ece_from_bins(confidences: torch.Tensor, accuracies: torch.Tensor, n_bins: int, device: torch.device):
+        ece = torch.zeros(1, device=device)
+        bin_boundaries = torch.linspace(0, 1, n_bins + 1, device=device)
+        
+        for bin_idx in range(n_bins):
+            in_bin = (confidences > bin_boundaries[bin_idx]) & (confidences <= bin_boundaries[bin_idx + 1])
+            prop_in_bin = in_bin.float().mean()
+            
+            if prop_in_bin > 0:
+                accuracy_in_bin = accuracies[in_bin].float().mean()
+                avg_confidence_in_bin = confidences[in_bin].mean()
+                ece += torch.abs(avg_confidence_in_bin - accuracy_in_bin) * prop_in_bin
+                
+        return ece.item()
+        
+    @staticmethod
+    def persuasion_matrix(b_all: list):
+        """
+        Calculates how much agents changed their beliefs over time.
+        Allows us to track if persuasion is actually occurring.
+        Returns tensor of shape (Rounds, Agents) tracking average shifts.
+        """
+        T = len(b_all)
+        if T < 2:
+            return None
+            
+        shifts = []
+        for t in range(1, T):
+            shift = torch.norm(b_all[t] - b_all[t-1], p=1, dim=-1).mean(dim=0) # (M,)
+            shifts.append(shift)
+            
+        return torch.stack(shifts)
+
+    @staticmethod
+    def belief_diversity_curve(b_all: list):
+        """
+        Measures disagreement level over time. Useful to plot against d_t schedule.
+        """
+        divergence = []
+        for b_t in b_all:
+            b_i = b_t.unsqueeze(2)
+            b_j = b_t.unsqueeze(1)
+            dist = torch.norm(b_i - b_j, p=1, dim=-1).mean()
+            divergence.append(dist.item())
+        return divergence
+
+    # ─── [v5] TTT Diagnostics ────────────────────────────────────────────────
+
+    @staticmethod
+    def ttt_adaptation_magnitude(ttt_info: dict) -> float:
+        """
+        Насколько сильно TTT изменил веса агентов (L2 norm of param shift).
+        
+        Полезно для мониторинга: слишком маленький сдвиг = TTT бесполезен,
+        слишком большой = риск catastrophic forgetting.
+        Хороший диапазон: 0.01 — 1.0.
+        """
+        if ttt_info is None:
+            return 0.0
+        return ttt_info.get("param_shift", 0.0)
+
+    @staticmethod
+    def ttt_loss_curve(ttt_info: dict) -> list:
+        """
+        Кривая self-supervised loss по шагам TTT.
+        
+        Должна убывать: loss[0] > loss[1] > ... > loss[K-1].
+        Если не убывает — lr слишком большой или K слишком маленький.
+        """
+        if ttt_info is None:
+            return []
+        return ttt_info.get("losses", [])
+
+    @staticmethod
+    def ttt_convergence_ratio(ttt_info: dict) -> float:
+        """
+        Отношение финального loss к начальному: loss[-1] / loss[0].
+        
+        < 0.5 = хорошая адаптация (loss уменьшился вдвое+).
+        ≈ 1.0 = TTT не помогает, стоит увеличить K или lr.
+        > 1.0 = расходимость, уменьшить lr.
+        """
+        if ttt_info is None:
+            return 1.0
+        losses = ttt_info.get("losses", [])
+        if len(losses) < 2:
+            return 1.0
+        return losses[-1] / max(losses[0], 1e-9)