ICTF - 题目 [AI/ML] Writeup

一个"字符级语言模型"藏着一个 flag，给模型喂对输入它就会吐出 <success>。要求逆向权重，恢复正确的 flag。

0x00 拿到手先看文件

题目给了这些东西：

inference.py      # 推理脚本，加载模型并对输入预测一个字符
model_def.py      # 模型定义（ICTFCausalLM / ICTFTokenizer）
ictf_model/       # 权重：config.json + model.safetensors
题目说明.txt

先看 inference.py，逻辑很简单：

user_input = input("Enter your prompt: ")
inputs = tokenizer(user_input, return_tensors="pt")
input_ids = inputs["input_ids"]
outputs = model(input_ids)
logits = outputs["logits"]
next_token_id = torch.argmax(logits[:, -1, :], dim=-1).item()
predicted_char = tokenizer.id2char.get(next_token_id, "<unk>")

也就是：输入一个长度 ≤16 的字符串，模型预测一个字符。我们的目标就是找到那个让模型预测出 <success>（token id = 62）的输入——它就是 flag。

0x01 模型结构其实非常朴素

model_def.py 里模型长这样：

self.dense = nn.Linear(16, 21, bias=True)
self.act = nn.ReLU()
self.lm_head = nn.Linear(21, 64, bias=True)

前向就是：

x  (16,)  --dense-->  h (21,)  --ReLU-->  h  --lm_head-->  logits (64,)

词表是 0-9a-zA-Z（62 个字符），外加 id 62 = <success>、id 63 = <fail>。padding 用 id 0。

环境里没装 torch，但模型太小了，我直接手搓前向：用 struct 解析 safetensors，ReLU 和两层 Linear 纯 Python / numpy 就能复现，和官方逻辑完全等价。

0x02 突破口：lm_head 的权重几乎全是 0

把 lm_head.weight 打出来一看，瞬间就明朗了：

• 第 62 行（对应 <success>）是：[1.0, -1e10, -1e10, -1e10, ...]（只有第 0 维系数为 1，其余 20 维全是 -1e10）
• 其余 63 行全是 0
• lm_head.bias 里，前 62 类都是 -10000，第 62 类是 -376131.21875，第 63 类（<fail>）是 0.4

所以 logits[62] 可以写成：

logits[62] = h[0] - 1e10 * (h[1] + h[2] + ... + h[20]) - 376131.21875

其它类的 logits 要么是 -10000，要么（<fail>）是 0.4。

关键推理：只要任意一个 h[i] (i>=1) 是正整数（ReLU 之后 ≥1），-1e10 的惩罚就会把 logits[62] 砸到负几千万，必然小于 0.4，模型输出 <fail>。

因此，要让模型输出 <success>，充要条件是：

1. 所有 h[i] = 0（i = 1..20），即 val_i = (Wd·x)[i] + bd[i] <= 0 全部成立；
2. h[0] > 376131.6，即 (Wd·x)[0] > 64573（因为 bd[0] = 311558.71875）。

这里 x 就是长度 16、每个分量在 0..61 的输入向量（对应字符 id）。

0x03 先随机扫了一波，发现根本碰不到

直觉告诉我这是约束满足问题，先用随机/坐标上升去搜。结果：

• 随机采样 300 万次，0 命中；
• 坐标上升跑了 2000 次重启，也是 0 命中。

这说明可行域极小，纯靠瞎猜是撞不上的。但 h[i] 都是整数（Wd、x、bd 全是整数，ReLU 出来非 0 即 ≥1），所以"关闭其余单元"是一个严格的线性不等式组。

于是换个思路：把问题建模成线性规划。

0x04 线性规划一把梭

我想最大化首隐单元的激活 (Wd·x)[0]，约束是其余所有单元关闭：

maximize   (Wd[0] · x)
s.t.       (Wd[i] · x) + bd[i] <= 0     for i = 1..20
           0 <= x_j <= 61                for j = 0..15

用 scipy.optimize.linprog 直接解（连续松弛）：

LP max (Wd x)[0] = 64573.00000000008
LP x = [15.  1. 10. 16.  2. -0.  2.  6. 12.  7. 15. 10.  1.  6.  6.  6.]

有意思的地方来了：

• 目标函数的最大值恰好就是 64573，卡在约束边界上；
• 最优解几乎全是整数（只有第 5 维是 -0.0）。

也就是说，在满足"其余单元全关"的前提下，(Wd·x)[0] 顶天只能到 64573。而我们要的阈值是 > 64572.9，整数下就是 >= 64573。这个边界解正好就是成功点，而且整数化后 h[i] 会精确等于 0（不会有浮点残余），logits[62] = 0.5 > 0.4。

把最优解转回字符（id → 0-9a-zA-Z）：

[15,1,10,16,2,0,2,6,12,7,15,10,1,6,6,6]
 f  1  a  g  2 0 2 6  c  7  f  a  1 6 6 6

得到字符串：f1ag2026c7fa1666

0x05 验证

手搓前向跑一遍：

x = [15,1,10,16,2,0,2,6,12,7,15,10,1,6,6,6]
h = [376131.71875, 0, 0, 0, ..., 0]      # 只有 h[0] 非零
logits[62] = 0.5000
logits[63] = 0.4000
argmax = 62  ->  <success>

h[1..20] 全是 0，logits[62] = 0.5 稳稳压过 <fail> 的 0.4，模型确认输出 <success>。

（顺带一提，因为最优值 64573 是 LP 的上界，任何成功的输入都必须满足 (Wd·x)[0] = 64573 且其余单元严格关闭，所以这个 leetspeak 风格的 f1ag2026c7fa1666 就是题目要的 flag。）

0x06 总结

• 模型的"秘密"不在复杂训练里，而直接写在 lm_head.weight 的异常结构上：[1, -1e10, ...] 把问题变成"第 0 个隐单元要开、其余 20 个必须全关"的硬约束；
• 一旦看清这点，整个题就退化成一个整数线性约束 + 单目标最大化，丢给 LP 求解器几毫秒出结果；
• 随机搜/梯度法都会因为可行域是 polytope 的一个零测度面而扑空，必须靠结构分析 + 精确求解。

Flag: f1ag2026c7fa1666