TypeShade Compute로 MNIST 학습

손글씨 숫자를 분류하는 소프트맥스 모델을 학습합니다. 모델 계산은 TypeShade가 맡고, 호스트는 데이터를 읽어 배치를 제어합니다.

실행 범위: 실습 도구는 실제 MNIST 데이터 일부를 브라우저에서 학습합니다. 아래 전체 데이터 수치는 별도로 진행한 재현 실험 결과입니다.

브라우저에서 MNIST 직접 학습

실제 TypeShade Compute 프로그램으로 학습하고, 테스트 숫자를 확인한 뒤 직접 그린 숫자를 분류합니다.

실제로 실행되는 TypeShade 코드: 순전파

아래 forward는 28×28 픽셀을 열 개 숫자의 점수로 변환하는 실제 Compute 함수입니다. 학습과 추론이 이 코드를 함께 사용합니다.

@compute([64])
export function forward(@builtin("global_invocation_id") gid: vec3u) {
  const row = gid.x;
  if (row >= batch.count) { return; }
  for (let c: u32 = 0; c < 10; c++) {
    let z = bias[c];
    for (let p: u32 = 0; p < 784; p++) {
      z += pixels[(batch.offset + row) * 784 + p] * weights[p * 10 + c];
    }
    logits[row * 10 + c] = z;
  }
}

전체 학습 과정과 코드 해설 7단계 보기

실행 중인 TypeShade 원본 코드 보기

위 MNIST 모델을 WebGPU/WebGL2용으로 컴파일한 실제 softmax.shade.ts 원본입니다. 학습과 추론에 사용되는 여섯 개의 Compute 함수를 모두 볼 수 있습니다.

"use typeshade";

class Batch {
  count: u32;
  offset: u32;
  rate: f32;
}
declare const batch: uniform<Batch>;
declare const pixels: storage<array<f32>>;
declare const labels: storage<array<u32>>;
declare const weights: storage<array<f32>, "read_write">;
declare const bias: storage<array<f32>, "read_write">;
declare const logits: storage<array<f32>, "read_write">;
declare const delta: storage<array<f32>, "read_write">;
declare const losses: storage<array<f32>, "read_write">;
declare const stats: storage<array<f32>, "read_write">;
declare const gradW: storage<array<f32>, "read_write">;
declare const gradB: storage<array<f32>, "read_write">;
declare const probabilities: storage<array<f32>, "read_write">;
declare const predicted: storage<array<u32>, "read_write">;

function stableProbability(z: f32, peak: f32, total: f32): f32 {
  return exp(z - peak) / total;
}

@compute([64])
export function forward(@builtin("global_invocation_id") gid: vec3u) {
  const row = gid.x;
  if (row >= batch.count) { return; }
  for (let c: u32 = 0; c < 10; c++) {
    let z = bias[c];
    for (let p: u32 = 0; p < 784; p++) {
      z += pixels[(batch.offset + row) * 784 + p] * weights[p * 10 + c];
    }
    logits[row * 10 + c] = z;
  }
}

@compute([64])
export function objective(@builtin("global_invocation_id") gid: vec3u) {
  const row = gid.x;
  if (row >= batch.count) { return; }
  let peak = logits[row * 10];
  for (let c: u32 = 1; c < 10; c++) { peak = max(peak, logits[row * 10 + c]); }
  let total: f32 = 0.;
  for (let c: u32 = 0; c < 10; c++) { total += exp(logits[row * 10 + c] - peak); }
  const label = labels[batch.offset + row];
  // Subtract before adding log(total), retaining small losses at large offsets.
  losses[row] = (peak - logits[row * 10 + label]) + log(total);
  for (let c: u32 = 0; c < 10; c++) {
    delta[row * 10 + c] = (stableProbability(logits[row * 10 + c], peak, total) - (c === label ? 1. : 0.)) / f32(batch.count);
  }
}

@compute([1])
export function reduce() {
  let total: f32 = 0.;
  let correct: f32 = 0.;
  for (let row: u32 = 0; row < batch.count; row++) {
    total += losses[row];
    let best: u32 = 0;
    for (let c: u32 = 1; c < 10; c++) {
      if (logits[row * 10 + c] > logits[row * 10 + best]) { best = c; }
    }
    if (best === labels[batch.offset + row]) { correct += 1.; }
  }
  stats[0] = total / f32(batch.count);
  stats[1] = correct;
}

// Explicit backward, not automatic differentiation. Storage-array reverse mode from
// change 0056 is not implemented on the inspected main. One writer per parameter.
@compute([64])
export function backward(@builtin("global_invocation_id") gid: vec3u) {
  const k = gid.x;
  if (k < 7840) {
    const p = k / 10;
    const c = k % 10;
    let sum: f32 = 0.;
    for (let row: u32 = 0; row < batch.count; row++) {
      sum += pixels[(batch.offset + row) * 784 + p] * delta[row * 10 + c];
    }
    gradW[k] = sum;
  }
  if (k < 10) {
    let sum: f32 = 0.;
    for (let row: u32 = 0; row < batch.count; row++) { sum += delta[row * 10 + k]; }
    gradB[k] = sum;
  }
}

@compute([64])
export function update(@builtin("global_invocation_id") gid: vec3u) {
  const k = gid.x;
  if (k < 7840) { weights[k] -= batch.rate * gradW[k]; }
  if (k < 10) { bias[k] -= batch.rate * gradB[k]; }
}

// One label-free inference entry, shared by the browser and the experiment host.
// The first forward row is the digit being classified. All probability math stays
// inside TypeShade, not in a second JavaScript implementation.
@compute([1])
export function predict() {
  let peak = logits[0];
  let best: u32 = 0;
  for (let c: u32 = 1; c < 10; c++) {
    const z = logits[c];
    if (z > peak) { peak = z; best = c; }
  }
  let total: f32 = 0.;
  for (let c: u32 = 0; c < 10; c++) { total += exp(logits[c] - peak); }
  for (let c: u32 = 0; c < 10; c++) {
    probabilities[c] = stableProbability(logits[c], peak, total);
  }
  predicted[0] = best;
}

GitHub에서 이 소스 버전 열기

백엔드 학습 이미지 에폭

1에서 100까지 정수를 입력하세요. 에폭마다 선택한 학습 이미지 전체를 한 번씩 처리합니다.

에폭을 늘리면 GPU 연산이 반복되어 실행 시간, 전력 소모와 발열이 늘어납니다. 실행 시간은 백엔드와 이미지 수에 따라 달라집니다. 에폭이 많아도 평가 정확도가 좋아진다는 보장은 없습니다. 중단은 현재 배치가 끝난 뒤 처리됩니다.

시작하면 공식 MNIST gzip 파일 약 12 MB를 다운로드합니다. 데이터는 브라우저 안에서 처리하며 이미지나 그림을 TypeShade 서버에 보내지 않습니다.

백엔드를 고르고 학습을 시작하세요.

막대는 전체 에폭의 학습 이미지 처리 수를 표시합니다. 평가는 에폭마다 별도로 진행합니다.

학습 과정 보기

학습 배치마다 숫자 점수를 예측하고 손실과 기울기를 계산해 가중치를 갱신합니다. 에폭이 끝나면 가중치를 바꾸지 않고 갱신된 모델을 평가합니다.

백엔드를 고르고 학습을 시작하세요.

그래프는 에폭마다 같은 테스트 이미지 64개로 측정한 결과입니다. 손실이 낮을수록 예측 오차가 작고, 정확도는 정답을 맞힌 비율입니다. 이 일부 이미지가 MNIST 테스트 전체를 대표하지는 않습니다.

에폭별 평가 교차 엔트로피 손실

에폭별 평가 정확도

숫자 그리기

학습 후 검은 사각형에 흰색 숫자를 그려 보세요. TypeShade 순전파와 예측 커널로 계산합니다.

숫자 0에서 9까지의 점수

학습 진행 상황

호스트는 이미지를 읽고 배치를 제어합니다. 학습과 추론 로짓은 선택한 TypeShade WebGPU 또는 WebGL2 백엔드에서 계산합니다. WebGL2는 GLSL ES 3.00 패스를 사용합니다.

학습 파이프라인

입력 특성은 784개, 출력 클래스는 10개입니다. 학습과 추론은 TypeShade Compute 진입점 여섯 개로 계산합니다.

아래 코드는 실제 브라우저 실습에서 컴파일해 실행하는 TypeShade 원본입니다. 모델은 28×28 픽셀(784개)을 입력받아 0~9까지의 점수 10개를 계산하는 소프트맥스 회귀입니다. CNN이나 사전 학습된 모델이 아닙니다.

학습: forward → objective → backward → update

평가: forward → objective → reduce

직접 그린 숫자 예측: forward → predict

각 단계는 계산의 목적과 구현을 설명합니다. 모든 코드 블록은 컴파일에 사용되는 softmax.shade.ts에서 자동 추출합니다.

1. TypeScript에서 GPU 데이터를 선언하기

"use typeshade"로 시작한 파일에서 픽셀, 정답, 가중치, 편향, 기울기와 결과 버퍼를 선언합니다. batch는 현재 미니배치 크기와 학습률을 전달합니다.

"use typeshade";

class Batch {
  count: u32;
  offset: u32;
  rate: f32;
}
declare const batch: uniform<Batch>;
declare const pixels: storage<array<f32>>;
declare const labels: storage<array<u32>>;
declare const weights: storage<array<f32>, "read_write">;
declare const bias: storage<array<f32>, "read_write">;
declare const logits: storage<array<f32>, "read_write">;
declare const delta: storage<array<f32>, "read_write">;
declare const losses: storage<array<f32>, "read_write">;
declare const stats: storage<array<f32>, "read_write">;
declare const gradW: storage<array<f32>, "read_write">;
declare const gradB: storage<array<f32>, "read_write">;
declare const probabilities: storage<array<f32>, "read_write">;
declare const predicted: storage<array<u32>, "read_write">;

function stableProbability(z: f32, peak: f32, total: f32): f32 {
  return exp(z - peak) / total;
}

픽셀 한 장은 784개의 f32이며 10개 클래스의 가중치는 784×10=7,840개입니다. storage 버퍼의 가중치와 편향은 학습 중 GPU에 상주하고 다음 배치에서도 재사용됩니다. stableProbability는 큰 로짓에도 안정적인 소프트맥스 값을 계산하는 공통 함수입니다.

2. 순전파: 픽셀을 클래스 점수로 바꾸기

forward는 각 이미지의 784개 픽셀과 10개 클래스 가중치를 곱해 더하고 편향을 더해 로짓(logit)을 만듭니다.

@compute([64])
export function forward(@builtin("global_invocation_id") gid: vec3u) {
  const row = gid.x;
  if (row >= batch.count) { return; }
  for (let c: u32 = 0; c < 10; c++) {
    let z = bias[c];
    for (let p: u32 = 0; p < 784; p++) {
      z += pixels[(batch.offset + row) * 784 + p] * weights[p * 10 + c];
    }
    logits[row * 10 + c] = z;
  }
}

gid.x는 처리할 이미지의 행 번호입니다. 각 클래스 c에 대해 z = bias[c] + sum pixels[p]×weights[p×10+c]를 계산합니다. @compute([64])는 작업 그룹의 호출 크기이며 이미지 수가 작업 그룹 크기와 반드시 같다는 뜻은 아닙니다. 마지막 배치의 범위를 검사합니다.

3. 손실과 오차 신호: 왜 학습이 필요한지 계산

objective는 로짓의 최대값을 먼저 빼는 안정적 소프트맥스를 사용해 정답 클래스의 교차 엔트로피 손실을 계산합니다.

@compute([64])
export function objective(@builtin("global_invocation_id") gid: vec3u) {
  const row = gid.x;
  if (row >= batch.count) { return; }
  let peak = logits[row * 10];
  for (let c: u32 = 1; c < 10; c++) { peak = max(peak, logits[row * 10 + c]); }
  let total: f32 = 0.;
  for (let c: u32 = 0; c < 10; c++) { total += exp(logits[row * 10 + c] - peak); }
  const label = labels[batch.offset + row];
  // Subtract before adding log(total), retaining small losses at large offsets.
  losses[row] = (peak - logits[row * 10 + label]) + log(total);
  for (let c: u32 = 0; c < 10; c++) {
    delta[row * 10 + c] = (stableProbability(logits[row * 10 + c], peak, total) - (c === label ? 1. : 0.)) / f32(batch.count);
  }
}

오차 신호 delta는 (예측 확률 - 정답 원-핫 값) / 배치 크기입니다. 이 식이 다음 backward에서 각 가중치가 어느 방향으로 변해야 하는지 알려 줍니다. 여기서 exp와 log도 TypeShade 커널 안에서 계산됩니다.

4. 역전파: 모든 가중치의 기울기를 합산

backward는 현재 배치의 픽셀과 delta를 사용해 7,840개 가중치 및 10개 편향의 기울기를 계산합니다.

@compute([64])
export function backward(@builtin("global_invocation_id") gid: vec3u) {
  const k = gid.x;
  if (k < 7840) {
    const p = k / 10;
    const c = k % 10;
    let sum: f32 = 0.;
    for (let row: u32 = 0; row < batch.count; row++) {
      sum += pixels[(batch.offset + row) * 784 + p] * delta[row * 10 + c];
    }
    gradW[k] = sum;
  }
  if (k < 10) {
    let sum: f32 = 0.;
    for (let row: u32 = 0; row < batch.count; row++) { sum += delta[row * 10 + k]; }
    gradB[k] = sum;
  }
}

각 가중치의 gradW[p×10+c] = sum pixels[row,p]×delta[row,c]입니다. gradB는 해당 클래스의 delta 합계입니다. 이 구현은 자동 미분이 아니라 TypeShade에 직접 작성한 명시적 기울기 계산입니다.

5. SGD: GPU 버퍼의 모델을 직접 갱신

update는 계산된 기울기에 학습률을 곱해 기존 가중치와 편향에서 빼는 SGD 단계를 실행합니다.

@compute([64])
export function update(@builtin("global_invocation_id") gid: vec3u) {
  const k = gid.x;
  if (k < 7840) { weights[k] -= batch.rate * gradW[k]; }
  if (k < 10) { bias[k] -= batch.rate * gradB[k]; }
}

// One label-free inference entry, shared by the browser and the experiment host.
// The first forward row is the digit being classified. All probability math stays
// inside TypeShade, not in a second JavaScript implementation.

weights[k] -= batch.rate × gradW[k]가 핵심입니다. 결과를 매 배치 JavaScript로 읽어 모델을 갱신하지 않습니다. 가중치는 GPU storage에 남고 다음 forward가 갱신된 값을 읽습니다.

6. 평가: 평균 손실과 맞힌 개수 집계

reduce는 각 이미지의 손실과 가장 높은 로짓을 확인해 배치의 평균 손실 및 정답 수를 계산합니다.

@compute([1])
export function reduce() {
  let total: f32 = 0.;
  let correct: f32 = 0.;
  for (let row: u32 = 0; row < batch.count; row++) {
    total += losses[row];
    let best: u32 = 0;
    for (let c: u32 = 1; c < 10; c++) {
      if (logits[row * 10 + c] > logits[row * 10 + best]) { best = c; }
    }
    if (best === labels[batch.offset + row]) { correct += 1.; }
  }
  stats[0] = total / f32(batch.count);
  stats[1] = correct;
}

// Explicit backward, not automatic differentiation. Storage-array reverse mode from
// change 0056 is not implemented on the inspected main. One writer per parameter.

stats[0]은 배치 평균 손실, stats[1]은 맞힌 개수입니다. 이 커널이 합산하는 것은 가중치 기울기가 아니라 평가 지표입니다. 호스트는 작은 stats 배열만 읽어 전체 평가 결과를 합산합니다.

7. 추론: 0~9 확률과 최종 숫자를 GPU에서 결정

predict는 forward가 기록한 10개 로짓을 정규화해 확률 분포를 만들고 가장 큰 값을 가진 숫자를 선택합니다.

@compute([1])
export function predict() {
  let peak = logits[0];
  let best: u32 = 0;
  for (let c: u32 = 1; c < 10; c++) {
    const z = logits[c];
    if (z > peak) { peak = z; best = c; }
  }
  let total: f32 = 0.;
  for (let c: u32 = 0; c < 10; c++) { total += exp(logits[c] - peak); }
  for (let c: u32 = 0; c < 10; c++) {
    probabilities[c] = stableProbability(logits[c], peak, total);
  }
  predicted[0] = best;
}

확률은 probabilities[c]에, 예측 숫자는 predicted[0]에 기록됩니다. 브라우저 JavaScript는 이 값을 읽어 막대그래프를 표시할 뿐, Math.exp로 소프트맥스를 다시 계산하지 않습니다. 정답 라벨도 추론에는 필요하지 않습니다.

이 TypeShade 커널들이 브라우저에서 어떻게 연결될까요?

실습의 JavaScript는 공식 IDX 파일을 다운로드하고 픽셀을 0~1로 정규화한 뒤 TypeShade 프로그램의 Compute 단계들을 호출합니다. 아래는 별도의 의사 코드가 아니라 실제 공통 브라우저 세션의 메서드입니다.

학습 배치 하나를 실행하는 실제 코드

trainBatch(pixels, labels, rate) {
      return queue(async () => {
        if (!Number.isFinite(rate) || rate <= 0) throw new RangeError('Invalid learning rate');
        setBatch(pixels, labels, labels.length);
        for (const entry of ['forward', 'objective', 'backward', 'update'] as const)
          await dispatch(entry, labels.length, rate);
      });
    },

setBatch가 픽셀과 정답을 GPU 입력 버퍼에 기록하고 forward → objective → backward → update를 순서대로 실행합니다. 동일한 가중치와 편향 버퍼가 다음 배치까지 유지됩니다.

손글씨 한 장을 분류하는 실제 코드

predict(pixels) {
      return queue(async () => {
        if (pixels.length !== INPUTS) throw new RangeError('MNIST inference requires 784 pixels');
        setBatch(pixels, new Uint32Array(1), 1);
        await dispatch('forward', 1, 0);
        await dispatch('predict', 1, 0);
        const probabilities = (await data.probabilities.read()) as Float32Array;
        const prediction = (await data.predicted.read()) as Uint32Array;
        const sum = probabilities.reduce((a, b) => a + b, 0);
        if (
          !probabilities.every((value) => Number.isFinite(value) && value >= 0) ||
          Math.abs(sum - 1) > 0.0002 ||
          prediction[0] >= CLASSES
        )
          throw new Error('Invalid TypeShade inference probability output');
        return { probabilities, predicted: prediction[0] };
      });
    },

forward → predict로 GPU에서 클래스별 확률과 최종 숫자를 계산하고, 10개 확률과 정수 결과만 브라우저로 읽습니다. 훈련·예측 요청은 직렬화되어 동시 GPU 버퍼 덮어쓰기를 방지합니다.

실행 경계: 모델 수학(순전파, 소프트맥스, 손실, 기울기, SGD, 추론)은 TypeShade가 담당합니다. 데이터 로딩, 배치 제어, 캔버스, 화면 표시, 최종 수치 집계는 JavaScript 호스트가 담당합니다. WebGL2는 네이티브 Compute가 아니라 TypeShade가 생성한 GLSL ES 3.00 패스로 실행됩니다.

검증된 결과

전체 데이터 결과는 저장소에 기록된 수치입니다. WebGL2 검증은 더 작은 데이터셋으로 별도 진행했습니다.

실행 백엔드	학습 / 평가, 에폭	평가 정확도
생성된 CPU	60,000 / 10,000, 5 에폭	91.67%
WebGPU (SwiftShader)	60,000 / 10,000, 5 에폭	91.67%
WebGL2 패스	1,024 train / 1,000 test subset	전체 데이터 결과 미검증

WebGPU 측정에는 소프트웨어 렌더러 Google SwiftShader를 사용했습니다. 물리 GPU 가속 성능을 입증한 결과가 아닙니다.

WebGL2 테스트는 모든 패스를 실제 실행해 독립 수치 기준과 대조하고 일부 실제 데이터도 학습했습니다. 근거: MNIST WebGL2 검증 워크플로.

실험 재현하기

컴파일러 저장소 루트에서 Node 24와 Bun을 준비하고 의존성을 설치한 뒤 WebGL2 테스트를 실행합니다.

bun install
bun run build
./node_modules/.bin/playwright install --only-shell chromium
node --experimental-strip-types journeys/mnist/run.ts --download --train 1 --test 1 --epochs 1
node --experimental-strip-types journeys/mnist/test-webgl2.mjs --mnist --count 1024 --test 1000 --epochs 5

데이터를 다운로드한 뒤 WebGPU 백엔드로 전체 학습을 실행하는 명령어:

node --experimental-strip-types journeys/mnist/run.ts --tier webgpu

이 명령어는 로컬 Chromium을 사용하며 필요할 때 MNIST 데이터를 다운로드합니다. 성능은 어댑터와 호스트 환경에 따라 달라집니다.

검증 범위와 한계

WebGL2는 네이티브 컴퓨트 셰이더가 아니라 생성된 GLSL ES 3.00 패스를 사용합니다.
WebGL2 CI는 수치 계산과 실제 데이터 일부를 검증했습니다. 전체 데이터 정확도와 속도는 아직 검증하지 않았습니다.
공개된 전체 WebGPU 학습 결과는 소프트웨어 렌더링으로 측정했습니다. 물리 GPU 가속 성능은 주장하지 않습니다.
실습 도구는 선택한 데이터 일부로 브라우저에서 TypeShade 계산을 실행합니다. 전체 데이터 실험 결과는 별도 기록입니다.

소스와 검증 기록

소스와 결과는 사이트 빌드에 사용한 컴파일러 커밋을 가리킵니다.

TypeNN의 MNIST

이 가이드는 직접 작성한 커널로 소프트맥스 회귀를 학습합니다. TypeNN은 TypeShade 위에 만든 신경망 라이브러리입니다. 모델은 레이어를 쌓아 정의하고, 각 레이어는 TypeShade 커널 함수로 실행됩니다. 아래 코드는 은닉층이 하나인 모델을 정의하고 배치 단위로 학습합니다.

import * as nn from "typenn";

const device = await nn.device();
const model = new nn.Sequential(
  new nn.Flatten(),
  new nn.Linear(784, 128),
  new nn.ReLU(),
  new nn.Dropout(0.2),
  new nn.Linear(128, 10),
).to(device);
const optimizer = nn.optim.Adam(model.parameters(), { lr: 0.001 });
const criterion = new nn.CrossEntropyLoss();

for await (const batch of batches) {
  await optimizer.zeroGrad();
  const images = device.tensor(batch.images, {
    shape: [batch.size, 1, 28, 28],
  });
  const loss = criterion.forward(model.forward(images), batch.labels);
  await loss.backward();
  await optimizer.step();
  await loss.disposeGraph(model.parameters());
}

이 페이지 편집 문제 보고