models/fixed_point_reasoning/model_utils.py

import math
import os
import torch
import torch.nn as nn
from typing import Dict

from ..common import trunc_normal_init_


def shift(input:torch.Tensor, shift, dim=0, fillval=0):
    # torch.roll without the copy of the wrap-around section
    size = input.size(dim)
    fill = torch.full_like(input.narrow(dim, 0, abs(shift)), fillval)
    if shift > 0:
        output = torch.cat([fill, input.narrow(dim, 0, size-shift)], dim=dim)
    if shift < 0:
        output = torch.cat([input.narrow(dim, -shift, size+shift), fill], dim=dim)
    return output

class WindowMix(nn.Module):
    """Learned convex mix of the last m iterate STATES, used as the damping
    anchor inside the fixed-point step. Weights are content-based (query =
    current iterate, keys = window states) plus a learnable per-age recency
    bias, passed through softmax -> convex (sum=1). Values are the RAW states
    (no v/o projection), so the read stays in the convex hull of the window;
    at convergence every slot == y* so the mix returns y* -> the fixed point is
    preserved exactly. recency_init large + q zero-init => starts one-hot on the
    newest state == byte-identical baseline damping. This is learned multi-point
    Krasnoselskii-Mann damping (a generalization of FPRM's fixed 2-point blend),
    NOT Anderson acceleration (which uses non-convex affine weights)."""

    def __init__(self, hidden: int, window: int, recency_init: float = 8.0):
        super().__init__()
        self.q = nn.Linear(hidden, hidden, bias=False)
        self.k = nn.Linear(hidden, hidden, bias=False)
        nn.init.zeros_(self.q.weight)   # content logit == 0 at init -> pure recency
        self.scale = hidden ** -0.5
        self.recency = nn.Parameter(-recency_init * torch.arange(window, dtype=torch.float32))

    def forward(self, window: torch.Tensor) -> torch.Tensor:   # window: [m, B, S, H], index 0 = newest
        m = window.shape[0]
        cur = window[0]
        q = self.q(cur).unsqueeze(0)                                  # [1, B, S, H]
        k = self.k(window)                                            # [m, B, S, H]
        logit = (q * k).sum(-1) * self.scale                          # [m, B, S]
        logit = logit + self.recency.to(logit.dtype).view(m, 1, 1)    # + learned age bias
        w = torch.softmax(logit, dim=0)                               # [m, B, S] convex
        return (w.unsqueeze(-1) * window).sum(dim=0).to(window.dtype)  # [B, S, H]


class FixedPointOptimizer(nn.Module):
    def __init__(self, config: dict):
        super().__init__()

        self.stepsize = config.stepsize
        self.stepsize_decay = config.stepsize_decay
        self.decay_patience = config.decay_patience
        self.eps = config.eps
        self.outlier_quantile = config.outlier_quantile
        self.max_iter = config.max_iter
        self.init_std = config.init_std
        self.additive_noise_std = config.additive_noise_std
        self.fp_thresh = config.fp_thresh

        # Multi-point convex damping on the loop axis. Window size from config
        # (learned mixer, for training) or env (zero-training probe). m<=1 ->
        # byte-identical baseline. See loop-attn plan.
        cfg_window = int(getattr(config, "loop_window", 0))
        env_learned = int(os.environ.get("FP_LEARNED_WINDOW", "0"))
        env_uniform = int(os.environ.get("FP_UNIFORM_WINDOW", "0"))
        self.loop_window = cfg_window or env_learned or env_uniform
        self.window_recency = float(os.environ.get("FP_WINDOW_RECENCY", "0"))  # only used by the env probe path
        learned = (cfg_window > 1) or (env_learned > 1)
        if learned and self.loop_window > 1:
            self.window_mix = WindowMix(config.hidden_size, self.loop_window,
                                        float(getattr(config, "loop_recency_init", 8.0)))
        else:
            self.window_mix = None

        self.fixed_init = config.fixed_init
        if self.fixed_init:
            fwd_dtype = getattr(torch, config.forward_dtype)
            self.init_vec = nn.Buffer(
                trunc_normal_init_(torch.empty(config.hidden_size, dtype=fwd_dtype), std=self.init_std),
                persistent=True,
            )

    def detach_state(self, state: dict):
        for k, v in state.items():
            if torch.is_tensor(v):
                state[k] = v.detach()
        return state
        
    def reset(self, reset_flag: torch.Tensor, shape: tuple, dtype: torch.dtype, 
              device: torch.device, state: dict, reset_metadata: bool = False):
        batch_size, seq_len, hidden_size = shape[0], shape[1], shape[2]
        reset_flag_1d = reset_flag.view(-1)
        reset_flag_3d = reset_flag_1d.view(-1, 1, 1)
        old_hist = state.get('hist') if state is not None else None

        if self.fixed_init:
            y = self.init_vec.to(dtype=dtype, device=device).expand(batch_size, seq_len, hidden_size).contiguous()
        else:
            y = trunc_normal_init_(torch.empty(batch_size, seq_len, hidden_size, dtype=dtype, device=device), std=self.init_std)
        residues = torch.inf * torch.ones(batch_size).to(device)
        stepsize = (self.stepsize * torch.ones(batch_size, 1, 1, dtype=dtype, device=device))
        patience = self.decay_patience * torch.ones(batch_size).to(device)
        iter_idx = torch.zeros(batch_size, dtype=torch.int32, device=device)
        best_residues = torch.inf * torch.ones(batch_size).to(device)

        if state is None:
            state = dict(y=y.contiguous(),
                         residues=residues,
                         stepsize=stepsize,
                         patience=patience,
                         iter_idx=iter_idx,
                         best_residues=best_residues)
        else:
            state = dict(y=torch.where(reset_flag_3d, y.contiguous(), state['y']),
                         residues=residues if reset_metadata else torch.where(reset_flag_1d, residues, state['residues']),
                         stepsize=stepsize if reset_metadata else torch.where(reset_flag_3d, stepsize, state['stepsize']),
                         patience=patience if reset_metadata else torch.where(reset_flag_1d, patience, state['patience']),
                         iter_idx=iter_idx if reset_metadata else torch.where(reset_flag_1d, iter_idx, state['iter_idx']),
                         best_residues=best_residues if reset_metadata else torch.where(reset_flag_1d, best_residues, state['best_residues']))

        if self.loop_window > 1:
            m = self.loop_window
            hist_new = y.detach().unsqueeze(0).repeat(m - 1, 1, 1, 1).contiguous()   # [m-1, B, S, H]
            if old_hist is None:
                state['hist'] = hist_new
            else:
                state['hist'] = torch.where(reset_flag_3d.unsqueeze(0), hist_new, old_hist)

        return state

    def step(self, state:Dict[str, torch.Tensor], y:torch.Tensor):
        state_dtype = state["y"].dtype
        if y.dtype != state_dtype:
            y = y.to(state_dtype)

        with torch.no_grad():
            residues = (state['y'].detach() - y.detach()).norm(p=torch.inf, dim=-1) / (y.detach().norm(p=torch.inf, dim=-1) + self.eps)
            residues = residues.max(dim=1)[0]

        stepsize = state['stepsize']
        if stepsize.dtype != state_dtype:
            stepsize = stepsize.to(state_dtype)
        if self.loop_window > 1 and 'hist' in state:
            prev_y = state['y']
            window = torch.cat([prev_y.unsqueeze(0), state['hist']], dim=0)          # [m, B, S, H]  (index 0 = newest)
            if self.window_mix is not None:
                blend_anchor = self.window_mix(window)                               # learned convex mix
            else:
                m = window.shape[0]
                ages = torch.arange(m, device=window.device, dtype=window.dtype)     # 0=newest .. m-1=oldest
                w = torch.softmax(-self.window_recency * ages, dim=0)                # env probe: fixed convex weights
                blend_anchor = (w.view(m, 1, 1, 1) * window).sum(dim=0)
            state['y'] = y * stepsize + blend_anchor * (1 - stepsize) + self.additive_noise_std * torch.randn_like(y)
            state['hist'] = torch.cat([prev_y.unsqueeze(0), state['hist'][:-1]], dim=0)  # roll: newest first, drop oldest
        else:
            state['y'] = y * stepsize + state['y'] * (1 - stepsize) + self.additive_noise_std * torch.randn_like(state['y'])

        # Track per-sample best iterate; require improvement of at least 0.01
        # so slow monotonic drift doesn't keep resetting patience.
        improved = residues < state['best_residues'] - 1e-2

        # update patience and lowest residue
        state['residues'] = residues
        state['best_residues'] = torch.where(improved, residues, state['best_residues'])
        state['patience'] = torch.where(improved, self.decay_patience, state['patience']-1)

        # update damping factor, reset patience
        adapt = (state['patience'] <= 0) & (state['residues'] >= self.fp_thresh)
        state['patience'] = torch.where(adapt, self.decay_patience, state['patience'])
        stepsize_dtype = state['stepsize'].dtype
        # if not self.training:
        state['stepsize'] = state['stepsize'] * torch.where(adapt, self.stepsize_decay, 1).to(stepsize_dtype).reshape(-1, 1, 1)

        state['iter_idx'] = state['iter_idx'] + 1
        return state

    def cont(self, state: Dict[str, torch.Tensor], thresh: float):
        if int(state['iter_idx'].max().item()) == 0:
            return self.max_iter > 0

        q = 1 - self.outlier_quantile if self.training else 1.0
        return (
            (torch.quantile(state['residues'].float(), q=q) >= thresh)
            & (torch.quantile(state['iter_idx'].float(), q=q) < self.max_iter)
            & (torch.quantile(state['stepsize'].float(), q=q) > 1e-3)
        )

class VariationalDropout(nn.Module):
    def __init__(self, dropout: float = 0.0):
        super().__init__()
        assert 0.0 <= dropout < 1.0, f"dropout must be in [0, 1), got {dropout}"
        self.dropout = dropout
        self._mask: torch.Tensor | None = None

    def sample_mask(self, x: torch.Tensor) -> None:
        if not self.training or self.dropout == 0.0:
            self._mask = None
            return
        keep_prob = 1.0 - self.dropout
        self._mask = torch.bernoulli(torch.full_like(x, keep_prob)) / keep_prob

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if not self.training or self.dropout == 0.0 or self._mask is None:
            return x
        return x * self._mask


class VariationalDropToken1d(nn.Module):
    def __init__(self, dropout: float = 0.0, token_first: bool = True):
        super().__init__()
        assert 0.0 <= dropout < 1.0, f"dropout must be in [0, 1), got {dropout}"
        self.dropout = dropout
        self.token_first = token_first
        self._mask: torch.Tensor | None = None

    def sample_mask(self, x: torch.Tensor) -> None:
        if not self.training or self.dropout == 0.0:
            self._mask = None
            return
        keep_prob = 1.0 - self.dropout
        if self.token_first:
            B, L, _ = x.shape
            self._mask = torch.bernoulli(torch.full((B, L, 1), keep_prob, device=x.device, dtype=x.dtype)) / keep_prob
        else:
            B, _, L = x.shape
            self._mask = torch.bernoulli(torch.full((B, 1, L), keep_prob, device=x.device, dtype=x.dtype)) / keep_prob

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if not self.training or self.dropout == 0.0 or self._mask is None:
            return x
        return x * self._mask