Az Amazon Trainium chipeken gyorsítaná a valós idejű videógenerálást

A hosszú videókhoz más modellre van szükség

Az Amazon Science szeptember 25-i beszámolója szerint a világmodellek az elmúlt nyolc évben jelentősen fejlődtek. A nagy tanítási adathalmazok, a variációs autoenkóderek, valamint a diffúziós és autoregresszív transzformerek lehetővé teszik, hogy ezek a rendszerek szöveges, képi, hang- és videóbemenetekből látens reprezentációkat hozzanak létre.

A világmodelleket az Amazon szerint többek között robotikában, közlekedésben, klímamodellezésben, videojáték-fejlesztésben, tudományos szimulációkban és tervezési prototípusok készítésében használják. A valós idejű működéshez azonban olyan hardverre és infrastruktúrára van szükség, amely alacsony késleltetés mellett képes kezelni a modellek számítási és memóriaigényét.

A Reactor olyan platformot üzemeltet, amelyen fejlesztők, tervezők és kutatók valós idejű, interaktív AI-modelleket telepíthetnek, használhatnak és méretezhetnek. A cég társalapítója és technológiai igazgatója, Bryce Schmidtchen szerint a hatékonyság része az inferencia ütemezése az adott chipen, valamint az egyes modellfuttatások lehető legjobb kihasználása.

A Rolling Forcing lehet a valós idejű videó alapja

Az Amazon Neuron Science csapata a videógenerálás fejlődésében az autoregresszív diffúziós modellek megjelenésére figyelt fel. Ezek a rendszerek a nagy nyelvi modellekből ismert sorozatos következő elem előrejelzését egyesítik a diffúziós modellek iteratív finomításával.

Ez a megközelítés különösen az interaktív videók esetében fontos. A felhasználó egy billentyűleütése bemenetként szolgálhat, a modell pedig az előzőleg létrehozott képkocka alapján dönti el, mi legyen a következő mozdulat vagy jelenet.

Az Amazon és a Reactor több videógenerálási technikát vizsgált, végül a Rolling Forcing mellett döntött. A forrás szerint a modell képes folyamatosan jó minőségű, 30 másodperces videókat készíteni, miközben viszonylag kis méretű.

A valós idejű működés legnagyobb nehézsége, hogy a képkockák folyamatosan, streamelve készülnek. Minden új képkockát azonnal meg kell jeleníteni, vagy a következő bemenetként kell felhasználni. Ha egy képkocka késik, a generálás lemarad a lejátszás mögött. A Rolling Forcing másodpercenként 16 képkocka létrehozására képes, ami a forrás szerint megfelel a videólejátszás szabványának. Az Amazon szerint a Trainium teljesítménye és memóriája lehetővé teszi, hogy a modell ezt a hosszú sorozatokból adódó terhelést 16 képkocka per másodperc feletti sebességgel kezelje.

A fejlesztők közvetlenül a számítási műveleteket hangolták

A Neuron Science és a Reactor alulról felfelé építkező, kernelközpontú módszertant választott. Három problémára összpontosítottak: a dinamikus alakzatokra, a szokatlan memóriahozzáférési mintákra és a nagy mennyiségű gyorsítótár-kezelésre.

A valós idejű videógenerálás során a jelenetek változásával az egyes előrehaladási lépésekhez tartozó figyelmi szakaszok hossza is eltérhet. A képkockák egy folyamatos, csúszóablakos folyamatban készülnek és kerülnek ki a memóriából. Emiatt minden ablakhoz két külön művelet tartozik: a zaj eltávolítása, majd a gyorsítótár megtisztítása.

A videógenerálásban használt forgó pozícióbeágyazás, vagyis RoPE, három tengelyt kezel: a magasságot, a szélességet és az időt. Az Amazon szerint a háromdimenziós RoPE általános fordítás esetén sok apró adatátvitelt eredményezhet. A kulcsérték-gyorsítótár minden rétegben olvasási és írási műveleteket igényel, ami további terhelést ró az eszköz memóriájára.

A csapat ezért a Neuron Kernel Interface, röviden NKI használatához fordult. Ez lehetővé teszi, hogy a fejlesztők közvetlenül a NeuronCore hardveren futó számítási kerneleket írjanak, és pontosabban szabályozzák az adatok mozgását a memória és a számítási egységek között.

Az Amazon szerint jelentősen rövidültek egyes műveletek

Az Amazon beszámolója szerint a célzott kerneloptimalizálás több szűk keresztmetszetet is kezelt. A 3D RoPE művelet futási ideje öt másodpercről 1,8 milliszekundumra csökkent, a gyorsítótár másolása pedig rétegenként 23 milliszekundumról 1,9 milliszekundumra rövidült. A közlemény szerint a figyelmi transzponálásokat teljesen megszüntették az összevonással.

A fejlesztés jelentősége a Reactor által támogatott, régiókon átívelő valós idejű alkalmazásokban mutatkozhat meg. A vállalat több száz, világszerte elhelyezett GPU-fürtben gondolkodik, és olyan rendszert épít, amely a hálózati, médiafolyam-kezelési, kodek- és felbontásigényeket is összehangolja. Az együttműködés alapot teremthet ahhoz, hogy az autoregresszív diffúziós videómodellek Trainium chipeken is valós időben fussanak.