| """rlhf_mapping.py — V6: learnable possibility mapping + RLHF. |
| |
| Implementa (item do PowerMachine/gru-ring-v13-9-2): |
| - **Learnable possibility mapping**: mapeia logits do modelo para uma |
| distribuição de "possibilidades" aprendível. Substitui o softmax fixo |
| por uma transformação aprendível que pode capturar distribuições não |
| uniformes (heavy-tailed, multimodais). |
| - **RLHF integration**: integra feedback humano (preference pairs) via |
| DPO-style loss. Diferente do DPO clássico, mantém a possibility mapping |
| aprendível durante o alignment. |
| |
| Matemática |
| ────────── |
| Possibility Mapping: |
| p_poss(y|x) = softmax(T_poss * logits + b_poss) |
| onde T_poss (temperatura aprendível) e b_poss (bias aprendível) são |
| parâmetros treináveis. T_poss permite que o modelo aprenda a "calibrar" |
| sua própria confiança. |
| |
| RLHF (DPO-style): |
| L_DPO = -log σ(β * (log p_θ(y_w|x) / p_ref(y_w|x) |
| - log p_θ(y_l|x) / p_ref(y_l|x))) |
| onde y_w = winner, y_l = loser, β = temperatura DPO. |
| """ |
| from __future__ import annotations |
| import math |
| from typing import Optional, Tuple, Dict, List |
| from dataclasses import dataclass |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
|
|
| @dataclass |
| class RLHFConfig: |
| """Configuração RLHF V6.""" |
| |
| init_temp: float = 1.0 |
| init_bias: float = 0.0 |
| temp_min: float = 0.1 |
| temp_max: float = 10.0 |
| |
| dpo_beta: float = 0.1 |
| dpo_loss_weight: float = 0.5 |
| |
| kl_penalty_weight: float = 0.01 |
|
|
|
|
| class LearnablePossibilityMapping(nn.Module): |
| """V6: mapeamento aprendível de logits para possibilidades. |
| |
| p_poss(y|x) = softmax(T_poss * logits + b_poss) |
| |
| T_poss é aprendível mas clipado em [temp_min, temp_max] via sigmoid. |
| """ |
| def __init__( |
| self, |
| init_temp: float = 1.0, |
| init_bias: float = 0.0, |
| temp_min: float = 0.1, |
| temp_max: float = 10.0, |
| ): |
| super().__init__() |
| self.temp_min = temp_min |
| self.temp_max = temp_max |
| |
| |
| init_raw = math.log( |
| (init_temp - temp_min) / max(temp_max - init_temp, 1e-8) |
| ) if temp_min < init_temp < temp_max else 0.0 |
| self.raw_temp = nn.Parameter(torch.tensor(float(init_raw))) |
| self.bias = nn.Parameter(torch.tensor(float(init_bias))) |
|
|
| @property |
| def temperature(self) -> torch.Tensor: |
| return self.temp_min + (self.temp_max - self.temp_min) * torch.sigmoid(self.raw_temp) |
|
|
| def forward(self, logits: torch.Tensor) -> torch.Tensor: |
| """logits: (..., V) → possibilities: (..., V) — distribuição suavizada.""" |
| T = self.temperature |
| return F.softmax(T * logits + self.bias, dim=-1) |
|
|
|
|
| class RLHFMapping(nn.Module): |
| """V6: RLHF com learnable possibility mapping. |
| |
| Forward (preference): |
| logits_w: (B, V) — logits do winner |
| logits_l: (B, V) — logits do loser |
| → dpo_loss: scalar, metrics: dict |
| |
| Forward (standard): |
| logits: (B, V) → possibilities: (B, V) |
| """ |
| def __init__(self, config: Optional[RLHFConfig] = None): |
| super().__init__() |
| self.cfg = config or RLHFConfig() |
| self.possibility = LearnablePossibilityMapping( |
| init_temp=self.cfg.init_temp, |
| init_bias=self.cfg.init_bias, |
| temp_min=self.cfg.temp_min, |
| temp_max=self.cfg.temp_max, |
| ) |
|
|
| def forward(self, logits: torch.Tensor) -> torch.Tensor: |
| """Standard forward: logits → possibilities.""" |
| return self.possibility(logits) |
|
|
| def dpo_loss( |
| self, |
| logits_w: torch.Tensor, |
| logits_l: torch.Tensor, |
| targets_w: torch.Tensor, |
| targets_l: torch.Tensor, |
| ref_logits_w: Optional[torch.Tensor] = None, |
| ref_logits_l: Optional[torch.Tensor] = None, |
| ) -> Tuple[torch.Tensor, Dict[str, float]]: |
| """Computa DPO loss. |
| |
| Se ref_logits não fornecido, usa logits_w/l como referência (toy mode). |
| """ |
| |
| log_p_w = F.log_softmax(logits_w, dim=-1) |
| log_p_l = F.log_softmax(logits_l, dim=-1) |
| log_p_w_target = log_p_w.gather(-1, targets_w.unsqueeze(-1)).squeeze(-1).sum(-1) |
| log_p_l_target = log_p_l.gather(-1, targets_l.unsqueeze(-1)).squeeze(-1).sum(-1) |
|
|
| if ref_logits_w is None or ref_logits_l is None: |
| |
| ref_log_p_w = torch.full_like(log_p_w_target, -math.log(logits_w.size(-1))) |
| ref_log_p_l = torch.full_like(log_p_l_target, -math.log(logits_l.size(-1))) |
| else: |
| ref_log_p_w = F.log_softmax(ref_logits_w, dim=-1).gather( |
| -1, targets_w.unsqueeze(-1) |
| ).squeeze(-1).sum(-1) |
| ref_log_p_l = F.log_softmax(ref_logits_l, dim=-1).gather( |
| -1, targets_l.unsqueeze(-1) |
| ).squeeze(-1).sum(-1) |
|
|
| |
| delta = self.cfg.dpo_beta * ( |
| (log_p_w_target - ref_log_p_w) - (log_p_l_target - ref_log_p_l) |
| ) |
| loss = -F.logsigmoid(delta).mean() |
|
|
| metrics = { |
| "dpo_loss": float(loss.item()), |
| "dpo_delta": float(delta.mean().item()), |
| "temperature": float(self.possibility.temperature.item()), |
| "bias": float(self.possibility.bias.item()), |
| "accuracy": float((delta > 0).float().mean().item()), |
| } |
| return self.cfg.dpo_loss_weight * loss, metrics |
|
|
| def kl_penalty( |
| self, |
| current_logits: torch.Tensor, |
| ref_logits: torch.Tensor, |
| ) -> torch.Tensor: |
| """KL(p_current || p_ref) — penalty PPO-style.""" |
| p_cur = F.softmax(current_logits, dim=-1) |
| log_p_cur = F.log_softmax(current_logits, dim=-1) |
| log_p_ref = F.log_softmax(ref_logits, dim=-1) |
| kl = (p_cur * (log_p_cur - log_p_ref)).sum(dim=-1).mean() |
| return self.cfg.kl_penalty_weight * kl |
|
|
|
|
| __all__ = ["RLHFMapping", "LearnablePossibilityMapping", "RLHFConfig"] |
|
|
