         }
         __syncthreads();
         if (threadIdx.x == 0) {                             // Save f to SHMEM
-          prods[threadIdx.y][j] = f;
-        }
-      }
-      __syncthreads();
-#pragma unroll
-      for (i = 1; i < NYDIM; i++) {                         // Reduce in SHMEM 
-        if (tid < NWINDOW) {
-          prods[0][tid] += prods[i][tid];
-        }
-        __syncthreads();
-      }
-      __syncthreads();
-#pragma unroll
-      for (j = 0; j < NWINDOW; j++) {     
-        f = prods[0][j];                                    // Get the value back to register memory
-
+          prods[threadIdx.y][jneg+NNEG*j] = f;
+        }
+      }
+    }
+    __syncthreads();
+#pragma unroll
+    for (i = 1; i < NYDIM; i++) {                         // Reduce in SHMEM 
+#pragma unroll
+      for (j = tid; j < NWINDOW * NNEG; j += dxy) {
+        prods[0][tid] += prods[i][tid];
+      }
+      __syncthreads();
+    }
+    for (j = tid; j < NWINDOW * NNEG; j += dxy) {
+      f = prods[0][tid];
         // Compute g from f
-        label = (jneg == 0);
-        if (f > 12.0f) {
-          g = 1.0f;
-        } else {
-          float expf = exp(f);
-          g = expf / (1.0f + expf);
-        }
-        g = (label - g) * lrate;
-
-        // Update the gradients
+      label = tid % NNEG;
+      if (f > 12.0f) {
+        g = 1.0f;
+      } else {
+        float expf = exp(f);
+        g = expf / (1.0f + expf);
+      }
+      g = (label - g) * lrate;
+      prods[0][tid] = g;
+    }
+    __syncthreads();
+
+#pragma unroll
+    for (j = 0; j < NWINDOW; j++) {     
+#pragma unroll
+      for (jneg = 0; jneg < NNEG; jneg++) {                   // Iterate over the negatives
+        g = prods[0][jneg + NNEG * i];
 #pragma unroll
         for (i = 0; i < NELTS; i++) {     
           daa[i][j] += g * bb[i];
           dbb[i] += g * aa[i][j];
         }
       }
+    }
+#pragma unroll
+    for (j = 0; j < NWINDOW; j++) {     
       __syncthreads();
 #pragma unroll
       for (i = 0; i < NELTS; i++) {                         // Save the update to the negative column


    // Get negative column indices
    __syncthreads();
    if (tid < NNEG) {                                  
      wb[tid] = WB[tid + icol * NNEG];
    }
    __syncthreads();
    // Iterate over the negatives
    for (jneg = 0; jneg < NNEG; jneg++) {                   // Iterate over the negatives
#pragma unroll
      for (i = 0; i < NELTS; i++) {                         // load the current negative column in NELTS sections
        if (tid + i*dxy < nrows) {
          bb[i] = B[tid + i*dxy + wb[jneg] * nrows]; 
        } else {
          bb[i] = 0;
        }
        dbb[i] = 0;
      }

      // Compute all the inner products with the current negative
#pragma unroll
      for (j = 0; j < NWINDOW; j++) {     
        f = 0;
#pragma unroll
        for (i = 0; i < NELTS; i++) {                       // load the current negative column in NELTS sections
          f += aa[i][j] * bb[i];                            // partial product
        }

        // This section reduces f over the column
#pragma unroll
        for (k = 1; k < 32; k = k+k) {                      // Reduce f in a warp

        aa[i][j] = aa[i][j+1];
        daa[i][j] = daa[i][j+1];
      }
    }
#pragma unroll
    for (i = 0; i < NELTS; i++) {                           // get the column data in NELTS sections
      if (good && tid + i*dxy < nrows) {
        aa[i][NWINDOW-1] = A[tid + i*dxy + wa * nrows];     // Load the new data
      } else {
