File size: 6,267 Bytes
79e8e52 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 | """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."""
# Learnable possibility mapping
init_temp: float = 1.0
init_bias: float = 0.0
temp_min: float = 0.1
temp_max: float = 10.0
# DPO
dpo_beta: float = 0.1
dpo_loss_weight: float = 0.5
# PPO-like penalty
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
# Parametrização: T = temp_min + (temp_max - temp_min) * sigmoid(raw_T)
# Inicializa raw_T tal que T(0) = init_temp
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 probs do winner e loser
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:
# Toy mode: assume ref = uniform
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)
# DPO: -log σ(β * (log p_w/p_ref_w - log p_l/p_ref_l))
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"]
|