Как устроена decision-модель Clef: разбор модели принятия решений

После токенизации алгоритм запоминает позиции в промпте вопросов и опций:

def encode_record(
    tokenizer: Any,
    record: dict[str, Any],
    max_length: int = 16384,
    max_state_tokens: int | None = None,
    processor: Any | None = None,
) -> EncodedRecord:
    schema_ids = _tokens(tokenizer, "\n\nSCHEMA FIELDS:\n")
    questions: list[EncodedQuestion] = []
    for question_index, (question_id, question) in enumerate(record["questions"].items()):
        schema_ids.extend(
            _tokens(
                tokenizer,
                f"\nFIELD {question_index + 1}\nID: {question_id}\nTYPE: {question['type']}\nINSTRUCTION: ",
            )
        )
        question_start = len(schema_ids)
        instructions = question.get("instructions")
        if instructions is None or instructions == "":
            instructions = str(question_id)
        schema_ids.extend(_tokens(tokenizer, render(instructions)))
        question_end = len(schema_ids)
        schema_ids.extend(_tokens(tokenizer, "\nALLOWED OPTIONS:\n"))

        option_spans: list[tuple[int, int]] = []
        option_ids: list[str] = []
        for option_index, (option_id, description) in enumerate(question_options(question)):
            schema_ids.extend(_tokens(tokenizer, f"OPTION {option_index + 1}: "))
            option_start = len(schema_ids)
            semantics = {"option_id": option_id}
            if description is not None:
                semantics["description"] = description
            schema_ids.extend(_tokens(tokenizer, render(semantics)))
            option_spans.append((option_start, len(schema_ids)))
            option_ids.append(option_id)
            schema_ids.extend(_tokens(tokenizer, "\n"))
        schema_ids.extend(_tokens(tokenizer, "END FIELD\n"))

Каждый кусок токенизируется отдельно для того, чтобы не происходило склеивание (если токенизировать весь промпт целиком, BPE может склеить конец "OPTION 1: " с началом JSON в о10. дин токен)

После этого происходит сдвиг всей схемы на длину префикса и state:

    schema_offset = len(prefix_ids) + len(state_ids)
    shifted_questions = tuple(
        EncodedQuestion(
            question_id=question.question_id,
            question_type=question.question_type,
            question_span=(
                question.question_span[0] + schema_offset,
                question.question_span[1] + schema_offset,
            ),
            option_spans=tuple(
                (start + schema_offset, end + schema_offset)
                for start, end in question.option_spans
            ),
            option_ids=question.option_ids,
        )
        for question in questions
    )
    input_ids = tuple(prefix_ids + state_ids + schema_ids + suffix_ids)

    def forward(
        self,
        hidden_states: torch.Tensor,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor,
        records: list[EncodedRecord],
        output_embedding_weight: torch.Tensor,
    ) -> list[list[torch.Tensor]]:
        results: list[list[torch.Tensor]] = []
        normalized_hidden = self.hidden_norm(hidden_states)
        for batch_index, record in enumerate(records):
            sequence_length = int(attention_mask[batch_index].sum().item())
            sequence_hidden = normalized_hidden[batch_index, :sequence_length]
            memory = self.memory_projection(sequence_hidden).unsqueeze(0)
            global_vector = sequence_hidden[-1]
            question_vectors = torch.stack(
                [
                    self._mean_span(sequence_hidden, question.question_span)
                    for question in record.questions
                ]
            )
            type_ids = torch.tensor(
                [question.question_type for question in record.questions],
                device=hidden_states.device,
            )
            option_contexts: list[torch.Tensor] = []
            lexical_options: list[torch.Tensor] = []
            option_counts = []
            for question in record.questions:
                context_vectors = torch.stack(
                    [
                        self._mean_span(sequence_hidden, span)
                        for span in question.option_spans
                    ]
                )
                lexical_vectors = []
                for start, end in question.option_spans:
                    token_ids = input_ids[batch_index, start:end]
                    lexical_vectors.append(output_embedding_weight[token_ids].mean(dim=0))
                lexical = torch.stack(lexical_vectors)
                option_contexts.append(context_vectors)
                lexical_options.append(lexical)
                option_counts.append(len(question.option_spans))
```

Вектор вопроса. Пусть question_span team = (5, 8), то есть строки 5, 6, 7:

строка 5  "Which"   [ 0.2,  1.0, -0.4,  0.6 ]
строка 6  "team"    [ 0.8,  0.6,  0.0,  0.2 ]
строка 7  "?"       [-0.1,  0.2,  0.4,  1.0 ]
                    ───────────────────────── среднее по столбцам
q_team            = [ 0.3,  0.6,  0.0,  0.6 ]

Каждая координата - среднее трёх чисел в столбце: например, (0.2+0.8-0.1)/3=0.3.

Контекстный вектор опции. Пусть option_span billing = (9, 11), то есть строки 9 и 10:

строка  9  "billing"    [ 1.2, -0.2,  0.4,  0.0 ]
строка 10  "payments"   [ 0.6,  0.4,  0.2,  0.8 ]
                        ─────────────────────────
ctx_billing           = [ 0.9,  0.1,  0.3,  0.4 ]

Лексический вектор опции. Позиции те же (9, 11), но строки берутся не из . Сначала по спану читаются номера токенов, потом по этим номерам берутся строки lm_head (номера условные):

input_ids[9]  = 3729    →  E[3729]  = [ 0.05, -0.02,  0.10,  0.01 ]
input_ids[10] = 14310   →  E[14310] = [ 0.03,  0.04,  0.02, -0.01 ]
                                       ───────────────────────────
lex_billing                          = [ 0.04,  0.01,  0.06,  0.00 ]

Разница между ctx и lex: строка 9 в - это слово «billing», прочитанное Qwen после тикета, и она меняется от запроса к запросу. Строка - слово «billing» само по себе, и она всегда одна и та же.

g и memory. g - просто последняя строка таблицы. memory - все строки, каждая умноженная на . Спаны тут не нужны.

class EvidenceRoutingLayer(torch.nn.Module):
    def __init__(
        self,
        width: int,
        heads: int,
        feedforward: int,
        dropout: float = 0.0,
    ) -> None:
        super().__init__()
        self.query_norm = torch.nn.LayerNorm(width)
        self.memory_norm = torch.nn.LayerNorm(width)
        self.attention = torch.nn.MultiheadAttention(
            width,
            heads,
            dropout=dropout,
            batch_first=True,
        )
        self.attention_dropout = torch.nn.Dropout(dropout)
        self.feedforward_norm = torch.nn.LayerNorm(width)
        self.feedforward = torch.nn.Sequential(
            torch.nn.Linear(width, feedforward),
            torch.nn.GELU(),
            torch.nn.Dropout(dropout),
            torch.nn.Linear(feedforward, width),
            torch.nn.Dropout(dropout),
        )

    def forward(self, queries: torch.Tensor, memory: torch.Tensor) -> torch.Tensor:
        normalized_queries = self.query_norm(queries)
        routed, _ = self.attention(
            normalized_queries,
            self.memory_norm(memory),
            self.memory_norm(memory),
            need_weights=False,
        )
        queries = queries + self.attention_dropout(routed)
        return queries + self.feedforward(self.feedforward_norm(queries))

Каждая опция обрабатывается независимо и на выходе получаем распределение «важности» токена всего запроса к конкретной опции.

u_billing   = W_c·ctx_billing   + W_ℓ·lex_billing   + W_oq·q_team
u_sales     = W_c·ctx_sales     + W_ℓ·lex_sales     + W_oq·q_team
u_technical = W_c·ctx_technical + W_ℓ·lex_technical + W_oq·q_team

Сборка запросов и вызов слоёв в JointSchemaHead.forward:

option_queries = []
for question_index, (context_vectors, lexical) in enumerate(
    zip(option_contexts, lexical_options)
):
    option_queries.append(
        self.option_context_projection(context_vectors)
        + self.option_lexical_projection(lexical)
        + self.option_question_projection(
            question_vectors[question_index]
        ).unsqueeze(0)
    )
routed_options = torch.cat(option_queries, dim=0).unsqueeze(0)
for layer in self.evidence_layers:
    routed_options = layer(routed_options, memory)
routed_options = routed_options[0]
split_options = list(torch.split(routed_options, option_counts, dim=0))

base_fields = self.question_projection(question_vectors)
option_summaries = []
for field, options in zip(base_fields, split_options):
    routing_weights = torch.softmax(
        torch.matmul(options, field) / math.sqrt(options.shape[-1]),
        dim=0,
    )
    option_summaries.append(
        torch.sum(routing_weights.unsqueeze(-1) * options, dim=0)
    )
fields = (
    base_fields
    + self.option_summary_norm(torch.stack(option_summaries))
    + self.global_projection(global_vector).unsqueeze(0)
    + self.type_embedding(type_ids)
)

self.layers = torch.nn.ModuleList(
    [
        torch.nn.TransformerDecoderLayer(
            d_model=width,
            nhead=heads,
            dim_feedforward=feedforward,
            dropout=dropout,
            activation="gelu",
            batch_first=True,
            norm_first=True,
        )
        for _ in range(layers)
    ]
)
self.field_norm = torch.nn.LayerNorm(width)

Вызов в forward:

fields = fields.unsqueeze(0)
for layer in self.layers:
    fields = layer(fields, memory)
fields = self.field_norm(fields[0])

TransformerDecoderLayer- базовый слой декодеров в архитектуре Transformer.

Объявление скорера и масштабов в _init_:

self.field_norm = torch.nn.LayerNorm(width)
self.option_norm = torch.nn.LayerNorm(width)
self.residual_scorer = torch.nn.Sequential(
    torch.nn.Linear(width * 4, width),
    torch.nn.GELU(),
    torch.nn.Dropout(dropout),
    torch.nn.Linear(width, 1),
)
self.prior_logit_scale = torch.nn.Parameter(torch.zeros(()))
self.joint_logit_scale = torch.nn.Parameter(torch.zeros(()))
self.residual_gate = torch.nn.Parameter(torch.zeros(()))

в forward:

record_logits: list[torch.Tensor] = []
for field, question, lexical, routed in zip(
    fields,
    record.questions,
    lexical_options,
    split_options,
):
    anchor = functional.normalize(
        question_vectors[len(record_logits)] + global_vector,
        dim=-1,
    )
    lexical_anchor = functional.normalize(lexical, dim=-1)
    prior_scale = self.prior_logit_scale.clamp(max=math.log(100.0)).exp()
    prior = prior_scale * torch.matmul(lexical_anchor, anchor)
    options = self.option_norm(routed)
    repeated_field = field.unsqueeze(0).expand_as(options)
    cosine = functional.cosine_similarity(repeated_field, options, dim=-1)
    features = torch.cat(
        [
            repeated_field,
            options,
            repeated_field * options,
            torch.abs(repeated_field - options),
        ],
        dim=-1,
    )
    residual = self.residual_scorer(features).squeeze(-1)
    joint_scale = self.joint_logit_scale.clamp(max=math.log(100.0)).exp()
    joint = joint_scale * cosine + residual
    record_logits.append(
        prior + torch.sigmoid(self.residual_gate) * joint
    )
results.append(record_logits)
return results