LLMはどう動くか:次のトークンを予測する仕組み

// 次の単語を「直前の1語」だけから予測する、おもちゃの言語モデル。
const corpus = `
the model reads the text and predicts the next token .
the model predicts the next word from the context .
the context holds the text the model can see .
the answer looks plausible and the answer can be wrong .
the next token is chosen from a list of probabilities .
`.trim().split(/\s+/);

// 1. 学習: 「ある単語の次に、どの単語が何回続いたか」を数える
const counts = new Map();
for (let i = 0; i < corpus.length - 1; i++) {
  const [a, b] = [corpus[i], corpus[i + 1]];
  if (!counts.has(a)) counts.set(a, new Map());
  counts.get(a).set(b, (counts.get(a).get(b) ?? 0) + 1);
}

// 2. 確率: 回数を、温度付きの確率に直す(温度が低いほど、一番多い単語に集中する)
function probabilities(word, temperature = 1) {
  const next = [...counts.get(word)];
  const weights = next.map(([, n]) => Math.pow(n, 1 / temperature));
  const total = weights.reduce((s, w) => s + w, 0);
  return next.map(([w], i) => [w, weights[i] / total]);
}

// 3. 生成: 確率に従って次の単語を引く。乱数は固定の種で、何度実行しても同じ結果になる
function rng(seed) {
  return () => {
    seed = (seed + 0x6D2B79F5) | 0;
    let t = Math.imul(seed ^ (seed >>> 15), 1 | seed);
    t = (t + Math.imul(t ^ (t >>> 7), 61 | t)) ^ t;
    return ((t ^ (t >>> 14)) >>> 0) / 4294967296;
  };
}
function generate(start, temperature, seed, length = 12) {
  const random = rng(seed);
  const out = [start];
  while (out.length < length && counts.has(out.at(-1))) {
    let r = random();
    for (const [w, p] of probabilities(out.at(-1), temperature)) {
      if ((r -= p) <= 0) { out.push(w); break; }
    }
  }
  return out.join(' ');
}

for (const t of [0.5, 1, 2]) {
  console.log(`温度 ${t}:`, probabilities('the', t).map(([w, p]) => `${w} ${(p * 100).toFixed(1)}%`).join(', '));
}
// 「the」の次に、最頻の単語(model か next)が選ばれる割合を、2000回の抽選で確かめる
for (const t of [0.2, 1, 3]) {
  const random = rng(42);
  let top = 0;
  for (let i = 0; i < 2000; i++) {
    let r = random();
    for (const [w, p] of probabilities('the', t)) {
      if ((r -= p) <= 0) {
        if (w === 'model' || w === 'next') top++;
        break;
      }
    }
  }
  console.log(`温度 ${t}: 最頻の単語が選ばれた割合 ${(top / 20).toFixed(2)}%`);
}
console.log('例(温度0.2):', generate('the', 0.2, 1));
console.log('例(温度3):', generate('the', 3, 1));

# 次の単語を「直前の1語」だけから予測する、おもちゃの言語モデル(toy.mjs と同じ計算を Python で書いたもの)。
corpus = """
the model reads the text and predicts the next token .
the model predicts the next word from the context .
the context holds the text the model can see .
the answer looks plausible and the answer can be wrong .
the next token is chosen from a list of probabilities .
""".split()

# 1. 学習: 「ある単語の次に、どの単語が何回続いたか」を数える
counts = {}
for a, b in zip(corpus, corpus[1:]):
    counts.setdefault(a, {})
    counts[a][b] = counts[a].get(b, 0) + 1


# 2. 確率: 回数を、温度付きの確率に直す(温度が低いほど、一番多い単語に集中する)
def probabilities(word, temperature=1):
    words = list(counts[word].items())
    weights = [n ** (1 / temperature) for _, n in words]
    total = sum(weights)
    return [(w, weight / total) for (w, _), weight in zip(words, weights)]


# 3. 生成: 確率に従って次の単語を引く。乱数は固定の種で、何度実行しても同じ結果になる
def rng(seed):
    state = seed & 0xFFFFFFFF

    def random():
        nonlocal state
        state = (state + 0x6D2B79F5) & 0xFFFFFFFF
        t = ((state ^ (state >> 15)) * (1 | state)) & 0xFFFFFFFF
        t = ((t + (((t ^ (t >> 7)) * (61 | t)) & 0xFFFFFFFF)) & 0xFFFFFFFF) ^ t
        return ((t ^ (t >> 14)) & 0xFFFFFFFF) / 4294967296

    return random


def generate(start, temperature, seed, length=12):
    random = rng(seed)
    out = [start]
    while len(out) < length and out[-1] in counts:
        r = random()
        for w, p in probabilities(out[-1], temperature):
            r -= p
            if r <= 0:
                out.append(w)
                break
    return ' '.join(out)


for t in [0.5, 1, 2]:
    print(f'温度 {t}:', ', '.join(f'{w} {p * 100:.1f}%' for w, p in probabilities('the', t)))

# 「the」の次に、最頻の単語(model か next)が選ばれる割合を、2000回の抽選で確かめる
for t in [0.2, 1, 3]:
    random = rng(42)
    top = 0
    for _ in range(2000):
        r = random()
        for w, p in probabilities('the', t):
            r -= p
            if r <= 0:
                if w in ('model', 'next'):
                    top += 1
                break
    print(f'温度 {t}: 最頻の単語が選ばれた割合 {top / 20:.2f}%')
print('例(温度0.2):', generate('the', 0.2, 1))
print('例(温度3):', generate('the', 3, 1))