File size: 6,371 Bytes
2e0601e | 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 | # SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Question-local evidence routing for SPEV decision checkpoints."""
from dataclasses import dataclass
import torch
from torch import nn
from torch.nn import functional as F
class EvidenceBlock(nn.Module):
def __init__(self, width: int, heads: int, ff: int):
super().__init__()
self.qnorm = nn.LayerNorm(width)
self.mnorm = nn.LayerNorm(width)
self.attn = nn.MultiheadAttention(width, heads, batch_first=True, dropout=0.0)
self.fnorm = nn.LayerNorm(width)
self.ff = nn.Sequential(nn.Linear(width, ff), nn.GELU(), nn.Linear(ff, width))
def forward(self, query, memory, memory_pad):
memory = self.mnorm(memory)
value, _ = self.attn(self.qnorm(query), memory, memory, key_padding_mask=memory_pad, need_weights=False)
query = query + value
return query + self.ff(self.fnorm(query))
class IsolatedDecisionHead(nn.Module):
def __init__(self, hidden: int, width: int, heads: int, ff: int, depth: int):
super().__init__()
self.norm = nn.LayerNorm(hidden)
self.memory = nn.Linear(hidden, width, bias=False)
self.qproj = nn.Linear(hidden, width, bias=False)
self.oqproj = nn.Linear(hidden, width, bias=False)
self.globalproj = nn.Linear(hidden, width, bias=False)
self.context = nn.Linear(hidden, width, bias=False)
self.lexical = nn.Linear(hidden, width, bias=False)
self.types = nn.Embedding(3, width)
self.evidence = nn.ModuleList([EvidenceBlock(width, heads, ff) for _ in range(depth)])
self.summarynorm = nn.LayerNorm(width)
self.joint = nn.ModuleList(
[
nn.TransformerDecoderLayer(
width, heads, ff, dropout=0.0, activation="gelu", batch_first=True, norm_first=True
)
for _ in range(depth)
]
)
self.fnorm = nn.LayerNorm(width)
self.onorm = nn.LayerNorm(width)
self.scorer = nn.Sequential(nn.Linear(width * 4, width), nn.GELU(), nn.Linear(width, 1))
self.prior_scale = nn.Parameter(torch.zeros(()))
self.joint_scale = nn.Parameter(torch.zeros(()))
self.gate = nn.Parameter(torch.zeros(()))
def forward(self, batch: dict[str, torch.Tensor]) -> torch.Tensor:
memory = self.memory(self.norm(batch["memory"]))
question = self.norm(batch["questions"])
context = self.norm(batch["options"])
global_feature = self.norm(batch["global"])
options = self.context(context) + self.lexical(batch["lexical"]) + self.oqproj(question)
for block in self.evidence:
options = block(options, memory, batch["memory_pad"])
fields = self.qproj(question)
membership = ~batch["option_pad"][:, None, :]
weights = torch.einsum("bqd,bcd->bqc", fields, options) / memory.shape[-1] ** 0.5
weights = weights.masked_fill(~membership, -1e9).softmax(-1) * membership
summary = torch.einsum("bqc,bcd->bqd", weights, options)
fields = fields + self.summarynorm(summary) + self.globalproj(global_feature)[:, None, :]
fields = fields + self.types(batch["types"])
for block in self.joint:
fields = block(fields, memory, memory_key_padding_mask=batch["memory_pad"])
fields = self.fnorm(fields).expand(-1, options.shape[1], -1)
options = self.onorm(options)
prior = (
F.normalize(question + global_feature[:, None, :], dim=-1) * F.normalize(batch["lexical"], dim=-1)
).sum(-1) * self.prior_scale.clamp(max=4.605170186).exp()
joint = F.cosine_similarity(fields, options, dim=-1) * self.joint_scale.clamp(max=4.605170186).exp()
residual = self.scorer(torch.cat([fields, options, fields * options, (fields - options).abs()], -1)).squeeze(-1)
return prior + self.gate.sigmoid() * (joint + residual)
@dataclass(frozen=True)
class QuestionFeatures:
memory: torch.Tensor
question: torch.Tensor
options: torch.Tensor
lexical: torch.Tensor
global_feature: torch.Tensor
question_type: int
def score_questions(head: IsolatedDecisionHead, records: list[QuestionFeatures]) -> list[torch.Tensor]:
"""Score all questions in a request batch with one head invocation."""
if not records:
raise ValueError("At least one question is required")
device = head.norm.weight.device
hidden = head.norm.normalized_shape[0]
memories = max(len(record.memory) for record in records)
candidates = max(len(record.options) for record in records)
batch = {
name: torch.zeros(shape, device=device, dtype=torch.float32)
for name, shape in {
"memory": (len(records), memories, hidden),
"questions": (len(records), 1, hidden),
"options": (len(records), candidates, hidden),
"lexical": (len(records), candidates, hidden),
"global": (len(records), hidden),
}.items()
}
batch["memory_pad"] = torch.ones(len(records), memories, device=device, dtype=torch.bool)
batch["option_pad"] = torch.ones(len(records), candidates, device=device, dtype=torch.bool)
batch["types"] = torch.tensor([[record.question_type] for record in records], device=device)
for index, record in enumerate(records):
if not len(record.memory) or not len(record.options) or record.question_type not in (0, 1, 2):
raise ValueError("Invalid question features")
if record.lexical.shape != record.options.shape:
raise ValueError("Criterion embeddings must match candidate features")
for name, value in (("memory", record.memory), ("options", record.options), ("lexical", record.lexical)):
batch[name][index, : len(value)] = value.to(device=device, dtype=torch.float32)
batch["questions"][index, 0] = record.question.to(device=device, dtype=torch.float32)
batch["global"][index] = record.global_feature.to(device=device, dtype=torch.float32)
batch["memory_pad"][index, : len(record.memory)] = False
batch["option_pad"][index, : len(record.options)] = False
values = head(batch)
return [values[index, : len(record.options)] for index, record in enumerate(records)]
|