Transformer Architecture

1/**2 * Example 40: Transformer Architecture3 *4 * New in v1.0.0: Full Transformer implementation including MultiheadAttention,5 * TransformerEncoder/Decoder layers, PositionalEncoding, and complete6 * encoder-decoder models for sequence-to-sequence tasks.7 */89import { randn, tensor } from "deepbox/ndarray";10import {11  Linear,12  MultiheadAttention,13  PositionalEncoding,14  TransformerDecoder,15  TransformerDecoderLayer,16  TransformerEncoder,17  TransformerEncoderLayer,18} from "deepbox/nn";1920console.log("=".repeat(60));21console.log("Example 40: Transformer Architecture");22console.log("=".repeat(60));2324function toFloat32PositionalEncoding(25  dModel: number,26  maxLen: number,27  dropout: number28): PositionalEncoding {29  const pe = new PositionalEncoding(dModel, { dropout, maxLen });30  const table: number[][] = [];3132  for (let pos = 0; pos < maxLen; pos++) {33    const row: number[] = [];34    for (let i = 0; i < dModel; i++) {35      const angle = pos / 10000 ** ((2 * Math.floor(i / 2)) / dModel);36      row.push(i % 2 === 0 ? Math.sin(angle) : Math.cos(angle));37    }38    table.push(row);39  }4041  // Align the internal buffer dtype with the rest of the example's float32 tensors.42  (pe as unknown as { peBuffer: ReturnType<typeof tensor> }).peBuffer = tensor(table, {43    dtype: "float32",44  });4546  return pe;47}4849// ============================================================================50// Part 1: MultiheadAttention51// ============================================================================52console.log("\n🔍 Part 1: MultiheadAttention");53console.log("-".repeat(60));5455// MultiheadAttention splits the embedding into multiple heads for parallel attention56const embedDim = 16;57const numHeads = 4;5859const mha = new MultiheadAttention(embedDim, numHeads);6061// Input: (batch=2, seqLen=5, embedDim=16)62const query = randn([2, 5, embedDim]);63const key = randn([2, 5, embedDim]);64const value = randn([2, 5, embedDim]);6566console.log(`MultiheadAttention(embedDim=${embedDim}, numHeads=${numHeads}):`);67console.log(`  Query shape:  ${query.shape}`);68console.log(`  Key shape:    ${key.shape}`);69console.log(`  Value shape:  ${value.shape}`);7071const attnOutput = mha.forward(query, key, value);72console.log(`  Output shape: ${attnOutput.shape}`);73console.log(`  Each head attends to ${embedDim / numHeads}D subspace`);7475// Self-attention: query = key = value76console.log("\n  Self-attention (Q=K=V):");77const selfAttnOut = mha.forward(query, query, query);78console.log(`  Output shape: ${selfAttnOut.shape}`);7980// ============================================================================81// Part 2: TransformerEncoderLayer82// ============================================================================83console.log("\n📦 Part 2: TransformerEncoderLayer");84console.log("-".repeat(60));8586// A single encoder layer: self-attention + feedforward with residual connections87const dModel = 16;88const nHead = 4;89const dFF = 64;9091const encoderLayer = new TransformerEncoderLayer({92  dModel,93  nHead,94  dFF,95  dropout: 0.1,96});9798console.log(`TransformerEncoderLayer(dModel=${dModel}, nHead=${nHead}, dFF=${dFF}):`);99100// Input: (batch=2, seqLen=8, dModel=16)101const encoderInput = randn([2, 8, dModel]);102console.log(`  Input shape:  ${encoderInput.shape}`);103104const encoderLayerOutput = encoderLayer.forward(encoderInput);105console.log(`  Output shape: ${encoderLayerOutput.shape}`);106console.log("  Components: SelfAttention → Add&Norm → FeedForward → Add&Norm");107108// ============================================================================109// Part 3: TransformerEncoder (Stacked Layers)110// ============================================================================111console.log("\n🏗️  Part 3: TransformerEncoder");112console.log("-".repeat(60));113114// Stack multiple encoder layers for deeper representations115const numEncoderLayers = 3;116const encoderLayerTemplate = new TransformerEncoderLayer({117  dModel,118  nHead,119  dFF,120  dropout: 0.1,121});122123const encoder = new TransformerEncoder(encoderLayerTemplate, numEncoderLayers);124125console.log(`TransformerEncoder(numLayers=${numEncoderLayers}):`);126const encSrc = randn([2, 10, dModel]);127console.log(`  Input shape:  ${encSrc.shape}`);128129const encoderOutput = encoder.forward(encSrc);130console.log(`  Output shape: ${encoderOutput.shape}`);131132// Count parameters133let encParams = 0;134for (const [, p] of encoder.namedParameters()) {135  encParams += p.size;136}137console.log(`  Total parameters: ${encParams}`);138139// ============================================================================140// Part 4: TransformerDecoderLayer141// ============================================================================142console.log("\n📦 Part 4: TransformerDecoderLayer");143console.log("-".repeat(60));144145// Decoder layer: self-attention + cross-attention + feedforward146const decoderLayer = new TransformerDecoderLayer(dModel, nHead, dFF, {147  dropout: 0.1,148});149150console.log(`TransformerDecoderLayer(dModel=${dModel}, nHead=${nHead}, dFF=${dFF}):`);151152// Target input: (batch=2, tgtLen=6, dModel=16)153const tgt = randn([2, 6, dModel]);154// Memory from encoder: (batch=2, srcLen=10, dModel=16)155const memory = encoderOutput;156157console.log(`  Target shape: ${tgt.shape}`);158console.log(`  Memory shape: ${memory.shape}`);159160const decoderLayerOutput = decoderLayer.forward(tgt, memory);161console.log(`  Output shape: ${decoderLayerOutput.shape}`);162console.log("  Components: SelfAttention → Add&Norm → CrossAttention → Add&Norm → FF → Add&Norm");163164// ============================================================================165// Part 5: TransformerDecoder (Stacked Layers)166// ============================================================================167console.log("\n🏗️  Part 5: TransformerDecoder");168console.log("-".repeat(60));169170const numDecoderLayers = 3;171const decoderLayerTemplate = new TransformerDecoderLayer(dModel, nHead, dFF, {172  dropout: 0.1,173});174175const decoder = new TransformerDecoder(decoderLayerTemplate, numDecoderLayers);176177console.log(`TransformerDecoder(numLayers=${numDecoderLayers}):`);178const decTgt = randn([2, 6, dModel]);179console.log(`  Target shape: ${decTgt.shape}`);180console.log(`  Memory shape: ${memory.shape}`);181182const decoderOutput = decoder.forward(decTgt, memory);183console.log(`  Output shape: ${decoderOutput.shape}`);184185let decParams = 0;186for (const [, p] of decoder.namedParameters()) {187  decParams += p.size;188}189console.log(`  Total parameters: ${decParams}`);190191// ============================================================================192// Part 6: PositionalEncoding193// ============================================================================194console.log("\n🌊 Part 6: PositionalEncoding");195console.log("-".repeat(60));196197// PositionalEncoding adds sinusoidal position information to embeddings198const pe = toFloat32PositionalEncoding(dModel, 100, 0.0);199200console.log(`PositionalEncoding(dModel=${dModel}, maxLen=100):`);201202const seqInput = randn([2, 8, dModel]);203console.log(`  Input shape:  ${seqInput.shape}`);204205const peOutput = pe.forward(seqInput);206console.log(`  Output shape: ${peOutput.shape}`);207console.log("  Adds sin/cos positional information to token embeddings");208console.log("  Even dimensions: sin(pos / 10000^(2i/d_model))");209console.log("  Odd dimensions:  cos(pos / 10000^(2i/d_model))");210211// ============================================================================212// Part 7: Complete Transformer Pipeline213// ============================================================================214console.log("\n🔄 Part 7: Complete Transformer Pipeline");215console.log("-".repeat(60));216217// Build a complete sequence-to-sequence transformer218const vocabSize = 100;219const seqLen = 12;220const batchSize = 2;221222// Source positional encoding for embedded token representations223const srcPE = toFloat32PositionalEncoding(dModel, seqLen, 0.1);224225// Target positional encoding for decoder token representations226const tgtPE = toFloat32PositionalEncoding(dModel, seqLen, 0.1);227228// Encoder229const fullEncoderLayer = new TransformerEncoderLayer({230  dModel,231  nHead,232  dFF,233  dropout: 0.1,234});235const fullEncoder = new TransformerEncoder(fullEncoderLayer, 2);236237// Decoder238const fullDecoderLayer = new TransformerDecoderLayer(dModel, nHead, dFF, {239  dropout: 0.1,240});241const fullDecoder = new TransformerDecoder(fullDecoderLayer, 2);242243// Output projection244const outputProj = new Linear(dModel, vocabSize);245246console.log("Complete Transformer Architecture:");247console.log(`  Vocabulary: ${vocabSize} tokens`);248console.log(`  Model dim: ${dModel}, Heads: ${nHead}, FF dim: ${dFF}`);249console.log(`  Encoder layers: 2, Decoder layers: 2`);250251// Forward pass252// Source tokens: (batch=2, srcLen=12)253const srcTokens = tensor(254  [255    [1, 5, 23, 42, 7, 15, 33, 8, 19, 2, 0, 0],256    [1, 12, 45, 3, 28, 9, 2, 0, 0, 0, 0, 0],257  ],258  { dtype: "int32" }259);260261// Target tokens: (batch=2, tgtLen=8)262const tgtTokens = tensor(263  [264    [1, 10, 25, 37, 14, 6, 2, 0],265    [1, 18, 44, 22, 2, 0, 0, 0],266  ],267  { dtype: "int32" }268);269270console.log(`\n  Source tokens shape: ${srcTokens.shape}`);271console.log(`  Target tokens shape: ${tgtTokens.shape}`);272273// Step 1: Start from embedded source token representations + positional encoding274const srcTokenEmbeddings = randn([batchSize, seqLen, dModel]);275const srcEmbedded = srcPE.forward(srcTokenEmbeddings);276console.log(`  Source embedded shape: ${srcEmbedded.shape}`);277278// Step 2: Encode279const fullEncoderOutput = fullEncoder.forward(srcEmbedded);280console.log(`  Encoder output shape: ${fullEncoderOutput.shape}`);281282// Step 3: Start from embedded target token representations + positional encoding283const tgtTokenEmbeddings = randn([batchSize, 8, dModel]);284const tgtEmbedded = tgtPE.forward(tgtTokenEmbeddings);285console.log(`  Target embedded shape: ${tgtEmbedded.shape}`);286287// Step 4: Decode with cross-attention to encoder output288const fullDecoderOutput = fullDecoder.forward(tgtEmbedded, fullEncoderOutput);289console.log(`  Decoder output shape: ${fullDecoderOutput.shape}`);290291// Step 5: Project to vocabulary logits292const logits = outputProj.forward(fullDecoderOutput);293console.log(`  Logits shape: ${logits.shape} (batch, tgtLen, vocab)`);294295// ============================================================================296// Summary297// ============================================================================298console.log("\n💡 Key Takeaways");299console.log("-".repeat(60));300console.log("• MultiheadAttention: parallel attention across multiple subspaces");301console.log("• TransformerEncoderLayer: self-attention + feedforward with residuals");302console.log("• TransformerDecoderLayer: self-attn + cross-attn + feedforward");303console.log("• TransformerEncoder/Decoder: stack of N identical layers");304console.log("• PositionalEncoding: sinusoidal position information for sequences");305console.log("• Full pipeline: Embed → PosEncode → Encode → Decode → Project");306console.log("• All components support batched inputs and gradient computation");307308console.log("\n✅ Transformer Architecture Example Complete!");309console.log("=".repeat(60));310