Modellierung extrem großer Bilder mit xT – The Berkeley Artificial Intelligence Research Blog

Als Computer-Vision-Forscher glauben wir, dass jedes Pixel eine Geschichte erzählen kann. Es scheint jedoch eine Schreibblockade zu geben, wenn es um den Umgang mit großen Bildern geht. Große Bilder sind keine Seltenheit mehr – die Kameras, die wir in unseren Taschen tragen und die unseren Planeten umkreisen, machen so große und detaillierte Bilder, dass sie unsere derzeit besten Modelle und Hardware bei der Handhabung bis an ihre Grenzen belasten. Generell sehen wir einen quadratischen Anstieg der Speichernutzung als Funktion der Bildgröße.

Heutzutage treffen wir beim Umgang mit großen Bildern eine von zwei suboptimalen Entscheidungen: Downsampling oder Zuschneiden. Diese beiden Methoden führen zu erheblichen Verlusten an Informationen und Kontext in einem Bild. Wir werfen einen weiteren Blick auf diese Ansätze und stellen $x$T vor, ein neues Framework, um große Bilder durchgängig auf modernen GPUs zu modellieren und gleichzeitig den globalen Kontext effektiv mit lokalen Details zu aggregieren.

Architektur für das $x$T-Framework.

Warum sich überhaupt mit großen Bildern beschäftigen?

Warum sollte man sich überhaupt die Mühe machen, mit großen Bildern umzugehen? Stellen Sie sich vor, Sie sitzen vor Ihrem Fernseher und schauen sich Ihre Lieblingsfußballmannschaft an. Das Spielfeld ist übersät mit Spielern, und die Action findet jeweils nur auf einem kleinen Teil des Bildschirms statt. Wären Sie jedoch zufrieden, wenn Sie nur einen kleinen Bereich um die Stelle sehen könnten, an der sich der Ball gerade befindet? Wären Sie alternativ auch damit zufrieden, das Spiel in niedriger Auflösung anzusehen? Jedes Pixel erzählt eine Geschichte, egal wie weit sie voneinander entfernt sind. Dies gilt in allen Bereichen, von Ihrem Fernsehbildschirm bis hin zu einem Pathologen, der ein Gigapixel-Objektiv betrachtet, um winzige Krebsherde zu diagnostizieren. Diese Bilder sind wahre Schatzkammern an Informationen. Wenn wir den Reichtum nicht vollständig erkunden können, weil unsere Werkzeuge die Karte nicht verarbeiten können, welchen Sinn hat es dann?

Sport macht Spaß, wenn man weiß, was los ist.

Genau darin liegt heute der Frust. Je größer das Bild, desto mehr müssen wir gleichzeitig herauszoomen, um das ganze Bild zu sehen, und umso mehr hineinzoomen, um die kleinsten Details zu erkennen. Das macht es zu einer Herausforderung, sowohl den Wald als auch die Bäume gleichzeitig zu erfassen. Die meisten aktuellen Methoden erfordern die Wahl, den Wald aus den Augen zu verlieren oder die Bäume zu übersehen, und keine der beiden Optionen ist großartig.

Wie $x$T versucht, dieses Problem zu beheben

Stellen Sie sich vor, Sie versuchen, ein riesiges Puzzle zu lösen. Anstatt die ganze Sache auf einmal in Angriff zu nehmen, was überwältigend wäre, beginnen Sie mit kleineren Abschnitten, schauen sich die einzelnen Teile genau an und finden dann heraus, wie sie in das Gesamtbild passen. Das ist im Grunde das, was wir mit großen Bildern mit $x$T machen.

$x$T nimmt diese riesigen Bilder und zerlegt sie hierarchisch in kleinere, besser verdauliche Stücke. Dabei geht es jedoch nicht nur darum, die Dinge kleiner zu machen. Es geht darum, jedes Teil für sich zu verstehen und dann mithilfe einiger cleverer Techniken herauszufinden, wie diese Teile im größeren Maßstab zusammenhängen. Es ist, als würde man mit jedem Teil des Bildes ein Gespräch führen, seine Geschichte erfahren und diese Geschichten dann mit den anderen Teilen teilen, um die vollständige Erzählung zu erhalten.

Verschachtelte Tokenisierung

Im Mittelpunkt von $x$T steht das Konzept der verschachtelten Tokenisierung. Vereinfacht ausgedrückt ist die Tokenisierung im Bereich der Computer Vision so, als würde man ein Bild in Stücke (Tokens) zerlegen, die ein Modell verdauen und analysieren kann. $x$T geht jedoch noch einen Schritt weiter, indem es eine Hierarchie in den Prozess einführt – daher verschachtelt.

Stellen Sie sich vor, Sie hätten die Aufgabe, einen detaillierten Stadtplan zu analysieren. Anstatt zu versuchen, die gesamte Karte auf einmal zu erfassen, unterteilen Sie sie in Stadtteile, dann in Stadtteile innerhalb dieser Stadtteile und schließlich in Straßen innerhalb dieser Stadtteile. Diese hierarchische Aufteilung macht es einfacher, die Details der Karte zu verwalten und zu verstehen und gleichzeitig den Überblick zu behalten, wo alles in das Gesamtbild passt. Das ist die Essenz der verschachtelten Tokenisierung – wir teilen ein Bild in Regionen auf, die jeweils in weitere Unterregionen aufgeteilt werden können, abhängig von der Eingabegröße, die von einem Vision-Backbone (was wir a nennen) erwartet wird Regions-Encoder), bevor sie gepatcht werden, um von diesem Regions-Encoder verarbeitet zu werden. Dieser verschachtelte Ansatz ermöglicht es uns, Features in verschiedenen Maßstäben auf lokaler Ebene zu extrahieren.

Koordinierende Regions- und Kontext-Encoder

Sobald ein Bild sauber in Token unterteilt ist, verwendet $x$T zwei Arten von Encodern, um diesen Teilen einen Sinn zu geben: den Regions-Encoder und den Kontext-Encoder. Jedes spielt eine besondere Rolle bei der Zusammenstellung der Gesamtgeschichte des Bildes.

Der Region Encoder ist ein eigenständiger „lokaler Experte“, der unabhängige Regionen in detaillierte Darstellungen umwandelt. Da jedoch jede Region isoliert verarbeitet wird, werden keine Informationen über das gesamte Bild verteilt. Der Regionsencoder kann ein beliebiges hochmodernes Vision-Backbone sein. In unseren Experimenten haben wir hierarchische Vision-Transformatoren wie Swin und Hiera sowie CNNs wie ConvNeXt verwendet!

Betreten Sie den Kontext-Encoder, den Guru für das große Ganze. Seine Aufgabe besteht darin, die detaillierten Darstellungen der Regionsencoder zu übernehmen und sie zusammenzufügen, um sicherzustellen, dass die Erkenntnisse aus einem Token im Kontext der anderen betrachtet werden. Der Kontextencoder ist im Allgemeinen ein Langsequenzmodell. Wir experimentieren mit Transformer-XL (und unserer Variante davon namens Hyper) und Mamba, obwohl Sie Longformer und andere neue Fortschritte in diesem Bereich verwenden könnten. Obwohl diese Langsequenzmodelle im Allgemeinen für die Sprache konzipiert sind, zeigen wir, dass es möglich ist, sie effektiv für Sehaufgaben zu verwenden.

Die Magie von $x$T liegt darin, wie diese Komponenten – die verschachtelte Tokenisierung, Regions-Encoder und Kontext-Encoder – zusammenkommen. Indem $x$T das Bild zunächst in überschaubare Teile zerlegt und diese Teile dann sowohl einzeln als auch in Verbindung systematisch analysiert, gelingt es ihm, die Detailtreue des Originalbilds beizubehalten und gleichzeitig den Fernkontext in den übergreifenden Kontext zu integrieren während gleichzeitig riesige Bilder durchgängig auf modernen GPUs angepasst werden.

Ergebnisse

Wir bewerten $x$T anhand anspruchsvoller Benchmark-Aufgaben, die von etablierten Computer-Vision-Basislinien bis hin zu anspruchsvollen Großbildaufgaben reichen. Insbesondere experimentieren wir mit iNaturalist 2018 zur feinkörnigen Artenklassifizierung, xView3-SAR zur kontextabhängigen Segmentierung und MS-COCO zur Erkennung.

Leistungsstarke Vision-Modelle, die mit $x$T verwendet werden, setzen neue Maßstäbe für nachgelagerte Aufgaben wie die feinkörnige Artenklassifizierung.

Unsere Experimente zeigen, dass $x$T bei allen Downstream-Aufgaben mit weniger Parametern eine höhere Genauigkeit erreichen kann und gleichzeitig viel weniger Speicher pro Region verbraucht als modernste Baselines^*. Wir sind in der Lage, Bilder mit einer Größe von bis zu 29.000 x 25.000 Pixeln auf 40-GB-A100-Modellen zu modellieren, während vergleichbare Basismodelle bereits bei 2.800 x 2.800 Pixeln nicht mehr genügend Speicher haben.

Leistungsstarke Vision-Modelle, die mit $x$T verwendet werden, setzen neue Maßstäbe für nachgelagerte Aufgaben wie die feinkörnige Artenklassifizierung.

^*Abhängig von Ihrer Wahl des Kontextmodells, z. B. Transformer-XL.

Warum das wichtiger ist, als Sie denken

Dieser Ansatz ist nicht nur cool; es ist notwendig. Für Wissenschaftler, die den Klimawandel verfolgen, oder Ärzte, die Krankheiten diagnostizieren, ist dies ein Wendepunkt. Es bedeutet, Modelle zu schaffen, die die gesamte Geschichte verstehen und nicht nur Teile davon. Bei der Umweltüberwachung beispielsweise kann die Fähigkeit, sowohl die umfassenderen Veränderungen in riesigen Landschaften als auch die Details bestimmter Gebiete zu erkennen, dabei helfen, das Gesamtbild der Klimaauswirkungen zu verstehen. Im Gesundheitswesen kann es den Unterschied ausmachen, ob man sich früh mit einer Krankheit infiziert oder nicht.

Wir behaupten nicht, alle Probleme der Welt auf einen Schlag gelöst zu haben. Wir hoffen, dass wir mit $x$T die Tür zu dem geöffnet haben, was möglich ist. Wir treten in eine neue Ära ein, in der wir keine Kompromisse bei der Klarheit oder Breite unserer Vision eingehen müssen. $x$T ist unser großer Schritt hin zu Modellen, die die Feinheiten großformatiger Bilder bewältigen können, ohne ins Schwitzen zu geraten.

Es gibt noch viel mehr zu tun. Die Forschung wird sich weiterentwickeln, und hoffentlich auch unsere Fähigkeit, noch größere und komplexere Bilder zu verarbeiten. Tatsächlich arbeiten wir an Nachfolgern von $x$T, die diese Grenze weiter erweitern werden.

Abschließend

Eine vollständige Abhandlung dieser Arbeit finden Sie im Artikel auf arXiv. Die Projektseite enthält einen Link zu unserem veröffentlichten Code und unseren Gewichten. Wenn Sie die Arbeit nützlich finden, zitieren Sie sie bitte wie folgt:

@article{xTLargeImageModeling,
  title={xT: Nested Tokenization for Larger Context in Large Images},
  author={Gupta, Ritwik and Li, Shufan and Zhu, Tyler and Malik, Jitendra and Darrell, Trevor and Mangalam, Karttikeya},
  journal={arXiv preprint arXiv:2403.01915},
  year={2024}
}