算術符号よりめっちゃ高速だとの噂の「非対称符号系(ANS)」のコードを JavaScript で書いて、速度を測ってみた

注意：エントロピー符号の話なので、かなりひとを選ぶ記事となっています。

Zstandard というデータ圧縮アルゴリズムが高速だとの噂で、採用している符号化アルゴリズムが「tANS」だというので、どのようなアルゴリズムなのか気になって調べてみた。
調べていくうちに「tANS」のようなテーブルを作るまでもなく、テーブル無しの「rANS」でも十分早いんじゃないか？と思えたので実際に動くかどうかコードを書いてみた。
結果として、計算コストを必要とする割り算や掛け算の数も少ない、短いコードになった。動かしてみて、やはり、問題ないと思えるくらいに高速に思えたので、ここで共有します。

ただ、算術符号とかレンジ符号の実行速度と直接おなじ環境で比較してみないと、なんとも言えないので、比較可能な JavaScript なコードがありましたら教えていただけると幸いです。

非対称符号系ANSに関しては、こちらの英語版 wiki がまとまっていると思われます。

rANS encode を JavaScript で書いてみた

具体的には、文字列を Uint8Array に変換（符号化、圧縮）しています。関数は文字列を入力し Object を返してます。Object には Uint8Array の他に頻度情報のテーブルが付いています。この頻度情報が無いと復号 decode できません。

const rANS_String_to_ByteArray_encode = _str =>
{
    const str_table  = _str.split('')
    const str_length = str_table.length

    const out = {
        frequency:          {},         // ex. { "a":5, "b":2, "d":3, "c":1 },
        input_str_length:   str_length,
        byte_array:         new Uint8Array( str_length ), // 出力サイズは未知なので大きめに確保している（そんなに使わない）
        output_length:      0,          // encode Byte
        lest_state:         0,          // decode時の開始 state
    }

    // 頻度を求める（高速化のために、個数を 2の累乗 に揃える）揃えない場合は、⭐️の位置の計算で、シフト演算は使えずに、割り算を使うことになる
    const freq_n = str_length.toString(2).length
    const max = 2 ** freq_n             // 2**12 = 4096
    for ( let i=0 ; i<max ; )
    {
        for ( const s of str_table )
        {
            if ( s in out.frequency ) out.frequency[s] ++
            else                      out.frequency[s] = 1

            if ( ++i >= max ) break
        }
    }

    const freq_symbols       = Object.keys( out.frequency )                 // ex. [ "a", "b", "d", "c" ],
    const freq_symbol_counts = freq_symbols.map( s => out.frequency[s] )    // ex. [ 5, 2, 3, 1 ],

    // cumulative frequencies
    const freq_cumul_counts = [ 0 ]     // ex. [ 0, 5, 7, 10, 11 ],
    let sum = 0
    for ( const c of freq_symbol_counts )
    {
        sum += c
        freq_cumul_counts.push( sum )
    }

    // encode
    const index_table = str_table.map( c => freq_symbols.findIndex( s => s == c ) )
    let state = 2 ** freq_n
    for ( let i=str_length-1 ; i >= 0 ; i-- ) // 文字列の後ろから前へ encode
    {
        const s  = index_table[i]
        const Fs = freq_symbol_counts[s]
        const Cs = freq_cumul_counts[s]

        while ( state >= (Fs << 8) )
        {
            out.byte_array[ out.output_length++ ] = state & 255
            state >>>= 8
        }

        state = ( Math.floor( state / Fs ) << freq_n ) + Cs + ( state % Fs )    // ⭐️
    }

    out.last_state = state

    return out
}

前半は頻度情報を計算している。後半（// encodeから後ろの16行程度で）、頻度情報を元に符号化している。
たったこれだけのコードで符号化できるなんて、、、。すごいです。もう、算術符号も、ハフマン符号も、誰も使わなくなるのでは？と心配になるほどです。

ポイントとしては、freq_n を算出しているところです。頻度の合計が 2の累乗になるように工夫してます。この freq_n を使うことで、割り算と余りの計算をビットシフトとビットアンドに置き換えることができ、encode/decode ともに処理の高速化に貢献しています。

（freq_symbol_counts の各値も 2の累乗に揃えることができると、さらなる高速化（割り算と余りの計算をビットシフトとビットアンドに置き換え）も可能と思えるが、、、若干圧縮率が犠牲になるし、どうやって計算したら両方を2の乗数に揃えられるのか思いつかないので保留中です。）

rANS decode を JavaScript で書いてみた

上の encoder で圧縮された object （頻度情報＋バイナリデータ　Uint8Array）を文字の配列に復号化します。

const rANS_ByteArray_to_StrArray_decode = _enc_obj =>
{
    const freq_symbols       = Object.keys( _enc_obj.frequency )   // ex. [ "a", "b", "d", "c" ],
    const freq_symbol_counts = freq_symbols.map( s => _enc_obj.frequency[s] )  // ex. // [ 5, 2, 3, 1 ],

    // cumulative frequencies
    const freq_cumul_counts = [ 0 ]     // ex. [ 0, 5, 7, 10, 11 ],
    let sum = 0
    for ( const c of freq_symbol_counts )
    {
        sum += c
        freq_cumul_counts.push( sum )
    }
    const freq_n = _enc_obj.input_str_length.toString(2).length

    // decode
    const max_length = _enc_obj.input_str_length
    const out        = Array( max_length )      // decode 結果

    let state = _enc_obj.last_state
    let pos   = _enc_obj.output_length - 1  // 後ろから前へ decode する
    const bit_mask  = (2 ** freq_n) - 1
    for ( let i = 0; i < max_length; i++)
    {
        const p  = state & bit_mask                                 // ⭐️
        const s  = freq_cumul_counts.findIndex( c => p < c ) - 1
        const Fs = freq_symbol_counts[s]
        const Cs = freq_cumul_counts[s]

        state = ( state >>> freq_n ) * Fs + p - Cs                  // ⭐️

        while ( state <= bit_mask )
        {
            state <<= 8
            state += _enc_obj.byte_array[ pos-- ]
        }

        out[i] = freq_symbols[s]  // 文字列連結すると時間がかかるので配列に格納
    }

    return out
}

// decode以下のやく19行で復号化しています。このコードの中で、コストのかかる割り算は無く、掛け算も１つだけですし、ループもメインの他に短いループが２つだけという、いや、もう、遅くなる要素が皆無だったりします。

実行速度

以下、CodePen で実際に測れます。

テキストエリアに試したい文字列を入力して、（試してみたい文字列が思いつかない場合は、「Create test data」のボタンを押すことで、1024*1024 文字のDNAふう文字列をランダム（正規化）生成します。）

「encoder」ボタンを押すと、１０回符号化して、その平均時間を表示します。

「decoder」ボタンを押すと、１０回復号化して、その平均時間を表示します。

手元の環境での記録

手元の環境での実測時は、encode が 200ms。decode が 30ms でした。👇

動作確認環境は 2013年ころに買った MacBookAir で、ブラウザは Chrome バージョン: 94.0.4606.71 です。参考まで。

tANS のようにテーブル引きにせずとも rANS でも十分に高速ではないか？と思った理由

rANS decode を手元の環境でプロファイルして見てみると、174: findIndex のところが 192.5 ms と若干ボトルネックになっていることが確認できる。

で、この部分を tANS のようにテーブルから引いたら早くなるのでは？と推測できるので、試してみると、

まず、テーブルを作るだけ(208行の Array() )で 210ms と配列の確保だけですでにコストオーバー。さらに、テーブルを引く部分(224行目)も 671.7 ms と異常に遅い。
テーブルにしても高速化に繋がらないと感じたのはこのあたりの実測値が理由です。おそらく、テーブルのサイズが大きいのが原因であると推測できるので。テーブルを小さくすることが必要なのだが、それには頻度情報を小さくするする必要があり、それはそのまま符号化に響くように思ったので、保留中。

なにか、よいアイデアがありましたら、募集中です。

この記事が、みなさまのプログラミングライフの一助になれば幸いです。

参考：

Zstandard

Asynmmetric numeral system(ANS)の紹介

What is Asymmetric Numeral Systems?