ICTFForCausalLM Writeup

题目信息

附件给出一个字符级「因果语言模型」ICTFForCausalLM、权重 ictf_model/，以及加载脚本 inference.py。

题面要求：分析模型代码与权重，找出藏在模型里的秘密，恢复正确 flag。

输入：长度 ≤ 16 的字符串（字符集为 [0-9a-zA-Z]）
输出：模型预测的下一个 token

一、模型结构

model_def.py 里的网络非常「不像」真 LM：

input_ids ∈ R^{16}     # token id 直接当浮点特征，不是 embedding
    ↓
Linear(16 → 21) + bias
    ↓
ReLU
    ↓
Linear(21 → 64) + bias   # lm_head
    ↓
logits ∈ R^{64}

Tokenizer 要点：

id	含义
0–61	`0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ`
62	`<success>`
63	`<fail>`

不足 16 字符会用 pad_token_id = 0（字符 '0'）右侧填充。

inference.py 对最后一维做 argmax，打印对应字符。真正的成功汇点是预测出 <success> (id=62)。

二、从权重读出判定逻辑

加载 model.safetensors 后看 lm_head：

token 0..61：权重全 0，bias ≈ -10000 → 永远选不中
token 63 (<fail>)：权重全 0，bias ≈ 0.4 → 默认失败基线
token 62 (<success>)：

w_success ≈ [1, -1e10, -1e10, ..., -1e10]   # 长度 21
b_success ≈ -376131.22

因此：

logit_62 = h[0] + (-1e10)·h[1] + … + (-1e10)·h[20] + b_success
logit_63 = 0.4

只要 h[1]…h[20] 里有任意一个 ReLU 后为正，logit_62 会被打到 −∞，必然失败。
要成功，必须同时满足：

校验神经元全灭：h[i] = 0（即 ReLU 前 pre[i] ≤ 0），i = 1..20
主通道够大：h[0] + b_success > 0.4，即 h[0] ≳ 376131.6

这不是训练出来的语言模型，而是把 口令校验硬编码进 MLP 的壳。

三、dense 层在干什么

记输入向量为 x ∈ {0,…,61}^{16}（各位置 token id），则

pre = W·x + b      # W ∈ R^{21×16}, b ∈ R^{21}
h   = ReLU(pre)

观察权重：

W 的元素都是精确整数（浮点误差为 0）
b[1:] 近似整数；b[0] 带一点小数，配合 lm_head bias 做阈值比较

前 16 个「校验」神经元（i = 1..16）构成线性方程组：

W[i] · x + b[i] = 0,   i = 1..16

16 个未知数、16 个方程，且系数矩阵满秩 → 唯一解。

用最小二乘或 Z3 整数求解均可，得到：

x = [15, 1, 10, 16, 2, 0, 2, 6, 12, 7, 15, 10, 1, 6, 6, 6]

映射回字符集：

f1ag2026c7fa1666

代入后：

pre[0]  ≈ 376131.72 > 阈值
pre[1..20] = 0
→ h = [376131.72, 0, 0, …, 0]
→ logit_62 ≈ 0.5 > logit_63 ≈ 0.4
→ argmax = 62 (<success>)

i = 17..20 的额外行在该解上同样为 0，属于冗余/加固约束，不改变唯一性。

四、求解流程（可复现）

Inventory
读 model_def.py / config.json：无 embedding、无 attention，本质是 16→21→64 的两层感知机。
钉汇点
成功 = argmax(logits) == 62，不是去「生成一段自然语言」。
切片 lm_head
从 success 行的 -1e10 惩罚读出：校验维必须被 ReLU 清零。
反推 dense
对 W[1:17]·x + b[1:17] = 0 求解（Z3 / numpy.linalg.solve / lstsq）。
Prove
把候选喂回 ICTFForCausalLM.from_pretrained("./ictf_model")，确认输出 <success>。

最小验证脚本：

import torch
from model_def import ICTFForCausalLM, ICTFTokenizer

text = "f1ag2026c7fa1666"
tok = ICTFTokenizer()
model = ICTFForCausalLM.from_pretrained("./ictf_model").eval()
ids = tok(text, return_tensors="pt")["input_ids"]
with torch.no_grad():
    logits = model(ids)["logits"]
    nid = int(torch.argmax(logits[:, -1, :], dim=-1))
print(nid, tok.id2char[nid])  # 62 <success>

五、原理小结

表象	实质
Causal LM + safetensors	口令比较器换皮
hidden_size = 21	1 个成功通道 + 20 个校验通道
ReLU	把「等式成立」变成「激活为 0」
`lm_head` 的 ±1e10	硬门控：任一校验维泄漏 → 强制 fail
字符 id 当特征	线性约束直接作用在 id 空间，便于精确构造