EDMI

EDMI - модуль для LLM, который динамически усиливает скрытые состояния модели. Он обнаруживает слабые токены, вычисляет небольшое улучшение через усилитель и внедряет его обратно, не нарушая генерацию.

Принцип работы прост: для каждого слоя выбираются несколько наиболее слабых токенов, их скрытые состояния корректируются через внутренние параметры, а сила влияния регулируется небольшим параметром scale.

Пример кода:

import torch
import torch.nn as nn

class EDMI(nn.Module):
    def __init__(self, d_model, top_k=2, layer_idx=0):
        super().__init__()
        self.top_k = top_k
        self.scale = nn.Parameter(torch.tensor(4e-4 * (0.75 ** layer_idx)))
        self.detector = nn.Linear(d_model, 1, bias=False)
        self.enhancer = nn.Sequential(
            nn.Linear(d_model, d_model // 4),
            nn.GELU(),
            nn.Linear(d_model // 4, d_model)
        )
        self.gate = nn.Sigmoid()

    def forward(self, h):
        b, seq_len = h.shape[0], h.shape[1]
        weakness = torch.sigmoid(self.detector(h)).squeeze(-1)
        k = min(self.top_k, seq_len)
        if k == 0:
            return h
        topk_idx = torch.topk(weakness, k=k, dim=1, largest=True).indices
        mask = torch.zeros_like(weakness).scatter_(1, topk_idx, 1.0).unsqueeze(-1)
        delta = self.scale * mask * self.gate(self.enhancer(h))
        return h + delta