In diesem Tutorial untersuchen wir, wie man mit JAX, Flax und Optax auf effiziente und modulare Weise ein fortschrittliches neuronales Netzwerk aufbaut und trainiert. Wir beginnen mit dem Entwurf einer tiefen Architektur, die Restverbindungen und Selbstaufmerksamkeitsmechanismen für das Lernen ausdrucksstarker Merkmale integriert. Im weiteren Verlauf implementieren wir ausgefeilte Optimierungsstrategien mit Lernratenplanung, Gradientenbeschneidung und adaptivem Gewichtsabfall. Während des gesamten Prozesses nutzen wir JAX-Transformationen wie jit, grad und vmap, um die Berechnung zu beschleunigen und eine reibungslose Trainingsleistung auf allen Geräten sicherzustellen. Schauen Sie sich das an VOLLSTÄNDIGE CODES hier.

!pip set up jax jaxlib flax optax matplotlib


import jax
import jax.numpy as jnp
from jax import random, jit, vmap, grad
import flax.linen as nn
from flax.coaching import train_state
import optax
import matplotlib.pyplot as plt
from typing import Any, Callable


print(f"JAX model: {jax.__version__}")
print(f"Gadgets: {jax.units()}")

Wir beginnen mit der Set up und dem Import von JAX, Flax und Optax sowie wichtigen Dienstprogrammen für numerische Operationen und Visualisierung. Wir überprüfen unser Geräte-Setup, um sicherzustellen, dass JAX auf der verfügbaren {Hardware} effizient läuft. Dieses Setup bildet die Grundlage für die gesamte Trainingspipeline. Schauen Sie sich das an VOLLSTÄNDIGE CODES hier.

class SelfAttention(nn.Module):
   num_heads: int
   dim: int
   @nn.compact
   def __call__(self, x):
       B, L, D = x.form
       head_dim = D // self.num_heads
       qkv = nn.Dense(3 * D)(x)
       qkv = qkv.reshape(B, L, 3, self.num_heads, head_dim)
       q, okay, v = jnp.break up(qkv, 3, axis=2)
       q, okay, v = q.squeeze(2), okay.squeeze(2), v.squeeze(2)
       attn_scores = jnp.einsum('bhqd,bhkd->bhqk', q, okay) / jnp.sqrt(head_dim)
       attn_weights = jax.nn.softmax(attn_scores, axis=-1)
       attn_output = jnp.einsum('bhqk,bhvd->bhqd', attn_weights, v)
       attn_output = attn_output.reshape(B, L, D)
       return nn.Dense(D)(attn_output)


class ResidualBlock(nn.Module):
   options: int
   @nn.compact
   def __call__(self, x, coaching: bool = True):
       residual = x
       x = nn.Conv(self.options, (3, 3), padding='SAME')(x)
       x = nn.BatchNorm(use_running_average=not coaching)(x)
       x = nn.relu(x)
       x = nn.Conv(self.options, (3, 3), padding='SAME')(x)
       x = nn.BatchNorm(use_running_average=not coaching)(x)
       if residual.form(-1) != self.options:
           residual = nn.Conv(self.options, (1, 1))(residual)
       return nn.relu(x + residual)


class AdvancedCNN(nn.Module):
   num_classes: int = 10
   @nn.compact
   def __call__(self, x, coaching: bool = True):
       x = nn.Conv(32, (3, 3), padding='SAME')(x)
       x = nn.relu(x)
       x = ResidualBlock(64)(x, coaching)
       x = ResidualBlock(64)(x, coaching)
       x = nn.max_pool(x, (2, 2), strides=(2, 2))
       x = ResidualBlock(128)(x, coaching)
       x = ResidualBlock(128)(x, coaching)
       x = jnp.imply(x, axis=(1, 2))
       x = x(:, None, :)
       x = SelfAttention(num_heads=4, dim=128)(x)
       x = x.squeeze(1)
       x = nn.Dense(256)(x)
       x = nn.relu(x)
       x = nn.Dropout(0.5, deterministic=not coaching)(x)
       x = nn.Dense(self.num_classes)(x)
       return x

Wir definieren ein tiefes neuronales Netzwerk, das Restblöcke und einen Selbstaufmerksamkeitsmechanismus für verbessertes Function-Lernen kombiniert. Wir bauen die Schichten modular auf und stellen so sicher, dass das Modell sowohl räumliche als auch kontextuelle Beziehungen erfassen kann. Dieses Design ermöglicht dem Netzwerk eine effektive Verallgemeinerung verschiedener Arten von Eingabedaten. Schauen Sie sich das an VOLLSTÄNDIGE CODES hier.

class TrainState(train_state.TrainState):
   batch_stats: Any


def create_learning_rate_schedule(base_lr: float = 1e-3, warmup_steps: int = 100, decay_steps: int = 1000) -> optax.Schedule:
   warmup_fn = optax.linear_schedule(init_value=0.0, end_value=base_lr, transition_steps=warmup_steps)
   decay_fn = optax.cosine_decay_schedule(init_value=base_lr, decay_steps=decay_steps, alpha=0.1)
   return optax.join_schedules(schedules=(warmup_fn, decay_fn), boundaries=(warmup_steps))


def create_optimizer(learning_rate_schedule: optax.Schedule) -> optax.GradientTransformation:
   return optax.chain(optax.clip_by_global_norm(1.0), optax.adamw(learning_rate=learning_rate_schedule, weight_decay=1e-4))

Wir erstellen einen benutzerdefinierten Trainingszustand, der Modellparameter und Batch-Statistiken verfolgt. Wir definieren außerdem einen Lernratenplan mit Aufwärmen und Kosinusabfall, gepaart mit einem AdamW-Optimierer, der Gradientenbeschneidung und Gewichtsabfall umfasst. Diese Kombination gewährleistet ein stabiles und adaptives Coaching. Schauen Sie sich das an VOLLSTÄNDIGE CODES hier.

@jit
def compute_metrics(logits, labels):
   loss = optax.softmax_cross_entropy_with_integer_labels(logits, labels).imply()
   accuracy = jnp.imply(jnp.argmax(logits, -1) == labels)
   return {'loss': loss, 'accuracy': accuracy}


def create_train_state(rng, mannequin, input_shape, learning_rate_schedule):
   variables = mannequin.init(rng, jnp.ones(input_shape), coaching=False)
   params = variables('params')
   batch_stats = variables.get('batch_stats', {})
   tx = create_optimizer(learning_rate_schedule)
   return TrainState.create(apply_fn=mannequin.apply, params=params, tx=tx, batch_stats=batch_stats)


@jit
def train_step(state, batch, dropout_rng):
   pictures, labels = batch
   def loss_fn(params):
       variables = {'params': params, 'batch_stats': state.batch_stats}
       logits, new_model_state = state.apply_fn(variables, pictures, coaching=True, mutable=('batch_stats'), rngs={'dropout': dropout_rng})
       loss = optax.softmax_cross_entropy_with_integer_labels(logits, labels).imply()
       return loss, (logits, new_model_state)
   grad_fn = jax.value_and_grad(loss_fn, has_aux=True)
   (loss, (logits, new_model_state)), grads = grad_fn(state.params)
   state = state.apply_gradients(grads=grads, batch_stats=new_model_state('batch_stats'))
   metrics = compute_metrics(logits, labels)
   return state, metrics


@jit
def eval_step(state, batch):
   pictures, labels = batch
   variables = {'params': state.params, 'batch_stats': state.batch_stats}
   logits = state.apply_fn(variables, pictures, coaching=False)
   return compute_metrics(logits, labels)

Wir implementieren JIT-kompilierte Trainings- und Bewertungsfunktionen, um eine effiziente Ausführung zu erreichen. Der Trainingsschritt berechnet Farbverläufe, aktualisiert Parameter und pflegt dynamisch Batch-Statistiken. Wir definieren auch Bewertungsmetriken, die uns helfen, Verlust und Genauigkeit während des gesamten Trainingsprozesses zu überwachen. Schauen Sie sich das an VOLLSTÄNDIGE CODES hier.

def generate_synthetic_data(rng, num_samples=1000, img_size=32):
   rng_x, rng_y = random.break up(rng)
   pictures = random.regular(rng_x, (num_samples, img_size, img_size, 3))
   labels = random.randint(rng_y, (num_samples,), 0, 10)
   return pictures, labels


def create_batches(pictures, labels, batch_size=32):
   num_batches = len(pictures) // batch_size
   for i in vary(num_batches):
       idx = slice(i * batch_size, (i + 1) * batch_size)
       yield pictures(idx), labels(idx)

Wir generieren synthetische Daten, um eine Bildklassifizierungsaufgabe zu simulieren, sodass wir das Modell trainieren können, ohne auf externe Datensätze angewiesen zu sein. Anschließend bündeln wir die Daten effizient für iterative Aktualisierungen. Mit diesem Ansatz können wir die gesamte Pipeline schnell testen und überprüfen, ob alle Komponenten ordnungsgemäß funktionieren. Schauen Sie sich das an VOLLSTÄNDIGE CODES hier.

def train_model(num_epochs=5, batch_size=32):
   rng = random.PRNGKey(0)
   rng, data_rng, model_rng = random.break up(rng, 3)
   train_images, train_labels = generate_synthetic_data(data_rng, num_samples=1000)
   test_images, test_labels = generate_synthetic_data(data_rng, num_samples=200)
   mannequin = AdvancedCNN(num_classes=10)
   lr_schedule = create_learning_rate_schedule(base_lr=1e-3, warmup_steps=50, decay_steps=500)
   state = create_train_state(model_rng, mannequin, (1, 32, 32, 3), lr_schedule)
   historical past = {'train_loss': (), 'train_acc': (), 'test_acc': ()}
   print("Beginning coaching...")
   for epoch in vary(num_epochs):
       train_metrics = ()
       for batch in create_batches(train_images, train_labels, batch_size):
           rng, dropout_rng = random.break up(rng)
           state, metrics = train_step(state, batch, dropout_rng)
           train_metrics.append(metrics)
       train_loss = jnp.imply(jnp.array((m('loss') for m in train_metrics)))
       train_acc = jnp.imply(jnp.array((m('accuracy') for m in train_metrics)))
       test_metrics = (eval_step(state, batch) for batch in create_batches(test_images, test_labels, batch_size))
       test_acc = jnp.imply(jnp.array((m('accuracy') for m in test_metrics)))
       historical past('train_loss').append(float(train_loss))
       historical past('train_acc').append(float(train_acc))
       historical past('test_acc').append(float(test_acc))
       print(f"Epoch {epoch + 1}/{num_epochs}: Loss: {train_loss:.4f}, Prepare Acc: {train_acc:.4f}, Check Acc: {test_acc:.4f}")
   return historical past, state


historical past, trained_state = train_model(num_epochs=5)


fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))
ax1.plot(historical past('train_loss'), label="Prepare Loss")
ax1.set_xlabel('Epoch'); ax1.set_ylabel('Loss'); ax1.set_title('Coaching Loss'); ax1.legend(); ax1.grid(True)
ax2.plot(historical past('train_acc'), label="Prepare Accuracy")
ax2.plot(historical past('test_acc'), label="Check Accuracy")
ax2.set_xlabel('Epoch'); ax2.set_ylabel('Accuracy'); ax2.set_title('Mannequin Accuracy'); ax2.legend(); ax2.grid(True)
plt.tight_layout(); plt.present()


print("n Tutorial full! This covers:")
print("- Customized Flax modules (ResNet blocks, Self-Consideration)")
print("- Superior Optax optimizers (AdamW with gradient clipping)")
print("- Studying charge schedules (warmup + cosine decay)")
print("- JAX transformations (@jit for efficiency)")
print("- Correct state administration (batch normalization statistics)")
print("- Full coaching pipeline with analysis")

Wir bringen alle Komponenten zusammen, um das Modell über mehrere Epochen zu trainieren, Leistungsmetriken zu verfolgen und die Verlust- und Genauigkeitstrends zu visualisieren. Wir überwachen den Lernfortschritt des Modells und validieren seine Leistung anhand von Testdaten. Letztendlich bestätigen wir die Stabilität und Wirksamkeit unseres JAX-basierten Trainingsworkflows.

Zusammenfassend haben wir eine umfassende Schulungspipeline unter Verwendung von JAX, Flax und Optax implementiert, die sowohl Flexibilität als auch Recheneffizienz demonstriert. Wir haben beobachtet, wie benutzerdefinierte Architekturen, fortschrittliche Optimierungsstrategien und präzises Zustandsmanagement zusammenkommen können, um einen leistungsstarken Deep-Studying-Workflow zu bilden. Durch diese Übung erlangen wir ein tieferes Verständnis für die Strukturierung skalierbarer Experimente in JAX und bereiten uns darauf vor, diese Techniken an reale Forschungs- und Produktionsaufgaben im Bereich maschinelles Lernen anzupassen.

Schauen Sie sich das an VOLLSTÄNDIGE CODES hier. Schauen Sie sich gerne bei uns um GitHub-Seite für Tutorials, Codes und Notebooks. Sie können uns auch gerne weiter folgen Twitter und vergessen Sie nicht, bei uns mitzumachen 100.000+ ML SubReddit und Abonnieren Unser E-newsletter. Warten! Bist du im Telegram? Jetzt können Sie uns auch per Telegram kontaktieren.

Asif Razzaq ist CEO von Marktechpost Media Inc.. Als visionärer Unternehmer und Ingenieur setzt sich Asif dafür ein, das Potenzial der künstlichen Intelligenz für das soziale Wohl zu nutzen. Sein jüngstes Unterfangen ist die Einführung einer Medienplattform für künstliche Intelligenz, Marktechpost, die sich durch eine ausführliche Berichterstattung über maschinelles Lernen und Deep-Studying-Nachrichten auszeichnet, die sowohl technisch fundiert als auch für ein breites Publikum leicht verständlich ist. Die Plattform verfügt über mehr als 2 Millionen monatliche Aufrufe, was ihre Beliebtheit beim Publikum verdeutlicht.

Folgen Sie MARKTECHPOST: Fügen Sie uns als bevorzugte Quelle bei Google hinzu.

Das ist etwas, was jeder, der viel Zeit mit Transformatoren und Selbstaufmerksamkeit verbracht hat, schon hundertmal gehört hat. Beides ist absolut wahr, wir haben das alle schon erlebt: Wenn man versucht, die Kontextgröße seines Modells zu erhöhen, kommt plötzlich alles zum Stillstand. Doch gleichzeitig gibt es, so scheint es, praktisch jede Woche ein neues hochmodernes Modell mit einer neuen rekordverdächtigen Kontextlänge. (Gemini hat eine Kontextlänge von 2 Millionen Token!)

Es gibt viele ausgefeilte Methoden wie RingAttention, die das Coaching unglaublich langer Kontextlängen in großen verteilten Systemen ermöglichen, aber was mich heute interessiert, ist eine einfachere Frage.

Wie weit kommen wir allein mit linearer Aufmerksamkeit?

Dies wird eine Artwork Blitztour, aber haben Sie etwas Geduld, denn wir werden ein paar wichtige Punkte ansprechen, bevor wir uns mit den Ergebnissen befassen.

Wir können den traditionellen Aufmerksamkeitsmechanismus im Wesentlichen mit zwei Kernpunkten zusammenfassen:

• Zunächst nimmt der typische Softmax-Aufmerksamkeitsausdruck das Produkt aus Abfrage- und Schlüsselmatrizen, normalisiert es zur Gewährleistung der Stabilität und nimmt dann das Softmax (zeilenweise), um die Aufmerksamkeitswerte zwischen den einzelnen Elementen der Sequenz zu erhalten.
• Zweitens wird die Zeitkomplexität von den N² Skalarprodukten dominiert, und das Produkt innerhalb des Softmax ist der limitierende Faktor. Dort berechnen wir die Aufmerksamkeitswerte.

Dies wird in der traditionellen Type wie folgt ausgedrückt:

Es stellt sich heraus, dass wir etwas anders darüber nachdenken können, wenn wir unsere Mathematikerfreunde fragen. Man kann sich den Softmax als eine von vielen Möglichkeiten vorstellen, die Wahrscheinlichkeitsverteilung zu beschreiben, die Token miteinander in Beziehung setzt. Wir können jedes beliebige Ähnlichkeitsmaß verwenden (das Skalarprodukt ist eines der einfachsten), und solange wir es normalisieren, ist alles in Ordnung.

Es ist ein bisschen schlampig, das zu sagen Ist Aufmerksamkeit, denn tatsächlich ist es nur die Aufmerksamkeit, die wir kennen und lieben, wenn die Ähnlichkeitsfunktion die Exponentialfunktion des Skalarprodukts von Abfragen und Schlüsseln (siehe unten) ist, wie wir es im Softmax finden. Aber hier wird es interessant: Was wäre, wenn wir statt dieses Ausdrucks eine Näherung vornehmen könnten?

Wir können davon ausgehen, dass es eine Characteristic-Map gibt.phi” was uns ein Ergebnis gibt quick das gleiche wie die Exponentialfunktion des Skalarprodukts. Und entscheidend ist, dass wir den Ausdruck auf diese Weise schreiben können, um mit der Reihenfolge der Matrixmultiplikationsoperationen zu spielen.

Im Papier Sie schlagen die Exponential Linear Unit (ELU) aufgrund einer Reihe nützlicher Eigenschaften als Merkmalskarte vor:

1. Für Werte über 0 liefert ELU(x) ein lineares Ergebnis, das zwar nicht mit der Exponentialfunktion identisch ist, aber die relative Reihenfolge zwischen den Wertungen beibehält.
2. Bei Werten kleiner oder gleich 0 bleibt der Exponentialterm der kontinuierliche Charakter der Funktion erhalten und stellt sicher, dass die Gradienten nicht einfach verschwinden.

Wir werden hier nicht allzu viel Zeit darauf verwenden, aber es ist empirisch ziemlich intestine bestätigt, dass dies eine gute Annäherung an die Softmax-Funktion ist.

Dadurch können wir die Reihenfolge der Operationen ändern. Wir können zuerst das Produkt unserer Characteristic-Map von Okay mit V nehmen, um einen KV-Block zu erstellen, und dann das Produkt mit Q. Das quadratische Produkt wird über die Modelldimensionsgröße und nicht über die Sequenzlänge bestimmt.

Wenn wir dies alles in den linearen Aufmerksamkeitsausdruck einfügen, erhalten wir:

Wobei wir die Begriffe in den Klammern nur einmal professional Abfragezeile berechnen müssen.

(Wenn Sie wissen möchten, wie die gelegentliche Maskierung hier hineinpasst und wie die Farbverläufe berechnet werden, werfen Sie einen Blick in das Dokument. Oder schauen Sie in diesem Bereich nach einem zukünftigen Weblog.)

Die mathematischen Argumente sind überzeugend, aber persönlich bin ich immer etwas misstrauisch, bis ich einige Benchmarks gesehen habe.

Sehen wir uns zunächst die Codeausschnitte an, die jeden dieser Begriffe beschreiben. Die Softmax-Aufmerksamkeit wird Ihnen sehr vertraut vorkommen, wir machen hier nichts Ausgefallenes.

class TraditionalAttention(nn.Module):
def __init__(self, d_k):
tremendous(TraditionalAttention, self).__init__()
self.d_k = d_k
def ahead(self, Q, Okay, V):
Z = torch.sqrt(torch.tensor(self.d_k, machine=Q.machine, dtype=torch.float32))
scores = torch.matmul(Q, Okay.transpose(-2, -1)) / Z
attention_weights = F.softmax(scores, dim=-1)
output = torch.matmul(attention_weights, V)
return output

Dann beginnen wir für die lineare Aufmerksamkeit mit dem Abrufen der Abfrage-, Schlüssel- und Wertematrizen und wenden dann die ELU(x)-Funktionszuordnung auf die Abfrage und die Schlüssel an. Dann verwenden wir die Einsummennotation, um die Multiplikationen durchzuführen.

class LinearAttention(nn.Module):
def __init__(self):
tremendous(LinearAttention, self).__init__()
self.eps = 1e-6
def elu_feature_map(self, x):
return F.elu(x) + 1
def ahead(self, Q, Okay, V):
Q = self.elu_feature_map(Q)
Okay = self.elu_feature_map(Okay)
KV = torch.einsum("nsd,nsd->ns", Okay, V)
# Compute the normalizer
Z = 1/(torch.einsum("nld,nd->nl", Q, Okay.sum(dim=1))+self.eps)
# Lastly compute and return the brand new values
V = torch.einsum("nld,ns,nl->nd", Q, KV, Z)
return V.contiguous()

Dies in Code geschrieben zu sehen ist schön und intestine, aber was bedeutet es tatsächlich experimentell? Von welcher Leistungssteigerung sprechen wir hier? Es kann schwierig sein, den Grad der Beschleunigung zu erfassen, wenn man von einem quadratischen zu einem linearen Engpass wechselt, additionally habe ich das folgende Experiment durchgeführt.

Wir werden eine einzelne Aufmerksamkeitsebene mit einer festen d_k-Modelldimension von 64 verwenden und die Zeit für einen Vorwärtsdurchlauf eines Sequenzsatzes mit 32 Batchgrößen als Benchmark verwenden. Die einzige zu ändernde Variable ist die Sequenzlänge, die zwischen 128 und 6000 liegt (die GPT-3-Kontextlänge als Referenz beträgt 2048). Jeder Durchlauf wird 100 Mal durchgeführt, um einen Mittelwert und eine Standardabweichung zu erhalten, und die Experimente werden mit einer Nvidia T4 GPU durchgeführt.

Für ein so einfaches Experiment sind die Ergebnisse ziemlich bemerkenswert.

Die Ergebnisse zeigen, dass wir sogar bei einem unglaublich kleinen Spielzeugbeispiel eine bis zu 60-fache Beschleunigung erreichen.

Diskussion

Hieraus lassen sich einige offensichtliche Erkenntnisse gewinnen:

1. Der Vorteil der linearen Aufmerksamkeit ist enorm – entweder in Bezug auf die Geschwindigkeit, ein höherer Durchsatz ist immer eine gute Sache. Oder in Bezug auf den Speicherbedarf zur Verarbeitung langer Sequenzen. In Umgebungen mit wenig Speicher kann dies ein großer Vorteil sein.
2. Das Verhältnisdiagramm weist einen überraschenden Knick auf – was uns vermuten lässt, dass hier eine zusätzliche Optimierung auf niedrigerer Ebene stattfindet, was bedeutet, dass das erwartete Verhältnis nicht ganz eintritt. Wir müssen dieses Ergebnis additionally mit Vorsicht genießen.

Der Vollständigkeit halber sei hier nicht verwechselt mit „lineare Aufmerksamkeit ist bei kleinen Modellen 60-mal schneller“. In Wirklichkeit machen die Feedforward-Schichten oft einen größeren Teil der Parameter in einem Transformer aus, und die Kodierung/Dekodierung ist oft auch eine Komponente mit limitierender Größe. Aber bei diesem eng definierten Downside ist das ziemlich beeindruckend!