以前書いた記事からしばらく実験を続けてアーキテクチャの改善があったので報告です。
レポジトリと学習済み重み（事前学習と対話ファインチューニング）を公開しています。
現状のアーキテクチャは下図（チャンクサイズ4の例）のブロックを積み上げた構造になっています：
トークン列をチャンク毎に圧縮して再帰的に処理するという基本的な構造は変わりません。推論時は直前チャンクに関係する情報のみを記憶することで空間計算量コンテキスト長nに対してO(log n)に抑えられます。
図の入力の高階クエリは2段目以降は前段の出力を使用しますが、初段のみ直前レベルチャンクの最終クエリを使用するようにします。図の例だと下側左から4本目の黒線入力が下側一番左の赤線入力に使用されます。
大きな変更点としてはFFN層を再帰構造から分離したことです。以前のモデルでは日英コーパスで学習したモデルが長期推論時に日本語からいきなり英語に切り替わるといった不安定性が見られましたが、現在のモデルの推論では同じ文字のループに陥らなければ数万トークン程度生成しても安定して日本語を出力するようになりました。
対話ファインチューニングを行ったモデルについても、以前のモデルだと先頭の質問に執着して話題を切り替えても同じような回答をする傾向がありましたが、現在のモデルだと話題の変更に追従できているようです。
上記事前学習モデル及び対話学習モデルはチャンクサイズ4で学習しました。私の計算資源でできたコンテキスト長2048の学習だと深さ方向の汎化（再帰深さ$\log_4 2048=5.5$）と計算コストのトレードオフでチャンクサイズ4が適切だと思っていますが、最小チャンクサイズ2で学習が可能か、現在訓練を回しているところです。
試しに学習途中のモデルに出力させてみたのが以下で、内容はともかく日本語文法は問題なく学習できているようでした。
>>> 日本の首都
日本の首都として知られる東京は、人口400万人当たりの都市圏に次ぐ人口の減少が懸念され、経済が逼迫した中、経済発展が進む中、若者の就業率は急速に落ち込みます。
その一方で、日本には、先進国において10年以上にわたる日本経済の激変が起きています。日本で人口の急増が予想されることや、急速に増加しているのが大都市であること、そして、新興国の人口増加が挙げられます。
地方都市に暮らす人々にとって、生活や雇用といった要素が非常に大きなインパクトを持つことで、人口の増加は、ますますます加速化していきます。
特に、高度成長期以降、人口の増加や地方都市への転換は、若者の生活を支援する大きな課題の一つになっています。
また、地域の文化やスポーツ振興に寄与する公共交通網が整備されていることも事実です。
このような状況の中で、自らの資産を蓄積し、地域社会の発展に貢献するためには、地域の活性化や人口減少が起きることによって、地域の経済発展を支えることになります。
しかし、現代においては、地域の持続的な発展、発展のためには、地域の経済・社会・文化の発展と相互に連携し、地域の自立を図るために、地域の課題解決を図り、持続可能な社会の実現を目指すものです。
[350 tokens, 25.55s, 13.70 tok/s]