分析

flag：f1ag2026c7fa1666

根据题目要求，可以得到，这是一道线性约束求解题

查看模型定义，有如下的代码片段：

    def forward(self, input_ids, **kwargs):
        x = input_ids.float()
        hidden_states = self.dense(x)
        hidden_states = self.act(hidden_states)
        logits = self.lm_head(hidden_states)
        logits = logits.unsqueeze(1)
        return {"logits": logits}

他先将输入的16字节字符的ID转换成了浮点数，然后经过了2轮神经网络计算，最终在 inference.py 中验证输出

ID 转换

模型使用如下代码将进行转化

    def encode(self, text, return_tensors=None):
        if len(text) > 16:
            raise ValueError("Input length cannot exceed 16 characters.")
        input_ids = [self.char2id[c] for c in text]
        pad_length = 16 - len(input_ids)
        if pad_length > 0:
            input_ids.extend([self.pad_token_id] * pad_length)
        if return_tensors == "pt":
            return torch.tensor([input_ids])
        return input_ids

可以看到这是一个查表的过程，char2id的定义如下：

self.charset = "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
self.char2id = {c: i for i, c in enumerate(self.charset)}

这里实现了一个映射关系，映射关系如下：

0  -> '0'
1  -> '1'
...
10 -> 'a'
11 -> 'b'
...
61 -> 'Z'

所以最终的输入经过

x = input_ids.float()

转化之后就会变成输入字符在该表中对应的下标数列，例如 f1ag2026c7fa1666 会被转化为 [15, 1, 10, 16, 2, 0, 2, 6, 12, 7, 15, 10, 1, 6, 6, 6]

第一层神经网络计算

hidden_states = self.dense(x)
hidden_states = self.act(hidden_states)

这是第一层神经网络的代码，实际作用便是

dense(x) = W*x + b
ReLU(y) = max(0, y)

它先做一堆线性加权求和，再把负数变成 0

第二层神经网络计算

logits = self.lm_head(hidden_states)

这是第二层神经网络变化的代码，它把 21 维的隐藏结果变成 64 个输出分数，如下：

0-61: 普通字符
62: <success>
63: <fail>

求解

经过分析，发现 <fail> 的权重固定为 0.4，<success> 的权重计算如下

success_logit =
hidden[0]
- 10000000000 * hidden[1]
- 10000000000 * hidden[2]
...
- 10000000000 * hidden[20]
- 376131.22

同时要满足输出success，还得满足 logit[62] > 所有其他 logit
又因为

hidden[i] = ReLU(dense_i(x))

所以每个 hidden[i] >= 0 。因此想输出 <success>，需要满足：

hidden[1] = hidden[2] = ... = hidden[20] = 0

否则只要其中一个大于 0，-1e10 * hidden[i] 就会让 <success> 分数变成巨大负数。
同时还需要：

hidden[0] - 376131.22 > 0.4

也就是：

hidden[0] > 376131.62

所以最终条件可以写成：

hidden[1..20] 全部为 0
hidden[0] > 376131.62

因此最终利用z3编写约束求解的exp，如下：

import json
import math
import struct
from functools import reduce
from operator import mul

from z3 import Int, Or, Solver, Sum, sat


MODEL_PATH = "ictf_model/model.safetensors"
CHARSET = "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
SUCCESS_ID = 62
FAIL_ID = 63


def prod(xs):
    return reduce(mul, xs, 1)


def load_safetensors(path):
    raw = open(path, "rb").read()
    header_len = struct.unpack("<Q", raw[:8])[0]
    header = json.loads(raw[8 : 8 + header_len])
    data_base = 8 + header_len

    tensors = {}
    for name, meta in header.items():
        if name == "__metadata__":
            continue
        if meta["dtype"] != "F32":
            raise ValueError(f"unsupported dtype: {meta['dtype']}")

        shape = meta["shape"]
        start, end = meta["data_offsets"]
        blob = raw[data_base + start : data_base + end]
        count = prod(shape)
        flat = list(struct.unpack("<" + "f" * count, blob))

        if len(shape) == 1:
            tensors[name] = flat
        elif len(shape) == 2:
            rows, cols = shape
            tensors[name] = [flat[i * cols : (i + 1) * cols] for i in range(rows)]
        else:
            raise ValueError(f"unsupported shape: {shape}")

    return tensors


def ceil_gt(value):
    """Return the smallest integer n satisfying n > value."""
    return math.floor(value) + 1


def solve_flag():
    t = load_safetensors(MODEL_PATH)

    dense_w = t["dense.weight"]
    dense_b = t["dense.bias"]
    lm_w = t["lm_head.weight"]
    lm_b = t["lm_head.bias"]

    print("[+] lm_head shape: 64 x 21")
    print("[+] success row:", lm_w[SUCCESS_ID])
    print("[+] success bias:", lm_b[SUCCESS_ID])
    print("[+] fail bias:", lm_b[FAIL_ID])

    xs = [Int(f"x{i}") for i in range(16)]
    s = Solver()

    for x in xs:
        s.add(x >= 0, x < len(CHARSET))

    # success logit = hidden[0] - 1e10 * sum(hidden[1:]) + success_bias
    # Since ReLU output is non-negative, success can win only when hidden[1..20] = 0.
    for row in range(1, 21):
        expr = Sum([int(dense_w[row][col]) * xs[col] for col in range(16)]) + int(dense_b[row])
        s.add(expr <= 0)

    # Need success_logit > fail_logit:
    # hidden[0] + success_bias > fail_bias
    # hidden[0] = dense_w[0] * x + dense_b[0], and it must be positive here.
    threshold = lm_b[FAIL_ID] - lm_b[SUCCESS_ID]
    need_dot0 = ceil_gt(threshold - dense_b[0])
    expr0 = Sum([int(dense_w[0][col]) * xs[col] for col in range(16)])
    s.add(expr0 >= need_dot0)

    if s.check() != sat:
        raise RuntimeError("no solution")

    model = s.model()
    ids = [model[x].as_long() for x in xs]
    flag = "".join(CHARSET[i] for i in ids)

    # Check uniqueness.
    s.add(Or([xs[i] != ids[i] for i in range(16)]))
    unique = s.check() != sat

    return flag, ids, unique, t


def forward(flag, t):
    ids = [CHARSET.index(c) for c in flag]
    dense_w = t["dense.weight"]
    dense_b = t["dense.bias"]
    lm_w = t["lm_head.weight"]
    lm_b = t["lm_head.bias"]

    hidden = []
    for row in range(21):
        v = sum(dense_w[row][col] * ids[col] for col in range(16)) + dense_b[row]
        hidden.append(max(0.0, v))

    logits = []
    for row in range(64):
        v = sum(lm_w[row][col] * hidden[col] for col in range(21)) + lm_b[row]
        logits.append(v)

    return hidden, logits


if __name__ == "__main__":
    flag, ids, unique, tensors = solve_flag()
    hidden, logits = forward(flag, tensors)

    print("[+] ids:", ids)
    print("[+] flag:", flag)
    print("[+] unique:", unique)
    print("[+] hidden[0]:", hidden[0])
    print("[+] max(hidden[1:]):", max(hidden[1:]))
    print("[+] success_logit:", logits[SUCCESS_ID])
    print("[+] fail_logit:", logits[FAIL_ID])
    print("[+] argmax:", max(range(64), key=lambda i: logits[i]))

求解出内容如下：
利用 inference.py 验证如下：