Update index.html

index.html

@@ -120,6 +120,9 @@
             const attentionBq = includeBias ? hiddenSize : 0;
             const attentionBk = includeBias ? numKeyValueHeads : 0;
             const attentionBv = includeBias ? numKeyValueHeads : 0;
+            const attentionQOutput = hiddenSize;
+            const attentionKOutput = numKeyValueHeads;
+            const attentionVOutput = numKeyValueHeads;
             const attentionTotal = 2 * (attentionQ + attentionBq) + 2 * (attentionK + attentionBk) + 2 * (attentionV + attentionBv);
 
             // Feed Forward calculations
@@ -145,23 +148,27 @@
 
             // Display results
             const attentionTable = document.getElementById('attention_table');
-            attentionTable.innerHTML = `
-                <tr><th>Type</th><th>Q</th><th>K</th><th>V</th><th>Total</th></tr>
-                <tr><td>Weights</td><td>${attentionQ.toLocaleString()}</td><td>${attentionK.toLocaleString()}</td><td>${attentionV.toLocaleString()}</td><td>${attentionTotal.toLocaleString()}</td></tr>
+            attentionTable.innerHTML = ` 
+                <tr><th>Type</th><th>Input Size</th><th>Output Size</th><th>Total Size</th></tr>
+                <tr><td>Q</td><td>${hiddenSize.toLocaleString()}</td><td>${attentionQOutput.toLocaleString()}</td><td>${(hiddenSize * attentionQOutput).toLocaleString()}</td></tr>
+                <tr><td>K</td><td>${hiddenSize.toLocaleString()}</td><td>${attentionKOutput.toLocaleString()}</td><td>${(hiddenSize * attentionKOutput).toLocaleString()}</td></tr>
+                <tr><td>V</td><td>${hiddenSize.toLocaleString()}</td><td>${attentionVOutput.toLocaleString()}</td><td>${(hiddenSize * attentionVOutput).toLocaleString()}</td></tr>
                 ${includeBias ? `<tr class="bias-row"><td>Bias</td><td>${attentionBq.toLocaleString()}</td><td>${attentionBk.toLocaleString()}</td><td>${attentionBv.toLocaleString()}</td><td></td></tr>` : ''}
             `;
 
             const feedForwardTable = document.getElementById('feed_forward_table');
             feedForwardTable.innerHTML = `
-                <tr><th>Type</th><th>Switch W</th><th>LU W</th><th>Proj W</th><th>Total</th></tr>
-                <tr><td>Weights</td><td>${switchW.toLocaleString()}</td><td>${luW.toLocaleString()}</td><td>${projW.toLocaleString()}</td><td>${feedForwardTotal.toLocaleString()}</td></tr>
+                <tr><th>Type</th><th>Input Size</th><th>Output Size</th><th>Total Size</th></tr>
+                <tr><td>Switch W</td><td>${hiddenSize.toLocaleString()}</td><td>${intermediateSize.toLocaleString()}</td><td>${(hiddenSize * intermediateSize).toLocaleString()}</td></tr>
+                <tr><td>LU W</td><td>${hiddenSize.toLocaleString()}</td><td>${intermediateSize.toLocaleString()}</td><td>${(hiddenSize * intermediateSize).toLocaleString()}</td></tr>
+                <tr><td>Proj W</td><td>${intermediateSize.toLocaleString()}</td><td>${hiddenSize.toLocaleString()}</td><td>${(intermediateSize * hiddenSize).toLocaleString()}</td></tr>
                 ${includeBias ? `<tr class="bias-row"><td>Bias</td><td>${switchB.toLocaleString()}</td><td>${luB.toLocaleString()}</td><td>${projB.toLocaleString()}</td><td></td></tr>` : ''}
             `;
 
             const embeddingTable = document.getElementById('embedding_table');
             embeddingTable.innerHTML = `
-                <tr><th>Type</th><th>Total</th></tr>
-                <tr><td>Embedding</td><td>${embeddingTotal.toLocaleString()}</td></tr>
+                <tr><th>Type</th><th>Input Size</th><th>Output Size</th><th>Total Size</th></tr>
+                <tr><td>Embedding</td><td>${vocabSize.toLocaleString()}</td><td>${hiddenSize.toLocaleString()}</td><td>${embeddingTotal.toLocaleString()}</td></tr>
             `;
 
             document.getElementById('one_layer_params').innerText = oneLayerParams.toLocaleString();
@@ -204,6 +211,10 @@
 <div class="output">
     <h3>Model Parameter Results</h3>
 
+    <div class="section-title">1 Layer Parameters</div>
+    <div><strong>1 Layer Parameters:</strong> <span id="one_layer_params">0</span></div>
+    <div><strong>Full Layers Parameters:</strong> <span id="full_layers_params">0</span></div>
+
     <div class="section-title">Attention Parameters</div>
     <table id="attention_table" class="sub-table"></table>
 
@@ -214,8 +225,6 @@
     <table id="embedding_table" class="sub-table"></table>
 
     <div class="highlight">
-        <strong>1 Layer Parameters:</strong> <span id="one_layer_params">0</span><br>
-        <strong>Full Layers Parameters:</strong> <span id="full_layers_params">0</span><br>
         <strong>Complete Model Size:</strong> <span id="full_size">0</span>
     </div>
 </div>