Body text

Pair frequencies

template<typename T>
class PairCount {
public:
    // Virtual destructor to ensure proper cleanup of derived classes.
    virtual ~PairCount() = default;

    // Gets the total number of unique pairs stored.
    virtual size_t get_count() = 0;

    // Retrieves the count for a specific pair.
    virtual optional<int> get_pair(pair<T,T> mp) = 0;

    // Creates a new pair or modifies the frequency of an existing one.
    virtual bool create_or_modify_pair(T a, T b, int freq) = 0;

    // Gets the pair with the highest count.
    virtual optional<pair<T,T>> get_top_pair_count() = 0;

    // Retrieves all pairs and their counts.
    virtual std::vector<std::vector<T>> get_all() = 0;
};

• Conflict resolution strategy: First seen in input
  

  Imagine a sequence as follows:

  1,2,8,9,3,4…

  After counting all the pairs we find that [1,2] and [3,4] have the same frequency.

  1. [1,2] => 20 
  2. [3,4] => 20

  In this case we pick the one added first, which means the one first seen in the input text.

  In Python this insertion order comes for free because of Raymond Hettinger's 2012 redesign of the Python dictionary. Implemented in Python 3.6 (released December 23, 2016), introduced compact dictionaries with key-sharing and faster performance. A side effect of this redesign was that dictionaries began preserving insertion order as an implementation detail. This was later formalized as a language guarantee in Python 3.7 (released June 27, 2018), meaning dictionaries officially maintain the order of key-value pairs as they are inserted.

  In Karpathy's code you can see that he simply relies on this behaviour to get the consistent result based on above.

  # count up the number of times every consecutive pair appears
  stats = get_stats(ids)
  # find the pair with the highest count
  pair = max(stats, key=stats.get)
  

  And from the Python documentation: https://docs.python.org/3/library/functions.html#max

  If multiple items are maximal, the function returns the first one encountered. This is consistent with other sort-stability preserving tools such as sorted(iterable, key=keyfunc, reverse=True)[0] and heapq.nlargest(1, iterable, key=keyfunc).

  In order to implement that we must track the insertion order. Rather than let the user deal with that I built it into the PairCount class. As elements are added, new ones get the current count and the count is incremented.

  Picking a data structure here is tricky because we want to be able to quickly store and modify pair frequencies (unordered_map), and a way to get the most frequent (priority_queue). Furthermore, we want to keep track of insertion order?

  Sometimes you need to use multiple data structures to support a use case with conflicting requirements. For this purpose I used the boost::multi_index.

  https://www.boost.org/doc/libs/1_88_0/libs/multi_index/doc/index.html

  There's nothing to stop you from using a set and a priority queue and tracking them yourself, but multi_index handles that for you based on the declaration of which indexes and access patterns you need.

  Let's take a look at the implementation of PairCountInsertOrder:

  First the data; we need to store pair, the count and the insert order.

  template<typename T>
  struct PairCountOrder {
      ::pair<T,T> pair;
      int count;
      size_t insert_order;
  
      PairCountOrder(::pair<T,T> p, int c, size_t fo) : pair(p), count(c), insert_order(fo) {}
      PairCountOrder(::pair<T,T> p, int c) : pair(p), count(c), insert_order(std::numeric_limits<size_t>::max()) {}
  };
  
  // Comparison struct for sorting. Sorts by count (descending), then by insertion order (ascending).
  template<typename T>
  struct CompareCountOrder {
      bool operator()(const PairCountOrder<T>& a, const PairCountOrder<T>& b) const {
          if(a.count == b.count) {
              return a.insert_order < b.insert_order;
          } else {
              return a.count > b.count; // higher count is greater
          }
      }
  };
  

  Next we define the container itself. We just specify the indexes required and Boost takes care of picking the underlying data structures.

  template<typename T>
  using PairCountStore = boost::multi_index_container<
      PairCountOrder<T>,
      indexed_by<
          // Index 0: Hashed unique index on the 'pair' member for fast lookups.
          hashed_unique<member<PairCountOrder<T>, pair<T,T>, &PairCountOrder<T>::pair>>,
          // Index 1: Ordered non-unique index for sorting by count and insertion order.
          ordered_non_unique<identity<PairCountOrder<T>>, CompareCountOrder<T>>
      >
  >;
  

  Index 0 explanation: It is hashed so we should get an O(1) lookup type, and unique meaning keys are unique, each pair can occur once only. The rest of the declaration explains how to get the key for this index (use the pair member).

  Index 1 explanation:  This needs to be an ordered collection so we can extract the highest frequency. It also needs to be non-unique (in its sort criteria), because we can have multiple elements with the same frequency.

  Now in our code we can grab the appropriate index depending on the current purpose and when we make modifications to the data the boost library will ensure the changes are synchronized across all the indexes in the container.

  auto& index_by_key = pcs.template get<0>();
  auto f = index_by_key.find(mp);
  if(f != pcs.end()) {
      index_by_key.modify(f, [freq](PairCountOrder<T>& pc) { pc.count += freq; });
      return false;
  } else {
      pcs.insert(PairCountOrder<T>(mp, freq, next_insert++));
      return true;
  }
  

• Conflict resolution strategy: Lexicographical
  

  Referred to as lexical in my implementation to save typing, this method means we pick from pairs based on which comes first. For example given the following two pairs:

  1. [1,2] => 20 
  2. [3,4] => 20

  They have the same frequency so we pick pair 1) as 1 < 3. The second member of the pair is used as the tie-breaker, and of course if both members are the same then they would be combined to a single entry in the PairCount.

  Again a multi_index container is needed here. Let's start with the data:

  template<typename T>
  struct PairCountLexical {
      ::pair<T,T> pair;
      int count;
  
      PairCountLexical(::pair<T,T> p, int c) : pair(p), count(c) {}
  };
  
  // Comparison struct for sorting. Sorts by count (descending), then by pair (lexical ascending).
  template<typename T>
  struct CompareLexicalOrder {
      bool operator()(const PairCountLexical<T>& a, const PairCountLexical<T>& b) const {
          if(a.count == b.count) {
              if (a.pair.first == b.pair.first) {
                  return a.pair.second < b.pair.second;
              } else {
                  return a.pair.first < b.pair.first;
              }
          } else {
              return a.count > b.count; // higher count is greater
          }
      }
  };
  

  And the container looks like this:

  template<typename T>
  using PairCountLexicalStore = boost::multi_index_container<
      PairCountLexical<T>,
      indexed_by<
          // Index 0: Hashed unique index on the 'pair' member for fast lookups.
          hashed_unique<member<PairCountLexical<T>, pair<T,T>, &PairCountLexical<T>::pair>>,
          // Index 1: Ordered non-unique index for sorting by count and lexical order.
          ordered_non_unique<identity<PairCountLexical<T>>, CompareLexicalOrder<T>>
      >
  >;
  

  Index 0 explanation: Same as above this gives us fast insert, modify and lookup for the pair frequencies.

  Index 1 explanation: Same as above except the outcome is different because of the implementation of CompareLexicalOrder.

• Optimization of frequency counts
  

  When running the code I see that the biggest cost is regenerating the frequency map each step. For example when churning through wikitext (a 500mb text corpus) it takes the Python code 28 seconds on my Macbook to count all the pairs.

  Let's work through Brendan Gregg's impactful optimizations:

  1. Don’t do it.
  2. Do it, but don’t do it again.
  3. Do it less.
  4. Do it later.
  5. Do it when they’re not looking.
  6. Do it concurrently.
  7. Do it more cheaply.

  Don't do it is not an option, we need those updated counts each step.  Do it but not again is fruitful though. 

  The key insight here is that we only need to do a full frequency count one time. Then we can incrementally update the pair frequencies as we walk through doing the merge process. Essentially we are removing and adding a number of pairs on each replacement.

  I noticed that the authors of the paper mentioned this too: 

  "In practice, we increase efficiency by indexing all pairs, and updating data structures incrementally."

  You can see their incremental update code here:

  https://github.com/rsennrich/subword-nmt/blob/92d6139d07d30e12735a0af9e7f7f925ebe62c54/subword_nmt/learn_bpe.py#L159

  In addition to this optimization they use a pruning technique that drops frequencies of pairs below some threshold. This makes sense because the Python max function iterates the whole collection. In my case our data structures do not, so pruning is probably not worth the additional complexity. Worth trying maybe?

  In any case, for my lexicographical conflict strategy I do implement this optimization and it is a huge win on performance as shown in the charts.

  Crucially, it is not implemented for the first occurring strategy, because the current implementation gives now way to easily keep track of the first occurence of a pair in the input corpus.