diff --git a/av1/common/entropy.h b/av1/common/entropy.h
index 7d0851a..33f287a 100644
--- a/av1/common/entropy.h
+++ b/av1/common/entropy.h
@@ -119,7 +119,12 @@
   // TODO(debargha): Does TX_64X64 require an additional extrabit?
   if (tx_size > TX_32X32) tx_size = TX_32X32;
 #endif
-  int bits = (int)bit_depth + 3 + (int)tx_size;
+#if CONFIG_CB4X4
+  int tx_offset = (tx_size < TX_4X4) ? 0 : (int)(tx_size - TX_4X4);
+#else
+  int tx_offset = (int)(tx_size - TX_4X4);
+#endif
+  int bits = (int)bit_depth + 3 + tx_offset;
   assert(bits <= (int)sizeof(av1_cat6_prob));
   return bits;
 }