diff --git a/src/nnue/layers/affine_transform_sparse_input.h b/src/nnue/layers/affine_transform_sparse_input.h
index effda826b..472da834f 100644
--- a/src/nnue/layers/affine_transform_sparse_input.h
+++ b/src/nnue/layers/affine_transform_sparse_input.h
@@ -254,6 +254,7 @@ class AffineTransformSparseInput {
     #elif defined(USE_AVX2)
         using invec_t  = __m256i;
         using outvec_t = __m256i;
+        #define vec_add_32 _mm256_add_epi32
         #define vec_set_32 _mm256_set1_epi32
         #define vec_add_dpbusd_32 SIMD::m256_add_dpbusd_epi32
     #elif defined(USE_SSSE3)
@@ -272,21 +273,19 @@ class AffineTransformSparseInput {
         #define vec_set_32(a) vreinterpretq_s8_u32(vdupq_n_u32(a))
         #define vec_add_dpbusd_32 SIMD::neon_m128_add_dpbusd_epi32
     #endif
-        static constexpr IndexType OutputSimdWidth = sizeof(outvec_t) / sizeof(OutputType);
-
+        constexpr IndexType OutputSimdWidth = sizeof(outvec_t) / sizeof(OutputType);
         constexpr IndexType NumChunks = ceil_to_multiple<IndexType>(InputDimensions, 8) / ChunkSize;
         constexpr IndexType NumAccums = OutputDimensions / OutputSimdWidth;
-        // If there's only one accumulator and we're using high-latency dot product instructions,
-        // split it to create three separate dependency chains and merge at the end
-        constexpr bool SplitAccums =
+        // If we're using high-latency dot product instructions, split the accumulators
+        // to create 3 separate dependency chains and merge at the end
+        constexpr IndexType NumRegs =
     #if defined(USE_VNNI)
-          NumAccums == 1;
+          3 * NumAccums;
     #else
-          false;
+          NumAccums;
     #endif
-        constexpr IndexType NumRegs = SplitAccums ? 3 * NumAccums : NumAccums;
-        std::uint16_t       nnz[NumChunks];
-        IndexType           count;
+        std::uint16_t nnz[NumChunks];
+        IndexType     count;
 
         const auto input32 = reinterpret_cast<const std::int32_t*>(input);
 
@@ -303,30 +302,34 @@ class AffineTransformSparseInput {
 
         // convince GCC to not do weird pointer arithmetic in the following loop
         const std::int8_t* weights_cp = weights;
+    #if defined(USE_VNNI)
+        for (IndexType k = NumAccums; k < NumRegs; ++k)
+            acc[k] = vec_zero();
 
-        if constexpr (SplitAccums)
+        while (start < end - 2)
         {
-            acc[1] = acc[2] = vec_set_32(0);
-            while (start < end - 2)
+            const std::ptrdiff_t i0  = *start++;
+            const std::ptrdiff_t i1  = *start++;
+            const std::ptrdiff_t i2  = *start++;
+            const invec_t        in0 = vec_set_32(input32[i0]);
+            const invec_t        in1 = vec_set_32(input32[i1]);
+            const invec_t        in2 = vec_set_32(input32[i2]);
+            const auto           col0 =
+              reinterpret_cast<const invec_t*>(&weights_cp[i0 * OutputDimensions * ChunkSize]);
+            const auto col1 =
+              reinterpret_cast<const invec_t*>(&weights_cp[i1 * OutputDimensions * ChunkSize]);
+            const auto col2 =
+              reinterpret_cast<const invec_t*>(&weights_cp[i2 * OutputDimensions * ChunkSize]);
+            for (IndexType k = 0; k < NumAccums; ++k)
             {
-                const std::ptrdiff_t i0  = *start++;
-                const std::ptrdiff_t i1  = *start++;
-                const std::ptrdiff_t i2  = *start++;
-                const invec_t        in0 = vec_set_32(input32[i0]);
-                const invec_t        in1 = vec_set_32(input32[i1]);
-                const invec_t        in2 = vec_set_32(input32[i2]);
-                const auto           col0 =
-                  reinterpret_cast<const invec_t*>(&weights_cp[i0 * OutputDimensions * ChunkSize]);
-                const auto col1 =
-                  reinterpret_cast<const invec_t*>(&weights_cp[i1 * OutputDimensions * ChunkSize]);
-                const auto col2 =
-                  reinterpret_cast<const invec_t*>(&weights_cp[i2 * OutputDimensions * ChunkSize]);
-                vec_add_dpbusd_32(acc[0], in0, *col0);
-                vec_add_dpbusd_32(acc[1], in1, *col1);
-                vec_add_dpbusd_32(acc[2], in2, *col2);
+                vec_add_dpbusd_32(acc[k], in0, col0[k]);
+                vec_add_dpbusd_32(acc[k + NumAccums], in1, col1[k]);
+                vec_add_dpbusd_32(acc[k + 2 * NumAccums], in2, col2[k]);
             }
-            acc[0] = vec_add_32(vec_add_32(acc[0], acc[1]), acc[2]);
         }
+        for (IndexType k = 0; k < NumAccums; ++k)
+            acc[k] = vec_add_32(vec_add_32(acc[k], acc[k + NumAccums]), acc[k + 2 * NumAccums]);
+    #endif
         while (start < end)
         {
             const std::ptrdiff_t i  = *start++;
@@ -340,8 +343,12 @@ class AffineTransformSparseInput {
         outvec_t* outptr = reinterpret_cast<outvec_t*>(output);
         for (IndexType k = 0; k < NumAccums; ++k)
             outptr[k] = acc[k];
+
     #undef vec_set_32
     #undef vec_add_dpbusd_32
+    #ifdef vec_add_32
+        #undef vec_add_32
+    #endif
 #else
         // Use dense implementation for the other architectures.
         affine_transform_non_ssse3<InputDimensions, PaddedInputDimensions, OutputDimensions>(